Diffusion Models

import torch
import torch.nn.functional as F

# 1. Noise schedule (β_t) — thường dùng linear hoặc cosine
T = 1000
betas = torch.linspace(1e-4, 0.02, T)  # linear schedule
alphas = 1.0 - betas
alphas_cumprod = torch.cumprod(alphas, dim=0)  # ᾱ_t

def sample_noisy(x0, t, noise):
    """Thêm nhiễu tới bước t bằng closed-form (không cần loop)."""
    a_bar = alphas_cumprod[t].view(-1, 1, 1, 1)
    return a_bar.sqrt() * x0 + (1 - a_bar).sqrt() * noise

# 2. Training step
def train_step(unet, x0, optimizer):
    batch = x0.size(0)
    # Lấy t ngẫu nhiên cho mỗi ảnh trong batch
    t = torch.randint(0, T, (batch,), device=x0.device)
    # ε ~ N(0, I)
    noise = torch.randn_like(x0)
    # Tạo x_t
    x_t = sample_noisy(x0, t, noise)
    # U-Net dự đoán nhiễu
    noise_pred = unet(x_t, t)
    # MSE loss giữa ε thực và ε dự đoán
    loss = F.mse_loss(noise_pred, noise)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    return loss.item()

# 3. Vòng lặp training
for epoch in range(n_epochs):
    for x0 in dataloader:
        loss = train_step(unet, x0.to(device), optimizer)

@torch.no_grad()
def ddim_sample(unet, shape, timesteps=50, eta=0.0, device="cuda"):
    """DDIM sampling: non-Markovian, deterministic khi eta=0."""
    x = torch.randn(shape, device=device)

    # Chọn 50 bước rải đều trong [0, T)
    step_ratio = T // timesteps
    ddim_steps = torch.arange(0, T, step_ratio).flip(0)

    for i, t in enumerate(ddim_steps):
        t_batch = torch.full((shape[0],), t, device=device)
        t_prev = ddim_steps[i + 1] if i + 1 < len(ddim_steps) else -1

        # U-Net predict noise
        eps = unet(x, t_batch)

        a_t = alphas_cumprod[t]
        a_prev = alphas_cumprod[t_prev] if t_prev >= 0 else torch.tensor(1.0)

        # Dự đoán x_0
        x0_pred = (x - (1 - a_t).sqrt() * eps) / a_t.sqrt()

        # Hệ số stochastic (eta=0 → fully deterministic)
        sigma = eta * ((1 - a_prev) / (1 - a_t)).sqrt() * (1 - a_t / a_prev).sqrt()
        direction = (1 - a_prev - sigma ** 2).sqrt() * eps
        noise = sigma * torch.randn_like(x) if eta > 0 else 0.0

        # Cập nhật x theo DDIM formula
        x = a_prev.sqrt() * x0_pred + direction + noise

    return x  # ảnh sinh ra

import torch
from diffusers import StableDiffusionPipeline, DDIMScheduler

# Load pipeline Stable Diffusion 1.5
pipe = StableDiffusionPipeline.from_pretrained(
    "runwayml/stable-diffusion-v1-5",
    torch_dtype=torch.float16,
).to("cuda")

# Thay scheduler mặc định bằng DDIM để deterministic
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)

# Inference với CFG scale + seed cố định
generator = torch.Generator("cuda").manual_seed(42)
image = pipe(
    prompt="a Vietnamese pho bowl, studio lighting, 8k photography",
    negative_prompt="blurry, low quality, deformed",
    num_inference_steps=30,       # DDIM chỉ cần 20-50 steps
    guidance_scale=7.5,            # CFG sweet spot
    generator=generator,
).images[0]

image.save("pho.png")

# --- CFG bên trong thực chất làm gì? (pseudo-code) ---
# for t in timesteps:
#     # 1. Duplicate latent: [cond, uncond]
#     latent_in = torch.cat([latent] * 2)
#     # 2. Text embedding: [cond_emb, uncond_emb]
#     text_emb = torch.cat([cond_emb, uncond_emb])
#     # 3. U-Net forward một lần cho cả 2
#     eps_cond, eps_uncond = unet(latent_in, t, text_emb).chunk(2)
#     # 4. Guided noise
#     eps_hat = eps_uncond + guidance_scale * (eps_cond - eps_uncond)
#     # 5. Scheduler step
#     latent = scheduler.step(eps_hat, t, latent).prev_sample

class StableDiffusion(nn.Module):
    def __init__(self):
        super().__init__()
        # 1. VAE: nén và giải nén ảnh
        self.vae = AutoencoderKL(
            in_channels=3, out_channels=3,
            latent_channels=4,   # 4 kênh latent
            downsample_factor=8, # 512 → 64
        )
        # 2. Text encoder: prompt → embedding
        self.text_encoder = CLIPTextModel.from_pretrained("clip-vit-large")
        # 3. U-Net: denoise latent có điều kiện text
        self.unet = UNet2DConditionModel(
            sample_size=64,          # latent resolution
            in_channels=4,           # latent channels
            cross_attention_dim=768, # khớp CLIP-L
        )

    def encode_image(self, image):
        """512×512×3 → 64×64×4 latent."""
        latent = self.vae.encode(image).latent_dist.sample()
        return latent * 0.18215  # scale factor của SD 1.5

    def decode_latent(self, latent):
        """64×64×4 → 512×512×3 ảnh."""
        return self.vae.decode(latent / 0.18215).sample

    def forward(self, latent, t, prompt_emb):
        """Predict noise từ latent có điều kiện prompt embedding."""
        return self.unet(latent, t, encoder_hidden_states=prompt_emb).sample

    @torch.no_grad()
    def generate(self, prompt, steps=50, cfg=7.5):
        # Encode prompt + empty prompt
        cond = self.text_encoder(tokenize(prompt))
        uncond = self.text_encoder(tokenize(""))
        # Random latent noise
        z = torch.randn(1, 4, 64, 64, device="cuda")
        # Denoise loop (DDIM)
        for t in reversed(range(0, 1000, 1000 // steps)):
            eps_c = self.forward(z, t, cond)
            eps_u = self.forward(z, t, uncond)
            eps = eps_u + cfg * (eps_c - eps_u)
            z = ddim_step(z, eps, t)
        # Decode latent → image
        return self.decode_latent(z)