π§ 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.
# 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.
# 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.
# 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.
# 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.
# 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.
# 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.
# 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.
# 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.
# 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.
# 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.
# 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.
# 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.
# 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.
# 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")