🧠 Diffusion Fundamentals

🎯 What are Diffusion Models?

Diffusion models are a class of generative models that learn to create data by reversing a gradual noising process. They work by: (1) progressively adding noise to data until it becomes pure noise (forward diffusion), and (2) learning to reverse this process to generate new samples from noise (reverse diffusion). They power Stable Diffusion, DALL-E 2, Imagen, and Midjourney.

Why Diffusion Models Matter

State-of-the-art: Best quality in image generation. Stable training: Unlike GANs, no adversarial training. Diverse outputs: Covers modes of data distribution. Controllable: Text, class, image conditioning. Versatile: Images, video, audio, 3D, molecules. Foundation: Base for modern generative AI.

πŸ“Œ GenerativeπŸ“Œ DenoisingπŸ“Œ Stable TrainingπŸ“Œ High Quality
# Diffusion model intuition (conceptual) class DiffusionModel: def __init__(self, T=1000): self.T = T # number of timesteps self.noise_schedule = self.setup_schedule() self.model = UNet() # predicts noise def forward_diffusion(self, x_0, t): """Add noise to data: x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1-alpha_bar_t) * epsilon""" alpha_bar = self.noise_schedule['alpha_bar'][t] epsilon = torch.randn_like(x_0) x_t = torch.sqrt(alpha_bar) * x_0 + torch.sqrt(1 - alpha_bar) * epsilon return x_t, epsilon def reverse_diffusion(self, x_T): """Denoise from pure noise to data.""" x = x_T for t in reversed(range(self.T)): # Predict noise predicted_noise = self.model(x, t) # Compute x_{t-1} from x_t alpha = self.noise_schedule['alpha'][t] alpha_bar = self.noise_schedule['alpha_bar'][t] beta = self.noise_schedule['beta'][t] # Denoising step mean = (1 / torch.sqrt(alpha)) * (x - (beta / torch.sqrt(1 - alpha_bar)) * predicted_noise) if t > 0: noise = torch.randn_like(x) x = mean + torch.sqrt(beta) * noise else: x = mean return x # Generated image def train_step(self, x_0): """Training: predict noise added at random timestep.""" # 1. Sample random timestep t = torch.randint(0, self.T, (x_0.shape[0],)) # 2. Add noise x_t, epsilon = self.forward_diffusion(x_0, t) # 3. Predict noise predicted_noise = self.model(x_t, t) # 4. MSE loss loss = F.mse_loss(predicted_noise, epsilon) return loss

➑️ Forward Diffusion Process

🎯 Adding Noise

The forward diffusion process gradually adds Gaussian noise to an image over T timesteps. At t=0, we have the original image; at t=T, we have pure noise. This is a fixed Markov chain β€” no learning involved.

The Forward Process

Formula: q(x_t | x_{t-1}) = N(x_t; √(1-Ξ²_t) x_{t-1}, Ξ²_t I). Closed form: x_t = √(αΎ±_t) x_0 + √(1-αΎ±_t) Ξ΅, where αΎ±_t = Ξ (1-Ξ²_i). Ξ² schedule: Linear, cosine, or sigmoid. Effect: Each step adds a small amount of noise. At t=T: x_T ~ N(0, I) β€” pure Gaussian noise.

πŸ“ˆ Markov ChainπŸ”Š Gaussian NoiseπŸ“‰ Beta Scheduleβœ… No Learning
# Forward diffusion process import torch import numpy as np import matplotlib.pyplot as plt class ForwardDiffusion: def __init__(self, T=1000, beta_start=1e-4, beta_end=0.02): self.T = T # Linear beta schedule self.betas = torch.linspace(beta_start, beta_end, T) self.alphas = 1.0 - self.betas self.alpha_bars = torch.cumprod(self.alphas, dim=0) def add_noise(self, x_0, t): """Add noise at timestep t. Works for any t directly (closed form).""" alpha_bar = self.alpha_bars[t] # Reshape for broadcasting if len(x_0.shape) == 4: # (B, C, H, W) alpha_bar = alpha_bar.view(-1, 1, 1, 1) epsilon = torch.randn_like(x_0) x_t = torch.sqrt(alpha_bar) * x_0 + torch.sqrt(1 - alpha_bar) * epsilon return x_t, epsilon def visualize_forward(self, image): """Visualize forward diffusion at different timesteps.""" timesteps = [0, 100, 250, 500, 750, 999] fig, axes = plt.subplots(1, len(timesteps), figsize=(15, 3)) for i, t in enumerate(timesteps): x_t, _ = self.add_noise(image, torch.tensor([t])) axes[i].imshow(x_t[0].permute(1, 2, 0).clamp(0, 1)) axes[i].set_title(f't = {t}') axes[i].axis('off') plt.tight_layout() plt.show() # Example forward = ForwardDiffusion(T=1000) # Load or create a sample image x_0 = torch.rand(1, 3, 64, 64) # random "image" # Show alpha_bar progression print("Alpha bar values:") for t in [0, 100, 250, 500, 750, 999]: print(f" t={t:4d}: αΎ±_t = {forward.alpha_bars[t]:.6f}") # Beta schedules def cosine_schedule(T=1000, s=0.008): """Cosine noise schedule (often better than linear).""" steps = T + 1 x = torch.linspace(0, T, steps) alphas_cumprod = torch.cos(((x / T) + s) / (1 + s) * torch.pi * 0.5) ** 2 alphas_cumprod = alphas_cumprod / alphas_cumprod[0] betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) return torch.clip(betas, 0.0001, 0.9999) def sigmoid_schedule(T=1000, start=-3, end=3): """Sigmoid noise schedule.""" steps = T + 1 x = torch.linspace(start, end, steps) alphas_cumprod = torch.sigmoid(-x) alphas_cumprod = alphas_cumprod / alphas_cumprod[0] betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) return torch.clip(betas, 0.0001, 0.9999) # Compare schedules linear_betas = torch.linspace(1e-4, 0.02, 1000) cosine_betas = cosine_schedule(1000) sigmoid_betas = sigmoid_schedule(1000) linear_alpha_bars = torch.cumprod(1 - linear_betas, dim=0) cosine_alpha_bars = torch.cumprod(1 - cosine_betas, dim=0) sigmoid_alpha_bars = torch.cumprod(1 - sigmoid_betas, dim=0) print("\nFinal alpha_bar (t=T):") print(f" Linear: {linear_alpha_bars[-1]:.6f}") print(f" Cosine: {cosine_alpha_bars[-1]:.6f}") print(f" Sigmoid: {sigmoid_alpha_bars[-1]:.6f}") print("(All should be close to 0 = pure noise)")

⬅️ Reverse Diffusion Process

🎯 Removing Noise

The reverse diffusion process learns to undo the forward process β€” starting from pure noise and gradually denoising to generate an image. This is where the neural network comes in: it learns to predict the noise at each step.

The Reverse Process

Goal: Learn p(x_{t-1} | x_t) to reverse q(x_t | x_{t-1}). Network predicts: The noise Ξ΅ that was added. Formula: x_{t-1} = (1/√α_t)(x_t - (Ξ²_t/√(1-αΎ±_t)) Ξ΅_ΞΈ(x_t, t)) + Οƒ_t z. Mean: Denoised estimate. Variance: Either fixed (DDPM) or learned.

🧠 Neural NetworkπŸ”Š Predicts NoiseπŸ“‰ Iterative✨ Generation
# Reverse diffusion process import torch import torch.nn as nn import torch.nn.functional as F class ReverseDiffusion: def __init__(self, model, T=1000, beta_start=1e-4, beta_end=0.02): self.model = model self.T = T self.betas = torch.linspace(beta_start, beta_end, T) self.alphas = 1.0 - self.betas self.alpha_bars = torch.cumprod(self.alphas, dim=0) self.sqrt_alpha_bars = torch.sqrt(self.alpha_bars) self.sqrt_one_minus_alpha_bars = torch.sqrt(1 - self.alpha_bars) @torch.no_grad() def sample(self, shape, device='cuda'): """Generate samples by reversing diffusion from pure noise.""" # Start from pure Gaussian noise x = torch.randn(shape, device=device) # Iteratively denoise for t in reversed(range(self.T)): t_tensor = torch.full((shape[0],), t, device=device, dtype=torch.long) # 1. Predict noise predicted_noise = self.model(x, t_tensor) # 2. Get schedule values alpha = self.alphas[t] alpha_bar = self.alpha_bars[t] beta = self.betas[t] # 3. Compute mean of p(x_{t-1} | x_t) mean = (1 / torch.sqrt(alpha)) * ( x - (beta / torch.sqrt(1 - alpha_bar)) * predicted_noise ) # 4. Add noise (except at t=0) if t > 0: noise = torch.randn_like(x) # Posterior variance variance = beta * (1 - self.alpha_bars[t-1]) / (1 - alpha_bar) x = mean + torch.sqrt(variance) * noise else: x = mean return x def sample_with_progress(self, shape, device='cuda', save_every=100): """Sample with intermediate saves for visualization.""" x = torch.randn(shape, device=device) intermediates = [x.cpu()] for t in reversed(range(self.T)): t_tensor = torch.full((shape[0],), t, device=device, dtype=torch.long) predicted_noise = self.model(x, t_tensor) alpha = self.alphas[t] alpha_bar = self.alpha_bars[t] beta = self.betas[t] mean = (1 / torch.sqrt(alpha)) * ( x - (beta / torch.sqrt(1 - alpha_bar)) * predicted_noise ) if t > 0: noise = torch.randn_like(x) variance = beta * (1 - self.alpha_bars[t-1]) / (1 - alpha_bar) x = mean + torch.sqrt(variance) * noise else: x = mean if t % save_every == 0: intermediates.append(x.cpu()) return x, intermediates # Simplified model (U-Net would be here) class SimpleDenoiser(nn.Module): def __init__(self, channels=3, hidden=64): super().__init__() self.net = nn.Sequential( nn.Conv2d(channels + 1, hidden, 3, padding=1), # +1 for timestep nn.ReLU(), nn.Conv2d(hidden, hidden, 3, padding=1), nn.ReLU(), nn.Conv2d(hidden, channels, 3, padding=1) ) def forward(self, x, t): # Add timestep as a channel t_channel = t.float().view(-1, 1, 1, 1) / 1000.0 t_channel = t_channel.expand(-1, 1, x.shape[2], x.shape[3]) x_with_t = torch.cat([x, t_channel], dim=1) return self.net(x_with_t) # Usage model = SimpleDenoiser() reverse = ReverseDiffusion(model, T=1000) # Generate samples samples = reverse.sample(shape=(4, 3, 32, 32)) print(f"Generated samples: {samples.shape}") print(f"Range: [{samples.min():.3f}, {samples.max():.3f}]") # The difference between forward and reverse: # FORWARD (fixed): x_0 β†’ x_1 β†’ ... β†’ x_T (adds noise) # REVERSE (learned): x_T β†’ x_{T-1} β†’ ... β†’ x_0 (removes noise) print("\nForward: x_0 β†’ ... β†’ x_T (deterministic noising)") print("Reverse: x_T β†’ ... β†’ x_0 (learned denoising)")

πŸŽ“ Training Objective

🎯 Learning to Denoise

The training objective is simple: predict the noise that was added at a random timestep. The loss is MSE between predicted and actual noise. This is called the simple loss or Ξ΅-prediction.

Training Loss

Objective: L = E[||Ρ - Ρ_θ(x_t, t)||²]. Sample t: Uniform random from 1 to T. Add noise: x_t = √(ᾱ_t)x_0 + √(1-ᾱ_t)Ρ. Predict: Network outputs Ρ_θ. Loss: MSE between Ρ and Ρ_θ. Variants: x_0-prediction, v-prediction.

πŸ“‰ MSE Loss🎲 Random tπŸ”Š Predict Ρ⚑ Simple
# Training diffusion model import torch import torch.nn as nn import torch.nn.functional as F from torch.utils.data import DataLoader from torch.optim import Adam class DiffusionTrainer: def __init__(self, model, T=1000, lr=2e-4): self.model = model self.T = T # Noise schedule self.betas = torch.linspace(1e-4, 0.02, T) self.alphas = 1.0 - self.betas self.alpha_bars = torch.cumprod(self.alphas, dim=0) # Optimizer self.optimizer = Adam(model.parameters(), lr=lr) def train_step(self, x_0, device='cuda'): """One training step.""" batch_size = x_0.shape[0] # 1. Sample random timesteps t = torch.randint(0, self.T, (batch_size,), device=device) # 2. Add noise alpha_bar = self.alpha_bars[t].to(device) alpha_bar = alpha_bar.view(-1, 1, 1, 1) epsilon = torch.randn_like(x_0) x_t = torch.sqrt(alpha_bar) * x_0 + torch.sqrt(1 - alpha_bar) * epsilon # 3. Predict noise predicted_noise = self.model(x_t, t) # 4. MSE loss loss = F.mse_loss(predicted_noise, epsilon) # 5. Backprop self.optimizer.zero_grad() loss.backward() self.optimizer.step() return loss.item() def train(self, dataloader, epochs=100, device='cuda'): """Full training loop.""" losses = [] for epoch in range(epochs): epoch_losses = [] for batch in dataloader: x_0 = batch[0] if isinstance(batch, tuple) else batch x_0 = x_0.to(device) loss = self.train_step(x_0, device) epoch_losses.append(loss) avg_loss = sum(epoch_losses) / len(epoch_losses) losses.append(avg_loss) print(f"Epoch {epoch+1}/{epochs}: Loss = {avg_loss:.6f}") return losses # Different prediction targets def get_training_targets(x_0, x_t, epsilon, alpha_bar): """Compute different training targets.""" return { 'epsilon': epsilon, # Predict noise (standard) 'x_0': x_0, # Predict original image 'v': torch.sqrt(alpha_bar) * epsilon - torch.sqrt(1 - alpha_bar) * x_0 # v-prediction } # v-prediction loss (better for some cases) class VPredictionTrainer(DiffusionTrainer): def train_step(self, x_0, device='cuda'): batch_size = x_0.shape[0] t = torch.randint(0, self.T, (batch_size,), device=device) alpha_bar = self.alpha_bars[t].to(device).view(-1, 1, 1, 1) epsilon = torch.randn_like(x_0) x_t = torch.sqrt(alpha_bar) * x_0 + torch.sqrt(1 - alpha_bar) * epsilon # v-target v_target = torch.sqrt(alpha_bar) * epsilon - torch.sqrt(1 - alpha_bar) * x_0 # Model predicts v v_predicted = self.model(x_t, t) loss = F.mse_loss(v_predicted, v_target) self.optimizer.zero_grad() loss.backward() self.optimizer.step() return loss.item() # Min-SNR weighting (improves training stability) def min_snr_weight(alpha_bars, t, gamma=5.0): """Min-SNR weighting for diffusion training.""" snr = alpha_bars[t] / (1 - alpha_bars[t]) return torch.minimum(snr, torch.tensor(gamma, device=snr.device)) / snr # Training with Min-SNR weighting def train_step_min_snr(self, x_0, device='cuda'): batch_size = x_0.shape[0] t = torch.randint(0, self.T, (batch_size,), device=device) alpha_bar = self.alpha_bars[t].to(device).view(-1, 1, 1, 1) epsilon = torch.randn_like(x_0) x_t = torch.sqrt(alpha_bar) * x_0 + torch.sqrt(1 - alpha_bar) * epsilon predicted_noise = self.model(x_t, t) # Weighted MSE weight = min_snr_weight(self.alpha_bars, t).view(-1, 1, 1, 1) loss = (weight * (predicted_noise - epsilon) ** 2).mean() self.optimizer.zero_grad() loss.backward() self.optimizer.step() return loss.item()

🧬 U-Net Architecture

🎯 The Denoising Network

The U-Net is the standard architecture for diffusion models. It's an encoder-decoder network with skip connections that preserves spatial information while extracting features at multiple scales.

U-Net Components

Encoder: Downsampling path, extracts features. Bottleneck: Lowest resolution, most abstract features. Decoder: Upsampling path, reconstructs image. Skip connections: Preserve fine details. Time embedding: Injects timestep info. Attention: Self-attention at low resolutions.

πŸ“₯ EncoderπŸ“€ DecoderπŸ”— Skip⏱️ Time
# U-Net for diffusion models import torch import torch.nn as nn import torch.nn.functional as F import math class SinusoidalPositionEmbedding(nn.Module): """Sinusoidal embeddings for timestep.""" def __init__(self, dim): super().__init__() self.dim = dim def forward(self, time): device = time.device half_dim = self.dim // 2 embeddings = math.log(10000) / (half_dim - 1) embeddings = torch.exp(torch.arange(half_dim, device=device) * -embeddings) embeddings = time[:, None] * embeddings[None, :] embeddings = torch.cat((embeddings.sin(), embeddings.cos()), dim=-1) return embeddings class ResidualBlock(nn.Module): """Residual block with time embedding.""" def __init__(self, in_channels, out_channels, time_emb_dim): super().__init__() self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1) self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1) self.time_mlp = nn.Linear(time_emb_dim, out_channels) self.norm1 = nn.GroupNorm(8, out_channels) self.norm2 = nn.GroupNorm(8, out_channels) self.residual = nn.Conv2d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity() def forward(self, x, t_emb): h = self.norm1(F.silu(self.conv1(x))) # Add time embedding t = self.time_mlp(t_emb)[:, :, None, None] h = h + t h = self.norm2(F.silu(self.conv2(h))) return h + self.residual(x) class AttentionBlock(nn.Module): """Self-attention block.""" def __init__(self, channels): super().__init__() self.norm = nn.GroupNorm(8, channels) self.qkv = nn.Conv2d(channels, channels * 3, 1) self.proj = nn.Conv2d(channels, channels, 1) def forward(self, x): B, C, H, W = x.shape h = self.norm(x) qkv = self.qkv(h).reshape(B, 3, C, H * W).permute(1, 0, 2, 3) q, k, v = qkv[0], qkv[1], qkv[2] scale = C ** -0.5 attn = torch.softmax(q.transpose(-2, -1) @ k * scale, dim=-1) out = (v @ attn.transpose(-2, -1)).reshape(B, C, H, W) return x + self.proj(out) class UNet(nn.Module): def __init__(self, in_channels=3, out_channels=3, time_emb_dim=256, base_channels=64, channel_multipliers=(1, 2, 4, 8)): super().__init__() # Time embedding self.time_embed = nn.Sequential( SinusoidalPositionEmbedding(base_channels), nn.Linear(base_channels, time_emb_dim), nn.SiLU(), nn.Linear(time_emb_dim, time_emb_dim) ) # Initial conv self.init_conv = nn.Conv2d(in_channels, base_channels, 3, padding=1) # Encoder self.encoder = nn.ModuleList() chs = [base_channels] for i, mult in enumerate(channel_multipliers): out_ch = base_channels * mult for _ in range(2): self.encoder.append(ResidualBlock(chs[-1], out_ch, time_emb_dim)) chs.append(out_ch) if i < len(channel_multipliers) - 1: self.encoder.append(nn.Conv2d(out_ch, out_ch, 3, stride=2, padding=1)) chs.append(out_ch) # Middle self.middle = nn.Sequential( ResidualBlock(chs[-1], chs[-1], time_emb_dim), AttentionBlock(chs[-1]), ResidualBlock(chs[-1], chs[-1], time_emb_dim) ) # Decoder self.decoder = nn.ModuleList() for i, mult in reversed(list(enumerate(channel_multipliers))): out_ch = base_channels * mult for _ in range(3): self.decoder.append(ResidualBlock(chs[-1] + out_ch, out_ch, time_emb_dim)) chs.append(out_ch) if i > 0: self.decoder.append(nn.ConvTranspose2d(out_ch, out_ch, 4, 2, 1)) chs.append(out_ch) # Final conv self.final = nn.Sequential( nn.GroupNorm(8, chs[-1]), nn.SiLU(), nn.Conv2d(chs[-1], out_channels, 3, padding=1) ) def forward(self, x, t): # Time embedding t_emb = self.time_embed(t) # Initial h = self.init_conv(x) # Encoder with skip connections skips = [h] for layer in self.encoder: if isinstance(layer, ResidualBlock): h = layer(h, t_emb) else: h = layer(h) skips.append(h) # Middle h = self.middle(h) # Decoder for layer in self.decoder: if isinstance(layer, ResidualBlock): skip = skips.pop() # Concatenate skip connection if h.shape[2:] != skip.shape[2:]: skip = F.interpolate(skip, size=h.shape[2:], mode='nearest') h = torch.cat([h, skip], dim=1) h = layer(h, t_emb) else: h = layer(h) return self.final(h) # Create U-Net unet = UNet(in_channels=3, out_channels=3, base_channels=64) x = torch.randn(2, 3, 64, 64) t = torch.randint(0, 1000, (2,)) out = unet(x, t) print(f"Input: {x.shape}, Output: {out.shape}") print(f"Parameters: {sum(p.numel() for p in unet.parameters()):,}")

🎲 Sampling & Denoising

🎯 Generating from Noise

Sampling is the process of generating new data by running the reverse diffusion process. Starting from random noise, we iteratively denoise over T steps to produce a final image.

Sampling Methods

DDPM: Standard, 1000 steps, stochastic. DDIM: Deterministic, 50-100 steps. DPM-Solver: Fast, 10-20 steps. Euler: Simple ODE solver. Heun: Second-order solver. Trade-off: Speed vs quality.

🐒 DDPM⚑ DDIMπŸš€ DPM-SolverπŸ“‰ Fewer Steps
# Sampling methods for diffusion models import torch import numpy as np class DDPMSampler: """Denoising Diffusion Probabilistic Models sampler.""" def __init__(self, model, T=1000): self.model = model self.T = T self.betas = torch.linspace(1e-4, 0.02, T) self.alphas = 1 - self.betas self.alpha_bars = torch.cumprod(self.alphas, dim=0) @torch.no_grad() def sample(self, shape, device='cuda'): x = torch.randn(shape, device=device) for t in reversed(range(self.T)): t_batch = torch.full((shape[0],), t, device=device, dtype=torch.long) eps = self.model(x, t_batch) alpha = self.alphas[t].to(device) alpha_bar = self.alpha_bars[t].to(device) beta = self.betas[t].to(device) mean = (x - (beta / torch.sqrt(1 - alpha_bar)) * eps) / torch.sqrt(alpha) if t > 0: var = beta * (1 - self.alpha_bars[t-1].to(device)) / (1 - alpha_bar) x = mean + torch.sqrt(var) * torch.randn_like(x) else: x = mean return x class DDIMSampler: """Denoising Diffusion Implicit Models sampler β€” deterministic and faster.""" def __init__(self, model, T=1000, num_steps=50, eta=0.0): self.model = model self.T = T self.num_steps = num_steps self.eta = eta # 0 = deterministic, 1 = DDPM self.betas = torch.linspace(1e-4, 0.02, T) self.alphas = 1 - self.betas self.alpha_bars = torch.cumprod(self.alphas, dim=0) @torch.no_grad() def sample(self, shape, device='cuda'): # Sub-sample timesteps step_size = self.T // self.num_steps timesteps = list(range(0, self.T, step_size))[::-1] x = torch.randn(shape, device=device) for i, t in enumerate(timesteps): t_batch = torch.full((shape[0],), t, device=device, dtype=torch.long) eps = self.model(x, t_batch) alpha_bar = self.alpha_bars[t].to(device) # Predict x_0 x_0_pred = (x - torch.sqrt(1 - alpha_bar) * eps) / torch.sqrt(alpha_bar) x_0_pred = x_0_pred.clamp(-1, 1) if i < len(timesteps) - 1: t_next = timesteps[i + 1] alpha_bar_next = self.alpha_bars[t_next].to(device) # DDIM update sigma = self.eta * torch.sqrt( (1 - alpha_bar_next) / (1 - alpha_bar) * (1 - alpha_bar / alpha_bar_next) ) c = torch.sqrt(1 - alpha_bar_next - sigma ** 2) x = torch.sqrt(alpha_bar_next) * x_0_pred + c * eps if self.eta > 0: x = x + sigma * torch.randn_like(x) else: x = x_0_pred return x class DPMSolver: """DPM-Solver: fast ODE solver for diffusion models.""" def __init__(self, model, T=1000, num_steps=20): self.model = model self.T = T self.num_steps = num_steps self.betas = torch.linspace(1e-4, 0.02, T) self.alphas = 1 - self.betas self.alpha_bars = torch.cumprod(self.alphas, dim=0) @torch.no_grad() def sample(self, shape, device='cuda'): timesteps = torch.linspace(self.T - 1, 0, self.num_steps + 1).long() x = torch.randn(shape, device=device) for i in range(self.num_steps): t_cur = timesteps[i].item() t_next = timesteps[i + 1].item() t_batch = torch.full((shape[0],), t_cur, device=device, dtype=torch.long) eps = self.model(x, t_batch) alpha_bar_cur = self.alpha_bars[t_cur].to(device) alpha_bar_next = self.alpha_bars[t_next].to(device) if t_next >= 0 else torch.tensor(1.0, device=device) # First-order DPM-Solver update x_0_pred = (x - torch.sqrt(1 - alpha_bar_cur) * eps) / torch.sqrt(alpha_bar_cur) x = torch.sqrt(alpha_bar_next) * x_0_pred + torch.sqrt(1 - alpha_bar_next) * eps return x # Comparison def compare_samplers(model, device='cuda'): import time shape = (1, 3, 64, 64) # DDPM (1000 steps) sampler = DDPMSampler(model, T=1000) t0 = time.time() _ = sampler.sample(shape, device) ddpm_time = time.time() - t0 # DDIM (50 steps) sampler = DDIMSampler(model, T=1000, num_steps=50) t0 = time.time() _ = sampler.sample(shape, device) ddim_time = time.time() - t0 # DPM-Solver (20 steps) sampler = DPMSolver(model, T=1000, num_steps=20) t0 = time.time() _ = sampler.sample(shape, device) dpm_time = time.time() - t0 print(f"DDPM (1000 steps): {ddpm_time:.2f}s") print(f"DDIM (50 steps): {ddim_time:.2f}s") print(f"DPM-Solver (20): {dpm_time:.2f}s") print(f"Speedup: {ddpm_time / dpm_time:.1f}x")

πŸ“Š DDPM (Denoising Diffusion)

🎯 The Foundation

DDPM (Ho et al., 2020) is the original diffusion model. It introduced the simple noise-prediction objective and demonstrated high-quality image generation. It requires 1000 sampling steps.

DDPM Key Ideas

Forward: Fixed Gaussian noise schedule. Reverse: Learned Gaussian transitions. Loss: Simple MSE on noise. Sampling: 1000 steps (slow). Quality: State-of-the-art at the time. Impact: Sparked modern diffusion era.

πŸ“„ PaperπŸ”Š Noise Pred🐒 1000 Steps🎯 Quality
# DDPM implementation from scratch import torch import torch.nn as nn import torch.nn.functional as F import numpy as np class DDPM(nn.Module): """Denoising Diffusion Probabilistic Model.""" def __init__(self, model, T=1000, beta_start=1e-4, beta_end=0.02): super().__init__() self.model = model self.T = T # Register noise schedule as buffers self.register_buffer('betas', torch.linspace(beta_start, beta_end, T)) self.register_buffer('alphas', 1 - self.betas) self.register_buffer('alpha_bars', torch.cumprod(self.alphas, dim=0)) # Posterior variance alpha_bars_prev = F.pad(self.alpha_bars[:-1], (1, 0), value=1.0) self.register_buffer('posterior_variance', self.betas * (1 - alpha_bars_prev) / (1 - self.alpha_bars)) def q_sample(self, x_0, t, noise=None): """Forward diffusion: q(x_t | x_0).""" if noise is None: noise = torch.randn_like(x_0) alpha_bar = self.alpha_bars[t].view(-1, 1, 1, 1) return torch.sqrt(alpha_bar) * x_0 + torch.sqrt(1 - alpha_bar) * noise def p_sample(self, x_t, t): """Reverse: p(x_{t-1} | x_t).""" t_batch = torch.full((x_t.shape[0],), t, device=x_t.device, dtype=torch.long) eps_pred = self.model(x_t, t_batch) alpha = self.alphas[t] alpha_bar = self.alpha_bars[t] beta = self.betas[t] # Mean mean = (1 / torch.sqrt(alpha)) * (x_t - (beta / torch.sqrt(1 - alpha_bar)) * eps_pred) # Add noise (except at t=0) if t > 0: noise = torch.randn_like(x_t) var = self.posterior_variance[t] return mean + torch.sqrt(var) * noise return mean @torch.no_grad() def sample(self, shape, device='cuda'): """Generate samples.""" x = torch.randn(shape, device=device) for t in reversed(range(self.T)): x = self.p_sample(x, t) return x def loss(self, x_0): """Training loss: predict the noise.""" batch_size = x_0.shape[0] t = torch.randint(0, self.T, (batch_size,), device=x_0.device) noise = torch.randn_like(x_0) x_t = self.q_sample(x_0, t, noise) eps_pred = self.model(x_t, t) return F.mse_loss(eps_pred, noise) # Training DDPM def train_ddpm(model, dataloader, epochs=100, lr=2e-4, device='cuda'): ddpm = DDPM(model).to(device) optimizer = torch.optim.Adam(ddpm.parameters(), lr=lr) for epoch in range(epochs): total_loss = 0 for batch in dataloader: x_0 = batch[0] if isinstance(batch, tuple) else batch x_0 = x_0.to(device) loss = ddpm.loss(x_0) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() print(f"Epoch {epoch+1}: Loss = {total_loss / len(dataloader):.6f}") return ddpm # DDPM variants class DDPMVariants: """Different DDPM prediction targets.""" @staticmethod def epsilon_prediction(model, x_t, t): """Predict noise (standard DDPM).""" return model(x_t, t) @staticmethod def x0_prediction(model, x_t, t, alpha_bars): """Predict original image.""" eps = model(x_t, t) alpha_bar = alpha_bars[t].view(-1, 1, 1, 1) x_0 = (x_t - torch.sqrt(1 - alpha_bar) * eps) / torch.sqrt(alpha_bar) return x_0.clamp(-1, 1) @staticmethod def v_prediction(model, x_t, t, alpha_bars): """v-prediction (improved stability).""" v = model(x_t, t) alpha_bar = alpha_bars[t].view(-1, 1, 1, 1) x_0 = torch.sqrt(alpha_bar) * x_t - torch.sqrt(1 - alpha_bar) * v eps = torch.sqrt(1 - alpha_bar) * x_t + torch.sqrt(alpha_bar) * v return x_0, eps # DDPM improvements class ImprovedDDPM: """Improved DDPM techniques.""" @staticmethod def cosine_schedule(T=1000, s=0.008): """Cosine noise schedule (better than linear).""" steps = T + 1 x = torch.linspace(0, T, steps) alphas_cumprod = torch.cos(((x / T) + s) / (1 + s) * torch.pi * 0.5) ** 2 alphas_cumprod = alphas_cumprod / alphas_cumprod[0] betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1]) return torch.clip(betas, 0.0001, 0.9999) @staticmethod def learned_variance(model, x_t, t): """Learn the variance instead of fixing it.""" output = model(x_t, t) # Split output into mean and variance eps, log_var = output.chunk(2, dim=1) return eps, log_var

⚑ DDIM (Faster Sampling)

🎯 Deterministic & Fast

DDIM (Song et al., 2021) is a deterministic sampling method that dramatically reduces the number of sampling steps. It's 10-50x faster than DDPM while maintaining quality.

DDIM Advantages

Speed: 20-100 steps instead of 1000. Deterministic: Same noise β†’ same output. Interpolation: Smooth latent space. Quality: Comparable to DDPM. Compatible: Uses same trained model. Ξ· parameter: Controls stochasticity.

⚑ Fast🎯 DeterministicπŸ”„ CompatibleπŸ“‰ 20-100 Steps
# DDIM sampler β€” detailed implementation import torch import numpy as np class DDIM: """Denoising Diffusion Implicit Models.""" def __init__(self, model, T=1000, num_steps=50, eta=0.0): self.model = model self.T = T self.num_steps = num_steps self.eta = eta # 0 = deterministic, 1 = DDPM-like # Noise schedule self.betas = torch.linspace(1e-4, 0.02, T) self.alphas = 1 - self.betas self.alpha_bars = torch.cumprod(self.alphas, dim=0) @torch.no_grad() def sample(self, shape, device='cuda'): """Generate samples with DDIM.""" # Create sub-sampled timestep schedule step_indices = torch.linspace(0, self.T - 1, self.num_steps + 1, dtype=torch.long) timesteps = step_indices.flip(0) # Descending # Start from pure noise x = torch.randn(shape, device=device) for i in range(len(timesteps) - 1): t_cur = timesteps[i].item() t_next = timesteps[i + 1].item() # Predict noise t_batch = torch.full((shape[0],), t_cur, device=device, dtype=torch.long) eps = self.model(x, t_batch) # Get alpha_bar values alpha_bar_cur = self.alpha_bars[t_cur].to(device) alpha_bar_next = self.alpha_bars[t_next].to(device) if t_next >= 0 else torch.tensor(1.0, device=device) # Predict x_0 x_0_pred = (x - torch.sqrt(1 - alpha_bar_cur) * eps) / torch.sqrt(alpha_bar_cur) x_0_pred = x_0_pred.clamp(-1, 1) # Clamp for stability # Compute sigma (variance) sigma = self.eta * torch.sqrt( (1 - alpha_bar_next) / (1 - alpha_bar_cur) * (1 - alpha_bar_cur / alpha_bar_next) ) # Direction pointing to x_t c = torch.sqrt(1 - alpha_bar_next - sigma ** 2) # DDIM update x = torch.sqrt(alpha_bar_next) * x_0_pred + c * eps # Add noise if eta > 0 if self.eta > 0: x = x + sigma * torch.randn_like(x) return x @torch.no_grad() def interpolate(self, x1, x2, alpha=0.5, device='cuda'): """Interpolate between two images in latent space.""" # Encode to noise t = torch.full((x1.shape[0],), self.T - 1, device=device, dtype=torch.long) z1 = self.encode(x1, t) z2 = self.encode(x2, t) # Interpolate z = alpha * z1 + (1 - alpha) * z2 # Decode return self.sample_from_noise(z, device) @torch.no_grad() def encode(self, x_0, t, device='cuda'): """Encode image to noise (forward process).""" alpha_bar = self.alpha_bars[t].view(-1, 1, 1, 1) noise = torch.randn_like(x_0) x_t = torch.sqrt(alpha_bar) * x_0 + torch.sqrt(1 - alpha_bar) * noise return x_t @torch.no_grad() def sample_from_noise(self, x_T, device='cuda'): """Sample from a given noise.""" return self.sample(x_T.shape, device) # Sampling with different eta values def compare_eta(model, device='cuda'): """Compare DDIM with different eta values.""" shape = (1, 3, 64, 64) results = {} for eta in [0.0, 0.5, 1.0]: ddim = DDIM(model, T=1000, num_steps=50, eta=eta) x = ddim.sample(shape, device) results[eta] = x print("eta=0.0: Fully deterministic") print("eta=0.5: Semi-stochastic") print("eta=1.0: Matches DDPM") return results # DDIM inversion (encode real image to noise) class DDIMInversion: """Invert a real image back to noise.""" def __init__(self, model, T=1000, num_steps=50): self.model = model self.T = T self.num_steps = num_steps self.betas = torch.linspace(1e-4, 0.02, T) self.alphas = 1 - self.betas self.alpha_bars = torch.cumprod(self.alphas, dim=0) @torch.no_grad() def invert(self, x_0, device='cuda'): """Invert image to noise (approximate).""" step_indices = torch.linspace(0, self.T - 1, self.num_steps + 1, dtype=torch.long) timesteps = step_indices # Ascending for inversion x = x_0 for i in range(len(timesteps) - 1): t_cur = timesteps[i].item() t_next = timesteps[i + 1].item() t_batch = torch.full((x.shape[0],), t_cur, device=device, dtype=torch.long) eps = self.model(x, t_batch) alpha_bar_cur = self.alpha_bars[t_cur].to(device) alpha_bar_next = self.alpha_bars[t_next].to(device) # Inversion step x_0_pred = (x - torch.sqrt(1 - alpha_bar_cur) * eps) / torch.sqrt(alpha_bar_cur) x = torch.sqrt(alpha_bar_next) * x_0_pred + torch.sqrt(1 - alpha_bar_next) * eps return x # Example usage class SimpleModel(torch.nn.Module): def __init__(self): super().__init__() self.conv = torch.nn.Conv2d(4, 3, 3, padding=1) def forward(self, x, t): t_ch = t.float().view(-1, 1, 1, 1) / 1000 t_ch = t_ch.expand(-1, 1, x.shape[2], x.shape[3]) return self.conv(torch.cat([x, t_ch], dim=1)) model = SimpleModel() ddim = DDIM(model, T=1000, num_steps=50, eta=0.0) x = ddim.sample((1, 3, 64, 64)) print(f"DDIM sample shape: {x.shape}") print(f"DDIM uses {ddim.num_steps} steps instead of {ddim.T}")

🎨 Latent Diffusion (Stable Diffusion)

🎯 Diffusion in Latent Space

Latent Diffusion Models (LDM) perform diffusion in a compressed latent space instead of pixel space. This makes them much faster and enables high-resolution generation. Stable Diffusion is based on LDM.

Latent Diffusion Key Ideas

1. Encode: VAE compresses image to latent (e.g., 64Γ—64Γ—4). 2. Diffuse: Run diffusion in latent space. 3. Decode: VAE decodes latent back to image. Benefits: 8-16x smaller, faster, high-res. Cross-attention: Text conditioning in U-Net. Scale: Can generate 1024Γ—1024+ images.

πŸ—œοΈ VAEπŸ“¦ Latent Space⚑ FastπŸ–ΌοΈ High-res
# Latent Diffusion Model (Stable Diffusion architecture) import torch import torch.nn as nn import torch.nn.functional as F class VAE(nn.Module): """Variational Autoencoder for encoding/decoding latents.""" def __init__(self, in_channels=3, latent_channels=4, hidden=128): super().__init__() # Encoder self.encoder = nn.Sequential( nn.Conv2d(in_channels, hidden, 3, stride=2, padding=1), nn.GroupNorm(8, hidden), nn.SiLU(), nn.Conv2d(hidden, hidden * 2, 3, stride=2, padding=1), nn.GroupNorm(8, hidden * 2), nn.SiLU(), nn.Conv2d(hidden * 2, hidden * 4, 3, stride=2, padding=1), nn.GroupNorm(8, hidden * 4), nn.SiLU(), nn.Conv2d(hidden * 4, latent_channels * 2, 3, padding=1) # mean + logvar ) # Decoder self.decoder = nn.Sequential( nn.Conv2d(latent_channels, hidden * 4, 3, padding=1), nn.GroupNorm(8, hidden * 4), nn.SiLU(), nn.Upsample(scale_factor=2, mode='nearest'), nn.Conv2d(hidden * 4, hidden * 2, 3, padding=1), nn.GroupNorm(8, hidden * 2), nn.SiLU(), nn.Upsample(scale_factor=2, mode='nearest'), nn.Conv2d(hidden * 2, hidden, 3, padding=1), nn.GroupNorm(8, hidden), nn.SiLU(), nn.Upsample(scale_factor=2, mode='nearest'), nn.Conv2d(hidden, in_channels, 3, padding=1) ) def encode(self, x): h = self.encoder(x) mean, logvar = h.chunk(2, dim=1) return mean, logvar def reparameterize(self, mean, logvar): std = torch.exp(0.5 * logvar) eps = torch.randn_like(std) return mean + eps * std def decode(self, z): return self.decoder(z) def forward(self, x): mean, logvar = self.encode(x) z = self.reparameterize(mean, logvar) return self.decode(z), mean, logvar class CrossAttention(nn.Module): """Cross-attention for text conditioning.""" def __init__(self, query_dim, context_dim, heads=8, dim_head=64): super().__init__() inner_dim = heads * dim_head self.heads = heads self.scale = dim_head ** -0.5 self.to_q = nn.Linear(query_dim, inner_dim, bias=False) self.to_k = nn.Linear(context_dim, inner_dim, bias=False) self.to_v = nn.Linear(context_dim, inner_dim, bias=False) self.to_out = nn.Linear(inner_dim, query_dim) def forward(self, x, context): # x: (B, N, query_dim), context: (B, M, context_dim) q = self.to_q(x) k = self.to_k(context) v = self.to_v(context) # Reshape for multi-head B, N, _ = q.shape M = k.shape[1] q = q.view(B, N, self.heads, -1).transpose(1, 2) k = k.view(B, M, self.heads, -1).transpose(1, 2) v = v.view(B, M, self.heads, -1).transpose(1, 2) # Attention attn = torch.softmax(q @ k.transpose(-2, -1) * self.scale, dim=-1) out = (attn @ v).transpose(1, 2).reshape(B, N, -1) return self.to_out(out) class LatentDiffusion(nn.Module): """Latent Diffusion Model.""" def __init__(self, unet, vae, text_encoder, T=1000): super().__init__() self.unet = unet self.vae = vae self.text_encoder = text_encoder self.T = T self.betas = torch.linspace(1e-4, 0.02, T) self.alphas = 1 - self.betas self.alpha_bars = torch.cumprod(self.alphas, dim=0) @torch.no_grad() def encode_image(self, x): """Encode image to latent.""" mean, logvar = self.vae.encode(x) return self.vae.reparameterize(mean, logvar) * 0.18215 @torch.no_grad() def decode_latent(self, z): """Decode latent to image.""" return self.vae.decode(z / 0.18215) @torch.no_grad() def generate(self, prompt, num_steps=50, guidance_scale=7.5): """Generate image from text prompt.""" # 1. Encode text text_emb = self.text_encoder(prompt) # 2. Start from random latent z = torch.randn(1, 4, 64, 64) # Latent shape # 3. Denoise in latent space for t in reversed(range(0, self.T, self.T // num_steps)): # Predict noise (with text conditioning) eps = self.unet(z, t, text_emb) # Classifier-free guidance eps_uncond = self.unet(z, t, None) eps = eps_uncond + guidance_scale * (eps - eps_uncond) # DDIM update alpha_bar = self.alpha_bars[t] z_0 = (z - torch.sqrt(1 - alpha_bar) * eps) / torch.sqrt(alpha_bar) if t > 0: alpha_bar_next = self.alpha_bars[t - self.T // num_steps] z = torch.sqrt(alpha_bar_next) * z_0 + torch.sqrt(1 - alpha_bar_next) * eps else: z = z_0 # 4. Decode to image image = self.decode_latent(z) return image # Stable Diffusion pipeline (simplified) class StableDiffusion: """Stable Diffusion pipeline.""" def __init__(self, model_id="runwayml/stable-diffusion-v1-5"): from diffusers import StableDiffusionPipeline self.pipe = StableDiffusionPipeline.from_pretrained( model_id, torch_dtype=torch.float16 ).to("cuda") def generate(self, prompt, num_inference_steps=50, guidance_scale=7.5, height=512, width=512, seed=None): """Generate image from prompt.""" if seed is not None: torch.manual_seed(seed) result = self.pipe( prompt=prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale, height=height, width=width ) return result.images[0] # Usage example # sd = StableDiffusion() # image = sd.generate("A majestic castle on a hill, sunset, oil painting") # image.save("castle.png") print("Stable Diffusion: Text β†’ Latent β†’ Image") print("Key components:") print(" 1. VAE: Compresses 512Γ—512Γ—3 β†’ 64Γ—64Γ—4") print(" 2. U-Net: Denoises latent with text conditioning") print(" 3. Text Encoder: CLIP encodes prompt") print(" 4. Scheduler: DDIM/DPM-Solver for fast sampling")

🎯 Conditioning (Text, Class, Image)

🎯 Controlling Generation

Conditioning allows us to control what the model generates β€” text prompts, class labels, reference images, depth maps, poses. This is what makes diffusion models so powerful and controllable.

Conditioning Types

Class conditioning: Generate specific classes (e.g., "cat"). Text conditioning: CLIP text embeddings. Image conditioning: img2img, inpainting. Structural: Depth, edge, pose. Cross-attention: How text guides U-Net. ControlNet: Additional conditioning.

πŸ“ Text🏷️ ClassπŸ–ΌοΈ Image🦴 Pose
# Conditioning techniques for diffusion models import torch import torch.nn as nn import torch.nn.functional as F # 1. Class conditioning (embedding lookup) class ClassConditionedUNet(nn.Module): def __init__(self, num_classes=10, class_emb_dim=128): super().__init__() self.class_embed = nn.Embedding(num_classes, class_emb_dim) self.class_proj = nn.Linear(class_emb_dim, 256) # ... rest of U-Net def forward(self, x, t, class_labels): # Add class embedding to time embedding class_emb = self.class_proj(self.class_embed(class_labels)) # Combine with time embedding # ... return x # 2. Text conditioning with CLIP class TextConditionedUNet(nn.Module): def __init__(self, clip_model_name='openai/clip-vit-base-patch32'): super().__init__() from transformers import CLIPTextModel, CLIPTokenizer self.text_encoder = CLIPTextModel.from_pretrained(clip_model_name) self.tokenizer = CLIPTokenizer.from_pretrained(clip_model_name) # Projection to U-Net dimension self.text_proj = nn.Linear(512, 768) def encode_text(self, prompts): """Encode text prompts.""" tokens = self.tokenizer( prompts, padding=True, truncation=True, return_tensors='pt', max_length=77 ) outputs = self.text_encoder(**tokens) return self.text_proj(outputs.last_hidden_state) # 3. Cross-attention conditioning class CrossAttentionBlock(nn.Module): """Cross-attention for text-to-image.""" def __init__(self, query_dim=320, context_dim=768, heads=8): super().__init__() self.attn = nn.MultiheadAttention( embed_dim=query_dim, num_heads=heads, batch_first=True ) self.context_proj = nn.Linear(context_dim, query_dim) self.norm = nn.LayerNorm(query_dim) def forward(self, x, context): """ x: (B, C, H, W) β€” image features context: (B, L, context_dim) β€” text embeddings """ B, C, H, W = x.shape # Flatten spatial dims x_flat = x.view(B, C, H * W).permute(0, 2, 1) # (B, HW, C) # Project context context = self.context_proj(context) # (B, L, C) # Cross-attention: image attends to text x_norm = self.norm(x_flat) attn_out, _ = self.attn(x_norm, context, context) # Residual x_flat = x_flat + attn_out # Reshape back x = x_flat.permute(0, 2, 1).view(B, C, H, W) return x # 4. Classifier-free guidance class ClassifierFreeGuidance: """Classifier-free guidance (CFG) for conditional generation.""" def __init__(self, model, guidance_scale=7.5): self.model = model self.guidance_scale = guidance_scale def predict_noise(self, x_t, t, condition): """Predict noise with CFG.""" # Conditional prediction eps_cond = self.model(x_t, t, condition) # Unconditional prediction (empty condition) eps_uncond = self.model(x_t, t, None) # Combine eps = eps_uncond + self.guidance_scale * (eps_cond - eps_uncond) return eps # 5. ControlNet (spatial conditioning) class ControlNet(nn.Module): """ControlNet for additional spatial conditioning.""" def __init__(self, base_unet, control_channels=3): super().__init__() # Copy of base U-Net encoder self.control_encoder = self.copy_encoder(base_unet) self.control_input = nn.Conv2d(control_channels, 64, 3, padding=1) self.zero_convs = nn.ModuleList([ nn.Conv2d(64, 64, 1) for _ in range(12) ]) def forward(self, x, t, control_image, text_emb): """ x: noisy latent control_image: edge map, depth, pose, etc. """ # Encode control image control = self.control_input(control_image) # Process through control encoder (with text conditioning) # ... # Output residual connections to base U-Net return residuals # 6. Image-to-image (img2img) class Img2Img: """Image-to-image translation with diffusion.""" def __init__(self, diffusion_model): self.model = diffusion_model @torch.no_grad() def generate(self, input_image, prompt, strength=0.75): """ input_image: source image strength: how much to change (0=no change, 1=complete generation) """ # 1. Encode input to latent z_0 = self.model.encode_image(input_image) # 2. Add noise (up to strength * T) start_step = int(self.model.T * strength) t_start = torch.full((1,), start_step, device=z_0.device, dtype=torch.long) alpha_bar = self.model.alpha_bars[start_step] noise = torch.randn_like(z_0) z_start = torch.sqrt(alpha_bar) * z_0 + torch.sqrt(1 - alpha_bar) * noise # 3. Denoise from start_step z = z_start for t in reversed(range(start_step)): eps = self.model.unet(z, t, prompt) # DDIM update # ... # 4. Decode return self.model.decode_latent(z) # 7. Inpainting class Inpainting: """Inpainting with diffusion models.""" def __init__(self, diffusion_model): self.model = diffusion_model @torch.no_grad() def generate(self, image, mask, prompt): """ image: source image mask: 1 = keep original, 0 = generate new """ z = self.model.encode_image(image) for t in reversed(range(self.model.T)): eps = self.model.unet(z, t, prompt) # Update # ... # Force known region to stay close to original z_known = self.model.encode_image(image) z = mask * z_known + (1 - mask) * z return self.model.decode_latent(z) # 8. Conditioning types summary CONDITIONING_TYPES = { 'text': 'CLIP text embeddings via cross-attention', 'class': 'Class embedding added to time embedding', 'image': 'Concatenate or cross-attend with reference', 'depth': 'ControlNet with depth maps', 'edge': 'ControlNet with Canny edges', 'pose': 'ControlNet with OpenPose skeletons', 'segmentation': 'Semantic maps for structure', 'style': 'IP-Adapter or reference images', 'composition': 'Multi-ControlNet for complex control', } for name, desc in CONDITIONING_TYPES.items(): print(f"{name:<15}: {desc}")

πŸ”€ Classifier-Free Guidance

🎯 Better Conditioning Without a Classifier

Classifier-Free Guidance (CFG) is a technique that improves conditional generation by interpolating between conditional and unconditional predictions. It's the standard in modern diffusion models.

How CFG Works

Formula: Ξ΅_guided = Ξ΅_uncond + w Γ— (Ξ΅_cond - Ξ΅_uncond). w: Guidance scale (typically 7.5). Training: Randomly drop conditioning (10% of time). Effect: Higher w = stronger adherence to prompt. Trade-off: Too high w = oversaturated, artifacts. No classifier: Unlike classifier guidance, no separate model needed.

🎯 Guidance ScaleπŸ”„ DropoutπŸ“ˆ Prompt Adherence✨ Quality
# Classifier-Free Guidance β€” detailed implementation import torch import torch.nn as nn import torch.nn.functional as F class CFGTrainer: """Training with classifier-free guidance.""" def __init__(self, model, T=1000, uncond_prob=0.1): self.model = model self.T = T self.uncond_prob = uncond_prob # probability of dropping condition self.betas = torch.linspace(1e-4, 0.02, T) self.alphas = 1 - self.betas self.alpha_bars = torch.cumprod(self.alphas, dim=0) def train_step(self, x_0, condition, device='cuda'): """One training step with CFG.""" batch_size = x_0.shape[0] # 1. Randomly drop condition drop_mask = torch.rand(batch_size, device=device) < self.uncond_prob condition_dropped = condition.clone() condition_dropped[drop_mask] = 0 # null token # 2. Sample timestep t = torch.randint(0, self.T, (batch_size,), device=device) # 3. Add noise alpha_bar = self.alpha_bars[t].to(device).view(-1, 1, 1, 1) noise = torch.randn_like(x_0) x_t = torch.sqrt(alpha_bar) * x_0 + torch.sqrt(1 - alpha_bar) * noise # 4. Predict noise (with dropped condition) eps_pred = self.model(x_t, t, condition_dropped) # 5. MSE loss loss = F.mse_loss(eps_pred, noise) return loss class CFGSampler: """Sampling with classifier-free guidance.""" def __init__(self, model, T=1000, guidance_scale=7.5): self.model = model self.T = T self.guidance_scale = guidance_scale self.betas = torch.linspace(1e-4, 0.02, T) self.alphas = 1 - self.betas self.alpha_bars = torch.cumprod(self.alphas, dim=0) @torch.no_grad() def sample(self, shape, condition, device='cuda', num_steps=50): """Sample with CFG.""" x = torch.randn(shape, device=device) step_size = self.T // num_steps timesteps = list(range(0, self.T, step_size))[::-1] for t in timesteps: t_batch = torch.full((shape[0],), t, device=device, dtype=torch.long) # 1. Conditional prediction eps_cond = self.model(x, t_batch, condition) # 2. Unconditional prediction (null condition) null_condition = torch.zeros_like(condition) eps_uncond = self.model(x, t_batch, null_condition) # 3. CFG combination eps = eps_uncond + self.guidance_scale * (eps_cond - eps_uncond) # 4. DDIM update alpha_bar = self.alpha_bars[t].to(device) x_0_pred = (x - torch.sqrt(1 - alpha_bar) * eps) / torch.sqrt(alpha_bar) x_0_pred = x_0_pred.clamp(-1, 1) if t > 0: t_next = t - step_size alpha_bar_next = self.alpha_bars[t_next].to(device) if t_next >= 0 else torch.tensor(1.0, device=device) x = torch.sqrt(alpha_bar_next) * x_0_pred + torch.sqrt(1 - alpha_bar_next) * eps else: x = x_0_pred return x # Effect of guidance scale def analyze_guidance(model, prompt, device='cuda'): """Analyze effect of guidance scale.""" scales = [1.0, 3.0, 7.5, 15.0, 30.0] results = {} for w in scales: sampler = CFGSampler(model, guidance_scale=w) x = sampler.sample((1, 3, 64, 64), prompt, device) results[w] = x print(f"w = {w:5.1f}: ", end="") if w < 2: print("Weak prompt adherence, diverse") elif w < 5: print("Moderate adherence, good diversity") elif w < 10: print("Strong adherence, standard quality") # sweet spot elif w < 20: print("Very strong adherence, some saturation") else: print("Oversaturated, artifacts") return results # Advanced: dynamic CFG class DynamicCFG: """Dynamic guidance scale over timesteps.""" def __init__(self, model, T=1000): self.model = model self.T = T def get_guidance_scale(self, t): """Guidance scale varies with timestep.""" # Higher guidance early (structure), lower late (details) progress = 1 - t / self.T return 5.0 + 5.0 * progress # 5 β†’ 10 # Rescaled CFG (prevents oversaturation) class RescaledCFG: """Rescaled classifier-free guidance.""" def __init__(self, model, guidance_scale=7.5, rescale=0.7): self.model = model self.guidance_scale = guidance_scale self.rescale = rescale def predict_noise(self, x_t, t, condition): eps_cond = self.model(x_t, t, condition) eps_uncond = self.model(x_t, t, None) eps_guided = eps_uncond + self.guidance_scale * (eps_cond - eps_uncond) # Rescale to prevent oversaturation std_cond = eps_cond.std() std_guided = eps_guided.std() eps_rescaled = eps_guided * (std_cond / std_guided) * self.rescale + \ eps_guided * (1 - self.rescale) return eps_rescaled # Perp-neg (perpendicular negative guidance) class PerpNegCFG: """Perpendicular negative guidance β€” alternative to CFG.""" def __init__(self, model, guidance_scale=7.5): self.model = model self.guidance_scale = guidance_scale def predict_noise(self, x_t, t, condition): eps_cond = self.model(x_t, t, condition) eps_uncond = self.model(x_t, t, None) # Project out component parallel to unconditional v = eps_cond - eps_uncond v_parallel = (v * eps_uncond).sum() / (eps_uncond ** 2).sum() * eps_uncond v_perp = v - v_parallel return eps_uncond + self.guidance_scale * v_perp print("Classifier-Free Guidance β€” Key Points:") print(" 1. Train with random condition dropout (10%)") print(" 2. At inference: combine cond + uncond predictions") print(" 3. Guidance scale w controls prompt adherence") print(" 4. Typical w: 7.5 (Stable Diffusion default)") print(" 5. No separate classifier needed")

πŸ–ΌοΈ Applications (Image, Video, Audio)

🎯 Beyond Images

Diffusion models are being applied to nearly every generative task: images, video, audio, 3D, molecules, text, and more. The same principles extend across modalities.

Application Areas

Image: Text-to-image, inpainting, super-resolution. Video: Text-to-video, video prediction. Audio: Text-to-speech, music generation. 3D: Text-to-3D, NeRF, Gaussian Splatting. Science: Protein design, drug discovery. Code: Code generation with diffusion.

πŸ–ΌοΈ Image🎬 Video🎡 Audio🧬 Science
# Diffusion applications across modalities import torch import torch.nn as nn # 1. IMAGE: Text-to-Image (Stable Diffusion) class TextToImage: """Generate images from text prompts.""" def __init__(self): from diffusers import StableDiffusionPipeline self.pipe = StableDiffusionPipeline.from_pretrained( "runwayml/stable-diffusion-v1-5", torch_dtype=torch.float16 ).to("cuda") def generate(self, prompt, negative_prompt="", steps=50, guidance=7.5): return self.pipe( prompt=prompt, negative_prompt=negative_prompt, num_inference_steps=steps, guidance_scale=guidance ).images[0] # 2. IMAGE: Super-resolution class SuperResolution: """Upscale images with diffusion.""" def __init__(self): from diffusers import StableDiffusionUpscalePipeline self.pipe = StableDiffusionUpscalePipeline.from_pretrained( "stabilityai/stable-diffusion-x4-upscaler" ).to("cuda") def upscale(self, image, prompt): return self.pipe(prompt=prompt, image=image).images[0] # 3. IMAGE: Inpainting class InpaintingPipeline: """Fill in missing parts of images.""" def __init__(self): from diffusers import StableDiffusionInpaintPipeline self.pipe = StableDiffusionInpaintPipeline.from_pretrained( "runwayml/stable-diffusion-inpainting" ).to("cuda") def inpaint(self, image, mask, prompt): return self.pipe(prompt=prompt, image=image, mask_image=mask).images[0] # 4. VIDEO: Text-to-Video class TextToVideo: """Generate videos from text.""" def __init__(self): # Example: AnimateDiff, Sora, Stable Video Diffusion pass def generate(self, prompt, num_frames=16): # 1. Generate initial image with text-to-image # 2. Use temporal diffusion for video # 3. Add motion modules pass # 5. AUDIO: Text-to-Speech class TextToSpeech: """Generate speech from text.""" def __init__(self): # Example: AudioLDM, Bark, VALL-E pass def generate(self, text): # Diffusion in mel-spectrogram space # Convert to waveform pass # 6. AUDIO: Music Generation class MusicGeneration: """Generate music with diffusion.""" def __init__(self): pass def generate(self, prompt, duration=30): # Generate audio in latent space # Example: MusicLM, AudioLDM pass # 7. 3D: Text-to-3D class TextTo3D: """Generate 3D models from text.""" def __init__(self): # Example: DreamFusion, Magic3D, Shap-E pass def generate(self, prompt): # 1. Use 2D diffusion as prior # 2. Optimize 3D representation (NeRF, mesh) # 3. Score distillation sampling pass # 8. SCIENCE: Protein Design class ProteinDesign: """Design proteins with diffusion.""" def __init__(self): # Example: RFdiffusion, Chroma pass def generate(self, length, constraints): # Diffusion over protein backbone coordinates pass # 9. CODE: Code Generation class CodeDiffusion: """Generate code with diffusion models.""" def __init__(self): pass def generate(self, prompt): # Diffusion over token embeddings # Example: DiffuCode pass # 10. MOLECULE: Drug Discovery class MoleculeGeneration: """Generate molecules with diffusion.""" def __init__(self): # Example: EDM (Equivariant Diffusion Model) pass def generate(self, num_atoms, constraints): # Diffusion over 3D molecular coordinates pass # Comparison of diffusion frameworks FRAMEWORKS = { 'Diffusers': 'HuggingFace β€” most popular, comprehensive', 'Stable Diffusion WebUI': 'Community UI with extensions', 'ComfyUI': 'Node-based workflow interface', 'InvokeAI': 'Professional creative tool', 'Fooocus': 'Simplified Stable Diffusion', 'Kohya SS': 'Training and fine-tuning', 'ControlNet': 'Spatial conditioning', 'LoRA': 'Low-rank adaptation for fine-tuning', } print("=== Diffusion Applications ===") for area in ['Image', 'Video', 'Audio', '3D', 'Science', 'Code']: print(f"\n{area}:") if area == 'Image': print(" β€’ Text-to-image (Stable Diffusion, DALL-E 2, Imagen)") print(" β€’ Inpainting, outpainting") print(" β€’ Super-resolution") print(" β€’ Image editing (InstructPix2Pix)") elif area == 'Video': print(" β€’ Text-to-video (Sora, Runway, Pika)") print(" β€’ Video prediction") print(" β€’ Video editing") elif area == 'Audio': print(" β€’ Text-to-speech (AudioLDM, Bark)") print(" β€’ Music generation (MusicLM)") print(" β€’ Audio inpainting") elif area == '3D': print(" β€’ Text-to-3D (DreamFusion, Shap-E)") print(" β€’ Image-to-3D (Zero-1-to-3)") print(" β€’ 3D scene generation") elif area == 'Science': print(" β€’ Protein design (RFdiffusion)") print(" β€’ Drug discovery (EDM)") print(" β€’ Material generation") elif area == 'Code': print(" β€’ Code generation (DiffuCode)") print(" β€’ Code infilling") print(" β€’ Test generation")

⚠️ Challenges & Limitations

🎯 What's Still Hard

Diffusion models face challenges: slow sampling, high compute, training instability, mode collapse, ethical concerns, and evaluation.

Key Challenges

Speed: 20-1000 steps per sample. Compute: Massive GPU requirements. Memory: High VRAM for high-res. Training: Weeks on multiple GPUs. Evaluation: No perfect metric. Ethics: Deepfakes, copyright, bias. Control: Precise control is still hard.

🐒 SlowπŸ’» Computeβš–οΈ EthicsπŸ“Š Evaluation
# Challenges and solutions in diffusion models import torch import torch.nn as nn # 1. SLOW SAMPLING β†’ Solutions class FastSamplingSolutions: """Techniques to speed up diffusion sampling.""" @staticmethod def ddim(num_steps=50): """DDIM: deterministic, 20-100 steps.""" return f"DDIM with {num_steps} steps" @staticmethod def dpm_solver(num_steps=20): """DPM-Solver: ODE-based, 10-20 steps.""" return f"DPM-Solver with {num_steps} steps" @staticmethod def distillation(): """Progressive distillation: train student to match teacher.""" return "Progressive distillation: 4-8 steps" @staticmethod def consistency_models(): """Consistency models: single-step generation.""" return "Consistency models: 1-2 steps" @staticmethod def latent_diffusion(): """Latent diffusion: 4-8x smaller space.""" return "Latent diffusion: 8x faster" # 2. HIGH COMPUTE β†’ Solutions class ComputeSolutions: """Reduce computational requirements.""" @staticmethod def mixed_precision(): """FP16/BF16 training and inference.""" return "Mixed precision: 2x memory savings" @staticmethod def gradient_checkpointing(): """Trade compute for memory.""" return "Gradient checkpointing: 4x memory savings" @staticmethod def lora(): """Low-Rank Adaptation for fine-tuning.""" return "LoRA: 1000x fewer trainable params" @staticmethod def quantization(): """8-bit/4-bit quantization.""" return "Quantization: 4x memory savings" # 3. TRAINING INSTABILITY β†’ Solutions class StabilitySolutions: """Techniques for stable training.""" @staticmethod def ema(): """Exponential Moving Average of weights.""" class EMA: def __init__(self, model, decay=0.9999): self.model = model self.decay = decay self.shadow = {k: v.clone() for k, v in model.state_dict().items()} def update(self): for k, v in self.model.state_dict().items(): self.shadow[k] = self.decay * self.shadow[k] + (1 - self.decay) * v @staticmethod def min_snr(): """Min-SNR weighting for loss.""" return "Min-SNR: balances learning across timesteps" @staticmethod def v_prediction(): """v-prediction for better stability.""" return "v-prediction: more stable at high noise" # 4. ETHICS β†’ Considerations class EthicalConsiderations: """Ethical issues with diffusion models.""" @staticmethod def deepfakes(): return { 'issue': 'Realistic fake images/videos', 'mitigation': 'Watermarking, detection, regulation' } @staticmethod def copyright(): return { 'issue': 'Training on copyrighted data', 'mitigation': 'Licensed data, opt-out, attribution' } @staticmethod def bias(): return { 'issue': 'Reinforcing stereotypes', 'mitigation': 'Diverse training, bias audits' } @staticmethod def misinformation(): return { 'issue': 'Fake news, propaganda', 'mitigation': 'Provenance, fact-checking, education' } @staticmethod def consent(): return { 'issue': 'Non-consensual imagery', 'mitigation': 'Consent requirements, filters' } # 5. EVALUATION β†’ Metrics class EvaluationMetrics: """Metrics for evaluating diffusion models.""" @staticmethod def fid(): """FrΓ©chet Inception Distance.""" return "FID: lower is better, measures realism + diversity" @staticmethod def is_score(): """Inception Score.""" return "IS: higher is better, measures quality + diversity" @staticmethod def clip_score(): """CLIP Score: text-image alignment.""" return "CLIP Score: measures prompt adherence" @staticmethod def human_eval(): """Human evaluation.""" return "Human eval: gold standard but expensive" # 6. CURRENT LIMITATIONS LIMITATIONS = { 'Sampling speed': '20-1000 steps; slow for real-time', 'Compute cost': 'Requires powerful GPUs', 'Training data': 'Needs billions of image-text pairs', 'Fine control': 'Hard to control exactly', 'Consistency': 'Different seeds give different outputs', 'Long-range coherence': 'Video and long-form struggle', 'Physics': 'Doesn\'t understand physics', 'Text rendering': 'Struggles with text in images', 'Hands/anatomy': 'Common failure mode', 'Counting': 'Can\'t reliably count objects', } print("=== Diffusion Model Challenges ===") for challenge, desc in LIMITATIONS.items(): print(f" β€’ {challenge}: {desc}") # 7. SOLUTIONS IN PROGRESS print("\n=== Active Research Areas ===") solutions = [ "Consistency models (1-step generation)", "Adversarial distillation (GAN + diffusion)", "Progressive distillation", "Rectified flow (straight ODE paths)", "Flow matching (simulation-free training)", "Consistency trajectory models", "Latent consistency models (4-step)", "Turbo/lightning models (1-4 steps)", ] for s in solutions: print(f" β€’ {s}")

πŸš€ Future: Unified Generative AI

🎯 Where Diffusion is Heading

The future of diffusion models: real-time generation, unified multimodal, world models, controllable, efficient, and integrated with LLMs.

Future Directions

Real-time: 1-step generation (LCM, SDXL-Turbo). Video: Sora-style long video generation. 3D: Native 3D diffusion. Multimodal: Any-to-any generation. World models: Interactive environments. LLM integration: Unified text + image. On-device: Mobile diffusion.

⚑ Real-time🎬 Video🧊 3D🌍 World Models
# Future directions in diffusion models import torch import torch.nn as nn # 1. LATENT CONSISTENCY MODELS (Real-time) class LatentConsistencyModel: """LCM: 1-4 step generation.""" def __init__(self, teacher_model, student_model): self.teacher = teacher_model self.student = student_model def distill(self, x_0, prompt): """Distill teacher to student for fast sampling.""" # Student learns to predict x_0 directly from x_t # Enables 1-4 step generation pass @torch.no_grad() def generate(self, prompt, num_steps=4): """Generate in 4 steps.""" # 4 steps instead of 50! pass # 2. RECTIFIED FLOW (Straight paths) class RectifiedFlow: """Rectified flow: straight ODE paths for fast sampling.""" def __init__(self, model): self.model = model def train(self, x_0, x_1): """Learn to transport x_0 to x_1 linearly.""" # x_t = (1-t) * x_0 + t * x_1 # Model predicts velocity pass @torch.no_grad() def sample(self, noise, num_steps=10): """Straight path allows fewer steps.""" pass # 3. FLOW MATCHING (Simulation-free) class FlowMatching: """Flow matching: efficient training without simulation.""" def __init__(self, model): self.model = model def train(self, x_0, x_1): """Match the flow between distributions.""" pass # 4. VIDEO DIFFUSION (Sora-style) class VideoDiffusion: """Long video generation.""" def __init__(self): # Patch-based video representation # Temporal attention # Long context modeling pass def generate(self, prompt, duration_seconds=60): """Generate long videos.""" # 1. Encode text # 2. Generate in latent space # 3. Maintain temporal coherence pass # 5. NATIVE 3D DIFFUSION class Native3DDiffusion: """Direct 3D generation.""" def __init__(self): # 3D-aware diffusion # Multi-view consistency pass def generate(self, prompt): """Generate 3D directly.""" # 1. Diffusion in 3D latent space # 2. Neural radiance fields # 3. Gaussian splatting pass # 6. UNIFIED MULTIMODAL class UnifiedMultimodalDiffusion: """Any-to-any generation.""" def __init__(self): # Single model for all modalities # Text, image, audio, video, 3D pass def generate(self, input_modality, output_modality, content): """Convert between any modalities.""" # Text β†’ Image # Image β†’ Video # Audio β†’ Image # etc. pass # 7. WORLD MODELS class WorldModel: """Interactive environment generation.""" def __init__(self): # Diffusion-based world model # Predicts future states pass def step(self, action): """Generate next state given action.""" # For game AI, robotics, simulation pass # 8. ON-DEVICE DIFFUSION class OnDeviceDiffusion: """Run diffusion on mobile/edge.""" def __init__(self): # Quantization # Pruning # Knowledge distillation pass def generate(self, prompt): """Generate on-device.""" # Snapdragon, Apple Neural Engine pass # FUTURE TIMELINE FUTURE_TIMELINE = { '2024': [ 'Real-time video generation', '1-step image generation', 'Native 3D diffusion', 'On-device diffusion', ], '2025': [ 'Unified multimodal models', 'Interactive world models', 'Professional video (Sora-level)', 'Diffusion + LLM integration', ], '2026+': [ 'Real-time 3D world generation', 'Full multimodal understanding', 'Personalized on-device models', 'AGI with generative capabilities', ], } print("=== Future of Diffusion Models ===") for year, milestones in FUTURE_TIMELINE.items(): print(f"\n{year}:") for m in milestones: print(f" β€’ {m}") # KEY TRENDS print("\n=== Key Trends ===") trends = [ "Speed: 1000 steps β†’ 50 β†’ 4 β†’ 1 step", "Resolution: 512 β†’ 1024 β†’ 4K β†’ 8K", "Length: Single image β†’ 16 frames β†’ 60s video", "Modalities: Images β†’ Video β†’ Audio β†’ 3D β†’ Any", "Size: Server β†’ Desktop β†’ Mobile β†’ On-chip", "Control: Text β†’ Text+Image β†’ Multi-condition", "Quality: Human-level β†’ Superhuman", ] for t in trends: print(f" {t}") # IMPACT print("\n=== Impact ===") print(" β€’ Content creation: Democratized") print(" β€’ Design: AI-assisted workflows") print(" β€’ Entertainment: Personalized content") print(" β€’ Science: Protein/drug discovery") print(" β€’ Education: Visual learning") print(" β€’ Accessibility: Text-to-speech, image description")