Upload inference-only latent diffusion model

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,63 @@
+---
+license: mit
+tags:
+- text-to-image
+- diffusion
+- latent-diffusion
+- pytorch
+- coco
+library_name: pytorch
+pipeline_tag: text-to-image
+---
+
+# anilegin/lightweight-diffusion-ldm
+
+Custom lightweight latent diffusion text-to-image model.
+
+This repository contains inference-only files:
+
+- VAE config and stripped VAE weights
+- LDM/UNet config and stripped LDM weights
+- diffusion/sampler code needed for DDPM and DDIM
+- a simple `inference.py` script
+- generation defaults in `generation_config.yaml`
+
+The checkpoints are stripped to contain model weights only; optimizer state, scheduler state, and training logs are not included.
+
+## Install
+
+```bash
+git clone https://huggingface.co/anilegin/lightweight-diffusion-ldm
+cd lightweight-diffusion-ldm
+pip install -r requirements.txt
+```
+
+## Generate images
+
+```bash
+python inference.py \
+  --prompt "a small dog sitting on a red couch" \
+  --sampler ddim \
+  --num-steps 50 \
+  --guidance-scale 3.0 \
+  --precision bf16 \
+  --output-dir outputs/example
+```
+
+For offline/local-only CLIP loading, make sure `openai/clip-vit-large-patch14` is cached locally and add:
+
+```bash
+--local-files-only
+```
+
+## Notes
+
+This is a custom PyTorch implementation, not a native Diffusers pipeline. The included source code is required for inference.
+
+## Training data
+
+Trained/evaluated with COCO-style image-caption data. Add more precise dataset, metrics, and limitations here before making the repo public.
+
+## Citation
+
+If you use this model, please cite the project/repository.

checkpoints/ldm_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:79a241aa2d3f67e89d8d9227e5022efc36d603ddf07c426c599b78304dd668b7
+size 1798549690

checkpoints/vae_model.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:179bfe367d29e1e726638b71a8a9dfc17651d9b368d4fade5029022c89df5b15
+size 426519034

configs/ldm_config.yaml

@@ -0,0 +1,102 @@
+project:
+  root: .
+data:
+  coco_root: /leonardo_scratch/large/userexternal/aegin000/datasets/coco2017
+  train_latent_dir: outputs/latents/coco_train2017_vae8_scaled032_allcaptions
+  val_latent_dir: outputs/latents/coco_val2017_vae8_scaled032_allcaptions
+  train_caption_mode: random
+  val_caption_mode: first
+  load_latents_to_memory: true
+outputs:
+  root: outputs
+  vae_dir: outputs/vae
+  ldm_dir: outputs/ldm
+  latent_dir: outputs/latents
+  sample_dir: outputs/samples
+  log_dir: logs
+cache:
+  root: cache
+model:
+  name: latent_diffusion_unet_strong
+  in_channels: 8
+  out_channels: 8
+  latent_size: 32
+  base_channels: 304
+  channel_multipliers:
+  - 1
+  - 2
+  - 4
+  num_res_blocks: 3
+  dropout: 0.0
+  attention_resolutions:
+  - 16
+  - 8
+  use_middle_attention: true
+  context_dim: 768
+  num_heads: 8
+  head_dim: 64
+  transformer_depth: 1
+  time_embedding_dim: null
+experiment:
+  name: ldm_coco_256_vae8_strong_vpred_ft
+text_encoder:
+  model_name: openai/clip-vit-large-patch14
+  max_length: 77
+  freeze: true
+  use_last_hidden_state: true
+conditioning:
+  cond_drop_prob: 0.05
+  empty_text: ''
+diffusion:
+  schedule_type: cosine
+  num_timesteps: 1000
+  prediction_type: v
+  loss_type: mse
+  beta_start: 0.0001
+  beta_end: 0.02
+  cosine_s: 0.008
+  max_beta: 0.999
+  snr_weighting: min_snr
+  snr_gamma: 5.0
+  normalize_snr_weights: false
+train:
+  seed: 42
+  precision: bf16
+  batch_size: 2
+  gradient_accumulation_steps: 12
+  num_workers: 2
+  pin_memory: true
+  max_epochs: 100
+  lr: 0.0001
+  min_lr: 5.0e-05
+  scheduler: cosine
+  warmup_steps: 2000
+  weight_decay: 0.01
+  betas:
+  - 0.9
+  - 0.999
+  grad_clip: 1.0
+  log_every: 50
+  validate_every: 1
+  save_every: 1
+  output_dir_key: outputs.ldm_dir
+  resume_from: null
+  finetune_from: outputs/ldm/ldm_coco_256_vae8_strong_vpred_ft/checkpoints/last.pt
+validation:
+  enabled: true
+  max_batches: 100
+  bucket_max_batches: 25
+  timestep_buckets:
+  - - 0
+    - 100
+  - - 100
+    - 300
+  - - 300
+    - 700
+  - - 700
+    - 1000
+optimizer:
+  name: adamw
+distributed:
+  enabled: true
+  backend: nccl

configs/vae_config.yaml

@@ -0,0 +1,78 @@
+project:
+  root: .
+data:
+  coco_root: /leonardo_scratch/large/userexternal/aegin000/datasets/coco2017
+outputs:
+  root: outputs
+  vae_dir: outputs/vae
+  ldm_dir: outputs/ldm
+  latent_dir: outputs/latents
+  sample_dir: outputs/samples
+  log_dir: logs
+cache:
+  root: cache
+dataset:
+  name: coco_captions
+  root_key: data.coco_root
+  train_split: train2017
+  val_split: val2017
+  resolution: 256
+  num_workers: 8
+  pin_memory: true
+model:
+  name: autoencoder_kl
+  in_channels: 3
+  out_channels: 3
+  latent_channels: 8
+  base_channels: 128
+  channel_multipliers:
+  - 1
+  - 2
+  - 4
+  - 4
+  num_res_blocks: 3
+  dropout: 0.0
+  use_attention: true
+  attention_heads: 4
+  scaling_factor: 1.0
+attention_resolutions:
+- 32
+experiment:
+  name: vae_coco_256_small
+train:
+  seed: 42
+  precision: bf16
+  batch_size: 16
+  gradient_accumulation_steps: 4
+  num_workers: 8
+  max_epochs: 100
+  lr: 0.0001
+  weight_decay: 0.0
+  betas:
+  - 0.9
+  - 0.999
+  grad_clip: 1.0
+  log_every: 100
+  validate_every: 1
+  save_every: 1
+  sample_every: 1
+  num_sample_images: 8
+  output_dir_key: outputs.vae_dir
+  initialize_from_scratch: true
+  resume_from: outputs/vae/vae_coco_256_small/checkpoints/last.pt
+  finetune_from: null
+early_stopping:
+  enabled: true
+  patience: 15
+  min_delta: 0.0
+  monitor_metric: val_total_loss
+loss:
+  recon_loss_type: l1
+  recon_weight: 1.0
+  use_lpips: true
+  lpips_net: vgg
+  perceptual_weight: 0.1
+  kl_weight: 1.0e-06
+  kl_warmup_steps: 10000
+optimizer:
+  name: adamw

generation_config.yaml

@@ -0,0 +1,20 @@
+ldm:
+  config: configs/ldm_config.yaml
+  checkpoint: checkpoints/ldm_model.pt
+vae:
+  config: configs/vae_config.yaml
+  checkpoint: checkpoints/vae_model.pt
+  scaling_factor: 1.032
+conditioning:
+  empty_text: ''
+generation:
+  seed: 42
+  precision: fp16
+  resolution: 256
+  batch_size: 4
+sampler:
+  type: ddim
+  num_steps: 50
+  eta: 0.0
+  guidance_scale: 3.0
+  clip_denoised: false

inference.py

@@ -0,0 +1,353 @@
+from __future__ import annotations
+
+import argparse
+import math
+import random
+import sys
+import warnings
+from pathlib import Path
+from typing import Any
+
+import torch
+import yaml
+from PIL import Image
+from tqdm import tqdm
+
+sys.path.append(str(Path(__file__).resolve().parent))
+
+from src.diffusion.gaussian_diffusion import GaussianDiffusion
+from src.diffusion.samplers import DDPMSampler, DDIMSampler
+from src.models.autoencoder.vae import AutoencoderKL
+from src.models.conditioning.clip_text import FrozenCLIPTextEncoder
+from src.models.diffusion.unet import build_latent_diffusion_unet_from_config
+
+
+def load_yaml(path: str | Path) -> dict[str, Any]:
+    with open(path, "r", encoding="utf-8") as f:
+        return yaml.safe_load(f)
+
+
+def safe_torch_load(path: str | Path, map_location="cpu"):
+    try:
+        return torch.load(path, map_location=map_location, weights_only=True)
+    except TypeError:
+        return torch.load(path, map_location=map_location)
+    except Exception:
+        return torch.load(path, map_location=map_location)
+
+
+def get_dtype(name: str) -> torch.dtype:
+    name = name.lower()
+    if name == "fp16":
+        return torch.float16
+    if name == "bf16":
+        return torch.bfloat16
+    if name == "fp32":
+        return torch.float32
+    raise ValueError(f"Unknown precision={name}")
+
+
+def autocast_context(device: torch.device, dtype: torch.dtype):
+    enabled = device.type == "cuda" and dtype in (torch.float16, torch.bfloat16)
+    if device.type == "cuda":
+        return torch.autocast("cuda", dtype=dtype, enabled=enabled)
+    return torch.autocast("cpu", enabled=False)
+
+
+def set_seed(seed: int):
+    random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def sanitize_filename(text: str, max_len: int = 80) -> str:
+    text = text.lower().strip()
+    keep = []
+    for ch in text:
+        if ch.isalnum():
+            keep.append(ch)
+        elif ch in {" ", "-", "_"}:
+            keep.append("_")
+    out = "".join(keep)
+    while "__" in out:
+        out = out.replace("__", "_")
+    out = out.strip("_") or "sample"
+    return out[:max_len]
+
+
+def save_image_tensor(image: torch.Tensor, path: str | Path):
+    image = image.detach().cpu().clamp(0.0, 1.0)
+    image = image.permute(1, 2, 0).float().numpy()
+    image = (image * 255).round().astype("uint8")
+    Image.fromarray(image).save(path)
+
+
+def save_image_grid(images: list[Image.Image], path: str | Path, nrow: int | None = None, padding: int = 2):
+    if not images:
+        return
+
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+
+    if nrow is None:
+        nrow = int(math.ceil(math.sqrt(len(images))))
+    ncol = int(math.ceil(len(images) / nrow))
+
+    widths, heights = zip(*(img.size for img in images))
+    cell_w, cell_h = max(widths), max(heights)
+    grid_w = nrow * cell_w + padding * (nrow - 1)
+    grid_h = ncol * cell_h + padding * (ncol - 1)
+
+    grid = Image.new("RGB", (grid_w, grid_h), color=(255, 255, 255))
+    for idx, img in enumerate(images):
+        row = idx // nrow
+        col = idx % nrow
+        x = col * (cell_w + padding)
+        y = row * (cell_h + padding)
+        grid.paste(img.convert("RGB"), (x, y))
+    grid.save(path)
+
+
+def load_model_state(module: torch.nn.Module, checkpoint_path: str | Path):
+    checkpoint = safe_torch_load(checkpoint_path, map_location="cpu")
+    if isinstance(checkpoint, dict):
+        state_dict = checkpoint.get("model", checkpoint.get("state_dict", checkpoint))
+    else:
+        state_dict = checkpoint
+    module.load_state_dict(state_dict, strict=True)
+
+
+def build_vae(vae_cfg: dict) -> AutoencoderKL:
+    model_cfg = dict(vae_cfg["model"])
+    model_cfg.pop("name", None)
+    return AutoencoderKL(**model_cfg)
+
+
+def build_diffusion(ldm_cfg: dict) -> GaussianDiffusion:
+    d = ldm_cfg["diffusion"]
+    return GaussianDiffusion(
+        schedule_type=str(d.get("schedule_type", "cosine")),
+        num_timesteps=int(d.get("num_timesteps", 1000)),
+        prediction_type=str(d.get("prediction_type", "v")),
+        loss_type=str(d.get("loss_type", "mse")),
+        beta_start=float(d.get("beta_start", 1e-4)),
+        beta_end=float(d.get("beta_end", 2e-2)),
+        cosine_s=float(d.get("cosine_s", 0.008)),
+        max_beta=float(d.get("max_beta", 0.999)),
+    )
+
+
+def build_text_encoder(ldm_cfg: dict, device: torch.device, local_files_only: bool):
+    warnings.filterwarnings("ignore", message=".*clean_up_tokenization_spaces.*", category=FutureWarning)
+
+    text_cfg = dict(ldm_cfg.get("text_encoder", {}))
+    text_encoder = FrozenCLIPTextEncoder(
+        model_name=str(text_cfg.get("model_name", "openai/clip-vit-large-patch14")),
+        max_length=int(text_cfg.get("max_length", 77)),
+        freeze=True,
+        use_last_hidden_state=bool(text_cfg.get("use_last_hidden_state", True)),
+        local_files_only=local_files_only,
+    )
+    text_encoder.to(device=device)
+    text_encoder.eval()
+    return text_encoder
+
+
+@torch.no_grad()
+def encode_contexts(text_encoder, prompts: list[str], empty_text: str, device: torch.device, dtype: torch.dtype):
+    context_dtype = dtype if device.type == "cuda" else torch.float32
+    cond_context = text_encoder.encode(prompts, device=device).to(dtype=context_dtype)
+    uncond_context = text_encoder.encode([empty_text] * len(prompts), device=device).to(dtype=context_dtype)
+    return cond_context, uncond_context
+
+
+@torch.no_grad()
+def decode_latents(vae, latents: torch.Tensor, scaling_factor: float, dtype: torch.dtype):
+    try:
+        with autocast_context(latents.device, dtype):
+            images = vae.decode(latents.float(), unscale=True)
+    except TypeError:
+        z = latents.float() / scaling_factor
+        with autocast_context(latents.device, dtype):
+            images = vae.decode(z)
+
+    if hasattr(images, "sample"):
+        images = images.sample
+
+    return ((images + 1.0) / 2.0).clamp(0.0, 1.0)
+
+
+def read_prompts(args) -> list[str]:
+    prompts: list[str] = []
+
+    if args.prompt is not None:
+        prompts.append(args.prompt)
+
+    if args.prompts_file is not None:
+        with open(args.prompts_file, "r", encoding="utf-8") as f:
+            prompts.extend([line.strip() for line in f if line.strip()])
+
+    if not prompts:
+        raise ValueError("Provide --prompt or --prompts-file.")
+
+    repeated = []
+    for prompt in prompts:
+        repeated.extend([prompt] * args.num_images_per_prompt)
+    return repeated
+
+
+@torch.no_grad()
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--config", type=str, default="generation_config.yaml")
+    parser.add_argument("--prompt", type=str, default=None)
+    parser.add_argument("--prompts-file", type=str, default=None)
+    parser.add_argument("--output-dir", type=str, default="outputs")
+    parser.add_argument("--sampler", type=str, default=None, choices=["ddim", "ddpm"])
+    parser.add_argument("--num-steps", type=int, default=None)
+    parser.add_argument("--guidance-scale", type=float, default=None)
+    parser.add_argument("--eta", type=float, default=None)
+    parser.add_argument("--precision", type=str, default=None, choices=["fp32", "bf16", "fp16"])
+    parser.add_argument("--seed", type=int, default=None)
+    parser.add_argument("--batch-size", type=int, default=None)
+    parser.add_argument("--num-images-per-prompt", type=int, default=1)
+    parser.add_argument("--local-files-only", action="store_true", help="Use only locally cached CLIP files.")
+    args = parser.parse_args()
+
+    repo_root = Path(__file__).resolve().parent
+    cfg = load_yaml(repo_root / args.config)
+
+    gen_cfg = cfg.get("generation", {})
+    sampler_cfg = cfg.get("sampler", {})
+
+    seed = int(args.seed if args.seed is not None else gen_cfg.get("seed", 42))
+    precision = str(args.precision if args.precision is not None else gen_cfg.get("precision", "bf16"))
+    batch_size = int(args.batch_size if args.batch_size is not None else gen_cfg.get("batch_size", 4))
+    sampler_name = str(args.sampler if args.sampler is not None else sampler_cfg.get("type", "ddim")).lower()
+    num_steps = int(args.num_steps if args.num_steps is not None else sampler_cfg.get("num_steps", 50))
+    guidance_scale = float(args.guidance_scale if args.guidance_scale is not None else sampler_cfg.get("guidance_scale", 3.0))
+    eta = float(args.eta if args.eta is not None else sampler_cfg.get("eta", 0.0))
+    clip_denoised = bool(sampler_cfg.get("clip_denoised", False))
+
+    set_seed(seed)
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    dtype = get_dtype(precision)
+
+    ldm_cfg_path = repo_root / cfg["ldm"]["config"]
+    ldm_ckpt_path = repo_root / cfg["ldm"]["checkpoint"]
+    vae_cfg_path = repo_root / cfg["vae"]["config"]
+    vae_ckpt_path = repo_root / cfg["vae"]["checkpoint"]
+
+    ldm_cfg = load_yaml(ldm_cfg_path)
+    vae_cfg = load_yaml(vae_cfg_path)
+
+    vae = build_vae(vae_cfg)
+    load_model_state(vae, vae_ckpt_path)
+    vae.to(device=device)
+    vae.eval()
+
+    unet = build_latent_diffusion_unet_from_config(ldm_cfg)
+    load_model_state(unet, ldm_ckpt_path)
+    unet.to(device=device)
+    unet.eval()
+
+    diffusion = build_diffusion(ldm_cfg).to(device)
+    text_encoder = build_text_encoder(ldm_cfg, device=device, local_files_only=args.local_files_only)
+
+    if sampler_name == "ddim":
+        sampler = DDIMSampler(diffusion)
+    elif sampler_name == "ddpm":
+        sampler = DDPMSampler(diffusion)
+    else:
+        raise ValueError(f"Unknown sampler: {sampler_name}")
+
+    latent_channels = int(ldm_cfg["model"].get("in_channels", 8))
+    image_size = int(gen_cfg.get("resolution", 256))
+    latent_size = int(ldm_cfg["model"].get("latent_size", image_size // 8))
+    scaling_factor = float(cfg["vae"].get("scaling_factor", getattr(vae, "scaling_factor", 1.0)))
+    empty_text = str(cfg.get("conditioning", {}).get("empty_text", ""))
+
+    prompts = read_prompts(args)
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    print("=============================================")
+    print("Custom latent diffusion inference")
+    print("Device:", device)
+    print("Precision:", precision)
+    print("Sampler:", sampler_name)
+    print("Steps:", num_steps if sampler_name == "ddim" else diffusion.num_timesteps)
+    print("Guidance scale:", guidance_scale)
+    print("Eta:", eta)
+    print("Seed:", seed)
+    print("Total images:", len(prompts))
+    print("Output dir:", output_dir)
+    print("=============================================")
+
+    saved_images: list[Image.Image] = []
+
+    for batch_start in tqdm(range(0, len(prompts), batch_size), desc="batches"):
+        batch_prompts = prompts[batch_start: batch_start + batch_size]
+        cond_context, uncond_context = encode_contexts(
+            text_encoder=text_encoder,
+            prompts=batch_prompts,
+            empty_text=empty_text,
+            device=device,
+            dtype=dtype,
+        )
+
+        shape = (len(batch_prompts), latent_channels, latent_size, latent_size)
+
+        with autocast_context(device, dtype):
+            if sampler_name == "ddim":
+                sample_out = sampler.sample(
+                    model=unet,
+                    shape=shape,
+                    device=device,
+                    context=cond_context,
+                    attention_mask=None,
+                    uncond_context=uncond_context,
+                    uncond_attention_mask=None,
+                    guidance_scale=guidance_scale,
+                    num_steps=num_steps,
+                    eta=eta,
+                    clip_denoised=clip_denoised,
+                    return_trajectory=False,
+                    progress=True,
+                )
+            else:
+                sample_out = sampler.sample(
+                    model=unet,
+                    shape=shape,
+                    device=device,
+                    context=cond_context,
+                    attention_mask=None,
+                    uncond_context=uncond_context,
+                    uncond_attention_mask=None,
+                    guidance_scale=guidance_scale,
+                    clip_denoised=clip_denoised,
+                    return_trajectory=False,
+                    progress=True,
+                )
+
+        latents = sample_out.latents if hasattr(sample_out, "latents") else sample_out
+        images = decode_latents(vae, latents, scaling_factor=scaling_factor, dtype=dtype)
+
+        for i, image in enumerate(images):
+            prompt = batch_prompts[i]
+            global_idx = batch_start + i
+            out_path = output_dir / f"{global_idx:04d}_{sanitize_filename(prompt)}.png"
+            save_image_tensor(image, out_path)
+
+            pil = Image.open(out_path).convert("RGB")
+            saved_images.append(pil)
+
+    save_image_grid(saved_images, output_dir / "grid.png")
+    print("Saved", len(saved_images), "images to", output_dir)
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,8 @@
+torch
+torchvision
+transformers
+huggingface_hub
+python-dotenv
+PyYAML
+Pillow
+tqdm

src/diffusion/gaussian_diffusion.py

@@ -0,0 +1,437 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+
+from src.losses.diffusion_loss import DiffusionLoss
+from src.diffusion.noise_schedule import NoiseSchedule, create_noise_schedule, extract
+from src.diffusion.prediction import (
+    get_training_target,
+    model_output_to_x0_and_eps,
+)
+
+
+@dataclass
+class DiffusionTrainingOutput:
+    loss: torch.Tensor
+    simple_loss: torch.Tensor
+    model_output: torch.Tensor
+    target: torch.Tensor
+    z_t: torch.Tensor
+    noise: torch.Tensor
+    timesteps: torch.Tensor
+
+
+class GaussianDiffusion:
+    """
+    Core latent diffusion utilities.
+
+    This handles:
+
+        - sampling timesteps
+        - adding noise q(z_t | z_0)
+        - creating v-prediction targets
+        - computing diffusion training loss
+        - computing DDPM posterior mean/variance for sampling
+    """
+
+    def __init__(
+        self,
+        schedule: NoiseSchedule | None = None,
+        schedule_type: str = "cosine",
+        num_timesteps: int = 1000,
+        prediction_type: str = "v",
+        loss_type: str = "mse",
+        beta_start: float = 1e-4,
+        beta_end: float = 2e-2,
+        cosine_s: float = 0.008,
+        max_beta: float = 0.999,
+        snr_gamma: float | None = None,
+        snr_weighting: str = "none",
+        normalize_snr_weights: bool = False,
+    ):
+        if schedule is None:
+            schedule = create_noise_schedule(
+                schedule_type=schedule_type,
+                num_timesteps=num_timesteps,
+                beta_start=beta_start,
+                beta_end=beta_end,
+                cosine_s=cosine_s,
+                max_beta=max_beta,
+            )
+
+        self.schedule = schedule
+        self.prediction_type = prediction_type.lower()
+        self.loss_type = loss_type.lower()
+        self.snr_gamma = snr_gamma
+        self.snr_weighting = snr_weighting.lower()
+        self.normalize_snr_weights = normalize_snr_weights
+
+        if self.prediction_type not in {"v", "v_prediction", "eps", "epsilon", "x0", "sample"}:
+            raise ValueError(
+                f"Unknown prediction_type={prediction_type}. "
+                "Use 'v', 'eps', or 'x0'."
+            )
+
+        if self.loss_type not in {"mse", "l1", "huber"}:
+            raise ValueError(
+                f"Unknown loss_type={loss_type}. "
+                "Use 'mse', 'l1', or 'huber'."
+            )
+
+        self.diffusion_loss = DiffusionLoss(
+            prediction_type=self.prediction_type,
+            loss_type=self.loss_type,
+            snr_gamma=self.snr_gamma,
+            snr_weighting=self.snr_weighting,
+            normalize_snr_weights=self.normalize_snr_weights,
+        )
+
+    @property
+    def num_timesteps(self) -> int:
+        return self.schedule.num_timesteps
+
+    def to(self, device: torch.device | str) -> "GaussianDiffusion":
+        self.schedule = self.schedule.to(device)
+        return self
+
+    def sample_timesteps(
+        self,
+        batch_size: int,
+        device: torch.device | str,
+    ) -> torch.Tensor:
+        """
+        Sample random diffusion timesteps.
+
+        Returns:
+            t: [B], values in [0, num_timesteps - 1]
+        """
+        return torch.randint(
+            low=0,
+            high=self.num_timesteps,
+            size=(batch_size,),
+            device=device,
+            dtype=torch.long,
+        )
+
+    def q_sample(
+        self,
+        z_0: torch.Tensor,
+        t: torch.Tensor,
+        noise: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Forward diffusion process:
+
+            q(z_t | z_0)
+
+        Formula:
+
+            z_t = sqrt(alpha_bar_t) * z_0
+                + sqrt(1 - alpha_bar_t) * eps
+
+        Args:
+            z_0:
+                Clean latent [B, C, H, W].
+
+            t:
+                Timesteps [B].
+
+            noise:
+                Optional epsilon noise. If None, sampled from N(0, I).
+
+        Returns:
+            z_t:
+                Noisy latent.
+
+            noise:
+                The epsilon noise used.
+        """
+        if noise is None:
+            noise = torch.randn_like(z_0)
+
+        sqrt_alpha_bar = extract(
+            self.schedule.sqrt_alphas_cumprod,
+            t,
+            z_0.shape,
+        )
+
+        sqrt_one_minus_alpha_bar = extract(
+            self.schedule.sqrt_one_minus_alphas_cumprod,
+            t,
+            z_0.shape,
+        )
+
+        z_t = sqrt_alpha_bar * z_0 + sqrt_one_minus_alpha_bar * noise
+
+        return z_t, noise
+
+    def training_target(
+        self,
+        z_0: torch.Tensor,
+        noise: torch.Tensor,
+        t: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Get target for current prediction type
+        """
+        return get_training_target(
+            z_0=z_0,
+            eps=noise,
+            t=t,
+            schedule=self.schedule,
+            prediction_type=self.prediction_type,
+        )
+
+    def p_losses(
+        self,
+        model,
+        z_0: torch.Tensor,
+        context: torch.Tensor | None = None,
+        t: torch.Tensor | None = None,
+        noise: torch.Tensor | None = None,
+        model_kwargs: dict | None = None,
+    ) -> DiffusionTrainingOutput:
+        """
+        Full diffusion training step using  loss module.
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+
+        batch_size = z_0.shape[0]
+        device = z_0.device
+
+        if t is None:
+            t = self.sample_timesteps(batch_size, device)
+
+        z_t, noise = self.q_sample(
+            z_0=z_0,
+            t=t,
+            noise=noise,
+        )
+
+        target = self.training_target(
+            z_0=z_0,
+            noise=noise,
+            t=t,
+        )
+
+        if context is None:
+            model_output = model(z_t, t, **model_kwargs)
+        else:
+            model_output = model(z_t, t, context=context, **model_kwargs)
+
+        alpha_t = extract(
+            self.schedule.sqrt_alphas_cumprod,
+            t,
+            z_0.shape,
+        )
+
+        sigma_t = extract(
+            self.schedule.sqrt_one_minus_alphas_cumprod,
+            t,
+            z_0.shape,
+        )
+
+        alpha_bar_t = self.schedule.alphas_cumprod.gather(
+            0,
+            t,
+        )
+
+        snr = alpha_bar_t / (1.0 - alpha_bar_t).clamp(min=1e-8)
+
+        loss_out = self.diffusion_loss(
+            model_output=model_output,
+            x0=z_0,
+            noise=noise,
+            alpha_t=alpha_t,
+            sigma_t=sigma_t,
+            snr=snr,
+            return_dict=True,
+        )
+
+        loss = loss_out["loss"]
+        raw_loss = loss_out["raw_loss"]
+
+        return DiffusionTrainingOutput(
+            loss=loss,
+            simple_loss=raw_loss.detach(),
+            model_output=model_output,
+            target=target,
+            z_t=z_t,
+            noise=noise,
+            timesteps=t,
+        )
+
+    def predict_x0_and_eps(
+        self,
+        model_output: torch.Tensor,
+        z_t: torch.Tensor,
+        t: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Convert model output to:
+
+            z_0 prediction
+            epsilon prediction
+        """
+        return model_output_to_x0_and_eps(
+            model_output=model_output,
+            z_t=z_t,
+            t=t,
+            schedule=self.schedule,
+            prediction_type=self.prediction_type,
+        )
+
+    def q_posterior(
+        self,
+        z_0: torch.Tensor,
+        z_t: torch.Tensor,
+        t: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Compute posterior:
+
+            q(z_{t-1} | z_t, z_0)
+
+        Returns:
+            posterior_mean
+            posterior_variance
+            posterior_log_variance_clipped
+        """
+        posterior_mean_coef1 = extract(
+            self.schedule.posterior_mean_coef1,
+            t,
+            z_t.shape,
+        )
+
+        posterior_mean_coef2 = extract(
+            self.schedule.posterior_mean_coef2,
+            t,
+            z_t.shape,
+        )
+
+        posterior_mean = (
+            posterior_mean_coef1 * z_0
+            + posterior_mean_coef2 * z_t
+        )
+
+        posterior_variance = extract(
+            self.schedule.posterior_variance,
+            t,
+            z_t.shape,
+        )
+
+        posterior_log_variance_clipped = extract(
+            self.schedule.posterior_log_variance_clipped,
+            t,
+            z_t.shape,
+        )
+
+        return (
+            posterior_mean,
+            posterior_variance,
+            posterior_log_variance_clipped,
+        )
+
+    @torch.no_grad()
+    def p_mean_variance(
+        self,
+        model,
+        z_t: torch.Tensor,
+        t: torch.Tensor,
+        context: torch.Tensor | None = None,
+        clip_denoised: bool = False,
+        model_kwargs: dict | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """
+        One reverse-process prediction.
+
+        Model predicts v/eps/x0.
+        We convert to predicted z_0
+        """
+        if model_kwargs is None:
+            model_kwargs = {}
+
+        if context is None:
+            model_output = model(
+                z_t,
+                t,
+                **model_kwargs,
+            )
+        else:
+            model_output = model(
+                z_t,
+                t,
+                context=context,
+                **model_kwargs,
+            )
+
+        pred_z0, pred_eps = self.predict_x0_and_eps(
+            model_output=model_output,
+            z_t=z_t,
+            t=t,
+        )
+
+        if clip_denoised:
+            pred_z0 = pred_z0.clamp(-1.0, 1.0)
+
+        (
+            posterior_mean,
+            posterior_variance,
+            posterior_log_variance,
+        ) = self.q_posterior(
+            z_0=pred_z0,
+            z_t=z_t,
+            t=t,
+        )
+
+        return {
+            "mean": posterior_mean,
+            "variance": posterior_variance,
+            "log_variance": posterior_log_variance,
+            "pred_z0": pred_z0,
+            "pred_eps": pred_eps,
+            "model_output": model_output,
+        }
+
+    @torch.no_grad()
+    def p_sample(
+        self,
+        model,
+        z_t: torch.Tensor,
+        t: torch.Tensor,
+        context: torch.Tensor | None = None,
+        clip_denoised: bool = False,
+        model_kwargs: dict | None = None,
+    ) -> torch.Tensor:
+        """
+        This is one reverse step
+        """
+        out = self.p_mean_variance(
+            model=model,
+            z_t=z_t,
+            t=t,
+            context=context,
+            clip_denoised=clip_denoised,
+            model_kwargs=model_kwargs,
+        )
+
+        noise = torch.randn_like(z_t)
+
+        # No noise when t == 0.
+        nonzero_mask = (t != 0).float()
+
+        while len(nonzero_mask.shape) < len(z_t.shape):
+            nonzero_mask = nonzero_mask[..., None]
+
+        z_prev = (
+            out["mean"]
+            + nonzero_mask
+            * torch.exp(0.5 * out["log_variance"])
+            * noise
+        )
+
+        return z_prev

src/diffusion/noise_schedule.py

@@ -0,0 +1,302 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import torch
+
+
+@dataclass
+class NoiseSchedule:
+    """
+    Precomputed DDPM noise schedule.
+
+    Main variables:
+
+        beta_t:
+            amount of noise added at timestep t
+
+        alpha_t:
+            1 - beta_t
+
+        alpha_bar_t:
+            cumulative product of alphas up to t
+
+        q(z_t | z_0):
+            z_t = sqrt(alpha_bar_t) * z_0
+                + sqrt(1 - alpha_bar_t) * eps
+    """
+
+    betas: torch.Tensor
+    alphas: torch.Tensor
+    alphas_cumprod: torch.Tensor
+    alphas_cumprod_prev: torch.Tensor
+
+    sqrt_alphas_cumprod: torch.Tensor
+    sqrt_one_minus_alphas_cumprod: torch.Tensor
+
+    log_one_minus_alphas_cumprod: torch.Tensor
+
+    sqrt_recip_alphas_cumprod: torch.Tensor
+    sqrt_recipm1_alphas_cumprod: torch.Tensor
+
+    posterior_variance: torch.Tensor
+    posterior_log_variance_clipped: torch.Tensor
+    posterior_mean_coef1: torch.Tensor
+    posterior_mean_coef2: torch.Tensor
+
+    num_timesteps: int
+    schedule_type: str
+
+    def to(self, device: torch.device | str) -> "NoiseSchedule":
+        device = torch.device(device)
+
+        return NoiseSchedule(
+            betas=self.betas.to(device),
+            alphas=self.alphas.to(device),
+            alphas_cumprod=self.alphas_cumprod.to(device),
+            alphas_cumprod_prev=self.alphas_cumprod_prev.to(device),
+            sqrt_alphas_cumprod=self.sqrt_alphas_cumprod.to(device),
+            sqrt_one_minus_alphas_cumprod=self.sqrt_one_minus_alphas_cumprod.to(device),
+            log_one_minus_alphas_cumprod=self.log_one_minus_alphas_cumprod.to(device),
+            sqrt_recip_alphas_cumprod=self.sqrt_recip_alphas_cumprod.to(device),
+            sqrt_recipm1_alphas_cumprod=self.sqrt_recipm1_alphas_cumprod.to(device),
+            posterior_variance=self.posterior_variance.to(device),
+            posterior_log_variance_clipped=self.posterior_log_variance_clipped.to(device),
+            posterior_mean_coef1=self.posterior_mean_coef1.to(device),
+            posterior_mean_coef2=self.posterior_mean_coef2.to(device),
+            num_timesteps=self.num_timesteps,
+            schedule_type=self.schedule_type,
+        )
+
+
+def make_beta_schedule(
+    schedule_type: str = "cosine",
+    num_timesteps: int = 1000,
+    beta_start: float = 1e-4,
+    beta_end: float = 2e-2,
+    cosine_s: float = 0.008,
+    max_beta: float = 0.999,
+) -> torch.Tensor:
+    """
+    Create beta schedule.
+
+    Supported:
+        linear:
+            Standard DDPM linear beta schedule.
+
+        cosine:
+            Improved DDPM cosine schedule.
+            Usually better behaved and good default for v-prediction.
+
+    Returns:
+        betas: [num_timesteps], float32
+    """
+    schedule_type = schedule_type.lower()
+
+    if schedule_type == "linear":
+        betas = torch.linspace(
+            beta_start,
+            beta_end,
+            num_timesteps,
+            dtype=torch.float64,
+        )
+
+    elif schedule_type == "cosine":
+        betas = cosine_beta_schedule(
+            num_timesteps=num_timesteps,
+            cosine_s=cosine_s,
+            max_beta=max_beta,
+        )
+
+    else:
+        raise ValueError(
+            f"Unknown schedule_type={schedule_type}. "
+            "Use 'linear' or 'cosine'."
+        )
+
+    return betas.float()
+
+
+def cosine_beta_schedule(
+    num_timesteps: int,
+    cosine_s: float = 0.008,
+    max_beta: float = 0.999,
+) -> torch.Tensor:
+    """
+    Cosine beta schedule
+    Instead of directly defining beta_t, we define alpha_bar(t)
+    using a cosine curve, then derive beta_t.
+    """
+    steps = num_timesteps + 1
+
+    x = torch.linspace(
+        0,
+        num_timesteps,
+        steps,
+        dtype=torch.float64,
+    )
+
+    alphas_cumprod = torch.cos(
+        ((x / num_timesteps) + cosine_s)
+        / (1.0 + cosine_s)
+        * math.pi
+        * 0.5
+    ) ** 2
+
+    alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
+
+    betas = 1.0 - (
+        alphas_cumprod[1:] / alphas_cumprod[:-1]
+    )
+
+    betas = torch.clamp(
+        betas,
+        min=1e-8,
+        max=max_beta,
+    )
+
+    return betas
+
+
+def create_noise_schedule(
+    schedule_type: str = "cosine",
+    num_timesteps: int = 1000,
+    beta_start: float = 1e-4,
+    beta_end: float = 2e-2,
+    cosine_s: float = 0.008,
+    max_beta: float = 0.999,
+) -> NoiseSchedule:
+    """
+    all precomputed schedule tensors needed for DDPM training and sampling.
+    """
+    betas = make_beta_schedule(
+        schedule_type=schedule_type,
+        num_timesteps=num_timesteps,
+        beta_start=beta_start,
+        beta_end=beta_end,
+        cosine_s=cosine_s,
+        max_beta=max_beta,
+    )
+
+    alphas = 1.0 - betas
+
+    alphas_cumprod = torch.cumprod(
+        alphas,
+        dim=0,
+    )
+
+    alphas_cumprod_prev = torch.cat(
+        [
+            torch.ones(1, dtype=alphas_cumprod.dtype),
+            alphas_cumprod[:-1],
+        ],
+        dim=0,
+    )
+
+    sqrt_alphas_cumprod = torch.sqrt(alphas_cumprod)
+
+    sqrt_one_minus_alphas_cumprod = torch.sqrt(
+        1.0 - alphas_cumprod
+    )
+
+    log_one_minus_alphas_cumprod = torch.log(
+        torch.clamp(
+            1.0 - alphas_cumprod,
+            min=1e-20,
+        )
+    )
+
+    sqrt_recip_alphas_cumprod = torch.sqrt(
+        1.0 / alphas_cumprod
+    )
+
+    sqrt_recipm1_alphas_cumprod = torch.sqrt(
+        1.0 / alphas_cumprod - 1.0
+    )
+
+    # Posterior q(z_{t-1} | z_t, z_0)
+    posterior_variance = (
+        betas
+        * (1.0 - alphas_cumprod_prev)
+        / (1.0 - alphas_cumprod)
+    )
+
+    posterior_log_variance_clipped = torch.log(
+        torch.clamp(
+            posterior_variance,
+            min=1e-20,
+        )
+    )
+
+    posterior_mean_coef1 = (
+        betas
+        * torch.sqrt(alphas_cumprod_prev)
+        / (1.0 - alphas_cumprod)
+    )
+
+    posterior_mean_coef2 = (
+        (1.0 - alphas_cumprod_prev)
+        * torch.sqrt(alphas)
+        / (1.0 - alphas_cumprod)
+    )
+
+    return NoiseSchedule(
+        betas=betas,
+        alphas=alphas,
+        alphas_cumprod=alphas_cumprod,
+        alphas_cumprod_prev=alphas_cumprod_prev,
+        sqrt_alphas_cumprod=sqrt_alphas_cumprod,
+        sqrt_one_minus_alphas_cumprod=sqrt_one_minus_alphas_cumprod,
+        log_one_minus_alphas_cumprod=log_one_minus_alphas_cumprod,
+        sqrt_recip_alphas_cumprod=sqrt_recip_alphas_cumprod,
+        sqrt_recipm1_alphas_cumprod=sqrt_recipm1_alphas_cumprod,
+        posterior_variance=posterior_variance,
+        posterior_log_variance_clipped=posterior_log_variance_clipped,
+        posterior_mean_coef1=posterior_mean_coef1,
+        posterior_mean_coef2=posterior_mean_coef2,
+        num_timesteps=num_timesteps,
+        schedule_type=schedule_type,
+    )
+
+
+def extract(
+    values: torch.Tensor,
+    timesteps: torch.Tensor,
+    broadcast_shape: tuple[int, ...],
+) -> torch.Tensor:
+    """
+    Extract values[t] and reshape for broadcasting.
+
+    Args:
+        values:
+            Schedule tensor with shape [T].
+
+        timesteps:
+            Long tensor with shape [B].
+
+        broadcast_shape:
+            Shape of target tensor, e.g. z_t.shape = [B, C, H, W].
+
+    Returns:
+        Tensor with shape [B, 1, 1, 1], broadcastable to broadcast_shape.
+
+    Example:
+        sqrt_alpha_bar_t = extract(
+            schedule.sqrt_alphas_cumprod,
+            t,
+            z_0.shape,
+        )
+    """
+    if timesteps.dtype != torch.long:
+        timesteps = timesteps.long()
+
+    out = values.gather(
+        dim=0,
+        index=timesteps,
+    )
+
+    while len(out.shape) < len(broadcast_shape):
+        out = out[..., None]
+
+    return out

src/diffusion/prediction.py

@@ -0,0 +1,259 @@
+from __future__ import annotations
+
+import torch
+
+from src.diffusion.noise_schedule import NoiseSchedule, extract
+
+
+def predict_x0_from_eps(
+    z_t: torch.Tensor,
+    t: torch.Tensor,
+    eps: torch.Tensor,
+    schedule: NoiseSchedule,
+) -> torch.Tensor:
+    """
+    Recover clean latent z_0 from epsilon prediction
+    """
+    sqrt_recip_alpha_bar = extract(
+        schedule.sqrt_recip_alphas_cumprod,
+        t,
+        z_t.shape,
+    )
+
+    sqrt_recipm1_alpha_bar = extract(
+        schedule.sqrt_recipm1_alphas_cumprod,
+        t,
+        z_t.shape,
+    )
+
+    return sqrt_recip_alpha_bar * z_t - sqrt_recipm1_alpha_bar * eps
+
+
+def predict_eps_from_x0(
+    z_t: torch.Tensor,
+    t: torch.Tensor,
+    z_0: torch.Tensor,
+    schedule: NoiseSchedule,
+) -> torch.Tensor:
+    """
+    Recover epsilon from clean latent z_0 and noisy latent z_t.
+
+    eps = (z_t - sqrt(alpha_bar_t) * z_0)
+          / sqrt(1 - alpha_bar_t)
+    """
+    sqrt_alpha_bar = extract(
+        schedule.sqrt_alphas_cumprod,
+        t,
+        z_t.shape,
+    )
+
+    sqrt_one_minus_alpha_bar = extract(
+        schedule.sqrt_one_minus_alphas_cumprod,
+        t,
+        z_t.shape,
+    )
+
+    return (z_t - sqrt_alpha_bar * z_0) / sqrt_one_minus_alpha_bar
+
+
+def get_v_target(
+    z_0: torch.Tensor,
+    eps: torch.Tensor,
+    t: torch.Tensor,
+    schedule: NoiseSchedule,
+) -> torch.Tensor:
+    """
+    Compute v-prediction target.
+
+    v-prediction target:
+
+        v = sqrt(alpha_bar_t) * eps
+          - sqrt(1 - alpha_bar_t) * z_0
+    """
+    sqrt_alpha_bar = extract(
+        schedule.sqrt_alphas_cumprod,
+        t,
+        z_0.shape,
+    )
+
+    sqrt_one_minus_alpha_bar = extract(
+        schedule.sqrt_one_minus_alphas_cumprod,
+        t,
+        z_0.shape,
+    )
+
+    v = sqrt_alpha_bar * eps - sqrt_one_minus_alpha_bar * z_0
+
+    return v
+
+
+def predict_x0_from_v(
+    z_t: torch.Tensor,
+    t: torch.Tensor,
+    v: torch.Tensor,
+    schedule: NoiseSchedule,
+) -> torch.Tensor:
+    """
+    Recover clean latent z_0 from v prediction
+
+    defs:
+
+        z_t = a * z_0 + b * eps
+        v   = a * eps - b * z_0
+
+        a = sqrt(alpha_bar_t)
+        b = sqrt(1 - alpha_bar_t)
+
+    it gives
+
+        z_0 = a * z_t - b * v
+    """
+    sqrt_alpha_bar = extract(
+        schedule.sqrt_alphas_cumprod,
+        t,
+        z_t.shape,
+    )
+
+    sqrt_one_minus_alpha_bar = extract(
+        schedule.sqrt_one_minus_alphas_cumprod,
+        t,
+        z_t.shape,
+    )
+
+    z_0 = sqrt_alpha_bar * z_t - sqrt_one_minus_alpha_bar * v
+
+    return z_0
+
+
+def predict_eps_from_v(
+    z_t: torch.Tensor,
+    t: torch.Tensor,
+    v: torch.Tensor,
+    schedule: NoiseSchedule,
+) -> torch.Tensor:
+    """
+    Recover epsilon from v prediction
+
+    defs:
+
+        z_t = a * z_0 + b * eps
+        v   = a * eps - b * z_0
+
+    it gives
+
+        eps = b * z_t + a * v
+    """
+    sqrt_alpha_bar = extract(
+        schedule.sqrt_alphas_cumprod,
+        t,
+        z_t.shape,
+    )
+
+    sqrt_one_minus_alpha_bar = extract(
+        schedule.sqrt_one_minus_alphas_cumprod,
+        t,
+        z_t.shape,
+    )
+
+    eps = sqrt_one_minus_alpha_bar * z_t + sqrt_alpha_bar * v
+
+    return eps
+
+
+def model_output_to_x0_and_eps(
+    model_output: torch.Tensor,
+    z_t: torch.Tensor,
+    t: torch.Tensor,
+    schedule: NoiseSchedule,
+    prediction_type: str = "v",
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """
+    Convert model output into both:
+        z_0 prediction
+        epsilon prediction
+    """
+    prediction_type = prediction_type.lower()
+
+    if prediction_type in {"v", "v_prediction"}:
+        z_0 = predict_x0_from_v(
+            z_t=z_t,
+            t=t,
+            v=model_output,
+            schedule=schedule,
+        )
+
+        eps = predict_eps_from_v(
+            z_t=z_t,
+            t=t,
+            v=model_output,
+            schedule=schedule,
+        )
+
+    elif prediction_type in {"eps", "epsilon"}:
+        eps = model_output
+
+        z_0 = predict_x0_from_eps(
+            z_t=z_t,
+            t=t,
+            eps=eps,
+            schedule=schedule,
+        )
+
+    elif prediction_type in {"x0", "sample"}:
+        z_0 = model_output
+
+        eps = predict_eps_from_x0(
+            z_t=z_t,
+            t=t,
+            z_0=z_0,
+            schedule=schedule,
+        )
+
+    else:
+        raise ValueError(
+            f"Unknown prediction_type={prediction_type}. "
+            "Use 'v', 'eps', or 'x0'."
+        )
+
+    return z_0, eps
+
+
+def get_training_target(
+    z_0: torch.Tensor,
+    eps: torch.Tensor,
+    t: torch.Tensor,
+    schedule: NoiseSchedule,
+    prediction_type: str = "v",
+) -> torch.Tensor:
+    """
+    Return the target the U-Net should learn
+
+    For our project, default is:
+
+        prediction_type = "v"
+
+    Then target is:
+
+        v = sqrt(alpha_bar_t) * eps
+          - sqrt(1 - alpha_bar_t) * z_0
+    """
+    prediction_type = prediction_type.lower()
+
+    if prediction_type in {"v", "v_prediction"}:
+        return get_v_target(
+            z_0=z_0,
+            eps=eps,
+            t=t,
+            schedule=schedule,
+        )
+
+    if prediction_type in {"eps", "epsilon"}:
+        return eps
+
+    if prediction_type in {"x0", "sample"}:
+        return z_0
+
+    raise ValueError(
+        f"Unknown prediction_type={prediction_type}. "
+        "Use 'v', 'eps', or 'x0'."
+    )

src/diffusion/samplers/__init__.py

@@ -0,0 +1,9 @@
+from src.diffusion.samplers.ddpm import DDPMSampler, DDPMSamplerOutput
+from src.diffusion.samplers.ddim import DDIMSampler, DDIMSamplerOutput
+
+__all__ = [
+    "DDPMSampler",
+    "DDPMSamplerOutput",
+    "DDIMSampler",
+    "DDIMSamplerOutput",
+]

src/diffusion/samplers/ddim.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from tqdm import tqdm
+
+from src.diffusion.gaussian_diffusion import GaussianDiffusion
+
+
+@dataclass
+class DDIMSamplerOutput:
+    latents: torch.Tensor
+    trajectory: list[torch.Tensor] | None = None
+
+
+class DDIMSampler:
+    """
+    DDIM sampler.
+
+    eta controls stochasticity:
+
+        eta = 0.0 -> deterministic DDIM
+        eta > 0.0 -> more stochastic
+    """
+
+    def __init__(
+        self,
+        diffusion: GaussianDiffusion,
+    ):
+        self.diffusion = diffusion
+
+    def make_timesteps(
+        self,
+        num_steps: int,
+        device: torch.device | str,
+    ) -> torch.Tensor:
+        """
+        Select evenly spaced timesteps from the original diffusion schedule.
+
+        Example:
+            original T = 1000
+            num_steps = 50
+
+        returns 50 timesteps descending from high noise to low noise.
+        """
+        if num_steps > self.diffusion.num_timesteps:
+            raise ValueError(
+                f"num_steps={num_steps} cannot be larger than "
+                f"num_timesteps={self.diffusion.num_timesteps}"
+            )
+
+        timesteps = torch.linspace(
+            0,
+            self.diffusion.num_timesteps - 1,
+            steps=num_steps,
+            device=device,
+        ).long()
+
+        timesteps = torch.flip(timesteps, dims=[0])
+
+        return timesteps
+
+    @torch.no_grad()
+    def predict_model_output(
+        self,
+        model,
+        z_t: torch.Tensor,
+        t: torch.Tensor,
+        context: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        uncond_context: torch.Tensor | None = None,
+        uncond_attention_mask: torch.Tensor | None = None,
+        guidance_scale: float = 1.0,
+    ) -> torch.Tensor:
+        """
+        Predict v/eps/x0 with optional classifier-free guidance.
+        """
+        if uncond_context is None or guidance_scale == 1.0:
+            if context is None:
+                return model(
+                    z_t,
+                    t,
+                )
+
+            return model(
+                z_t,
+                t,
+                context=context,
+                attention_mask=attention_mask,
+            )
+
+        cond_output = model(
+            z_t,
+            t,
+            context=context,
+            attention_mask=attention_mask,
+        )
+
+        uncond_output = model(
+            z_t,
+            t,
+            context=uncond_context,
+            attention_mask=uncond_attention_mask,
+        )
+
+        return uncond_output + guidance_scale * (cond_output - uncond_output)
+
+    @torch.no_grad()
+    def sample(
+        self,
+        model,
+        shape: tuple[int, int, int, int],
+        device: torch.device | str,
+        context: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        uncond_context: torch.Tensor | None = None,
+        uncond_attention_mask: torch.Tensor | None = None,
+        guidance_scale: float = 1.0,
+        num_steps: int = 50,
+        eta: float = 0.0,
+        clip_denoised: bool = False,
+        return_trajectory: bool = False,
+        progress: bool = True,
+    ) -> DDIMSamplerOutput:
+        """
+        DDIM sampling.
+
+        Returns:
+            clean latent estimate z_0 at the final step.
+        """
+        device = torch.device(device)
+        model.eval()
+
+        z_t = torch.randn(
+            shape,
+            device=device,
+        )
+
+        trajectory = [] if return_trajectory else None
+
+        ddim_timesteps = self.make_timesteps(
+            num_steps=num_steps,
+            device=device,
+        )
+
+        if progress:
+            iterator = tqdm(
+                range(len(ddim_timesteps)),
+                desc=f"DDIM sampling ({num_steps} steps)",
+            )
+        else:
+            iterator = range(len(ddim_timesteps))
+
+        for i in iterator:
+            step = ddim_timesteps[i]
+
+            t = torch.full(
+                (shape[0],),
+                int(step.item()),
+                device=device,
+                dtype=torch.long,
+            )
+
+            if i == len(ddim_timesteps) - 1:
+                prev_step = torch.tensor(
+                    -1,
+                    device=device,
+                    dtype=torch.long,
+                )
+            else:
+                prev_step = ddim_timesteps[i + 1]
+
+            model_output = self.predict_model_output(
+                model=model,
+                z_t=z_t,
+                t=t,
+                context=context,
+                attention_mask=attention_mask,
+                uncond_context=uncond_context,
+                uncond_attention_mask=uncond_attention_mask,
+                guidance_scale=guidance_scale,
+            )
+
+            pred_z0, pred_eps = self.diffusion.predict_x0_and_eps(
+                model_output=model_output,
+                z_t=z_t,
+                t=t,
+            )
+
+            if clip_denoised:
+                pred_z0 = pred_z0.clamp(-1.0, 1.0)
+
+            alpha_t = self.diffusion.schedule.alphas_cumprod[t]
+            alpha_t = alpha_t.view(shape[0], 1, 1, 1)
+
+            if prev_step.item() < 0:
+                alpha_prev = torch.ones_like(alpha_t)
+            else:
+                alpha_prev = self.diffusion.schedule.alphas_cumprod[
+                    torch.full(
+                        (shape[0],),
+                        int(prev_step.item()),
+                        device=device,
+                        dtype=torch.long,
+                    )
+                ]
+                alpha_prev = alpha_prev.view(shape[0], 1, 1, 1)
+
+            sigma_t = eta * torch.sqrt(
+                (1.0 - alpha_prev)
+                / (1.0 - alpha_t)
+                * (1.0 - alpha_t / alpha_prev)
+            )
+
+            # Direction pointing to z_t.
+            dir_xt = torch.sqrt(
+                torch.clamp(
+                    1.0 - alpha_prev - sigma_t ** 2,
+                    min=0.0,
+                )
+            ) * pred_eps
+
+            noise = sigma_t * torch.randn_like(z_t)
+
+            z_t = torch.sqrt(alpha_prev) * pred_z0 + dir_xt + noise
+
+            if return_trajectory:
+                trajectory.append(z_t.detach().cpu())
+
+        return DDIMSamplerOutput(
+            latents=z_t,
+            trajectory=trajectory,
+        )

src/diffusion/samplers/ddpm.py

@@ -0,0 +1,235 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from tqdm import tqdm
+
+from src.diffusion.gaussian_diffusion import GaussianDiffusion
+
+
+@dataclass
+class DDPMSamplerOutput:
+    latents: torch.Tensor
+    trajectory: list[torch.Tensor] | None = None
+
+
+class DDPMSampler:
+    """
+    DDPM sampler.
+
+    This sampler uses the learned reverse process:
+
+        z_T ~ N(0, I)
+        z_T -> z_{T-1} -> ... -> z_0
+
+    Supports classifier-free guidance if both conditional and unconditional
+    context are provided.
+    """
+
+    def __init__(
+        self,
+        diffusion: GaussianDiffusion,
+    ):
+        self.diffusion = diffusion
+
+    @torch.no_grad()
+    def predict_model_output(
+        self,
+        model,
+        z_t: torch.Tensor,
+        t: torch.Tensor,
+        context: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        uncond_context: torch.Tensor | None = None,
+        uncond_attention_mask: torch.Tensor | None = None,
+        guidance_scale: float = 1.0,
+    ) -> torch.Tensor:
+        """
+        Predict model output with optional classifier-free guidance.
+
+        If guidance_scale == 1 or uncond_context is None:
+            normal conditional prediction.
+
+        If guidance_scale > 1:
+            output = uncond + scale * (cond - uncond)
+        """
+        if uncond_context is None or guidance_scale == 1.0:
+            if context is None:
+                return model(
+                    z_t,
+                    t,
+                )
+
+            return model(
+                z_t,
+                t,
+                context=context,
+                attention_mask=attention_mask,
+            )
+
+        # Conditional prediction
+        cond_output = model(
+            z_t,
+            t,
+            context=context,
+            attention_mask=attention_mask,
+        )
+
+        # Unconditional prediction
+        uncond_output = model(
+            z_t,
+            t,
+            context=uncond_context,
+            attention_mask=uncond_attention_mask,
+        )
+
+        return uncond_output + guidance_scale * (cond_output - uncond_output)
+
+    @torch.no_grad()
+    def p_mean_variance_with_cfg(
+        self,
+        model,
+        z_t: torch.Tensor,
+        t: torch.Tensor,
+        context: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        uncond_context: torch.Tensor | None = None,
+        uncond_attention_mask: torch.Tensor | None = None,
+        guidance_scale: float = 1.0,
+        clip_denoised: bool = False,
+    ) -> dict[str, torch.Tensor]:
+        """
+        Same as GaussianDiffusion.p_mean_variance, but supports CFG.
+        """
+        model_output = self.predict_model_output(
+            model=model,
+            z_t=z_t,
+            t=t,
+            context=context,
+            attention_mask=attention_mask,
+            uncond_context=uncond_context,
+            uncond_attention_mask=uncond_attention_mask,
+            guidance_scale=guidance_scale,
+        )
+
+        pred_z0, pred_eps = self.diffusion.predict_x0_and_eps(
+            model_output=model_output,
+            z_t=z_t,
+            t=t,
+        )
+
+        if clip_denoised:
+            pred_z0 = pred_z0.clamp(-1.0, 1.0)
+
+        (
+            posterior_mean,
+            posterior_variance,
+            posterior_log_variance,
+        ) = self.diffusion.q_posterior(
+            z_0=pred_z0,
+            z_t=z_t,
+            t=t,
+        )
+
+        return {
+            "mean": posterior_mean,
+            "variance": posterior_variance,
+            "log_variance": posterior_log_variance,
+            "pred_z0": pred_z0,
+            "pred_eps": pred_eps,
+            "model_output": model_output,
+        }
+
+    @torch.no_grad()
+    def sample(
+        self,
+        model,
+        shape: tuple[int, int, int, int],
+        device: torch.device | str,
+        context: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        uncond_context: torch.Tensor | None = None,
+        uncond_attention_mask: torch.Tensor | None = None,
+        guidance_scale: float = 1.0,
+        clip_denoised: bool = False,
+        return_trajectory: bool = False,
+        progress: bool = True,
+    ) -> DDPMSamplerOutput:
+        """
+        Generate clean latents from pure noise.
+
+        Args:
+            shape:
+                Usually [B, 8, 32, 32] for your model.
+
+            context:
+                Conditional CLIP text context.
+
+            uncond_context:
+                Empty-prompt CLIP context for CFG.
+
+            guidance_scale:
+                CFG scale. Common values: 3.0 to 7.5.
+        """
+        device = torch.device(device)
+
+        model.eval()
+
+        z_t = torch.randn(
+            shape,
+            device=device,
+        )
+
+        trajectory = [] if return_trajectory else None
+
+        timesteps = reversed(range(self.diffusion.num_timesteps))
+
+        if progress:
+            timesteps = tqdm(
+                timesteps,
+                total=self.diffusion.num_timesteps,
+                desc="DDPM sampling",
+            )
+
+        for step in timesteps:
+            t = torch.full(
+                (shape[0],),
+                step,
+                device=device,
+                dtype=torch.long,
+            )
+
+            out = self.p_mean_variance_with_cfg(
+                model=model,
+                z_t=z_t,
+                t=t,
+                context=context,
+                attention_mask=attention_mask,
+                uncond_context=uncond_context,
+                uncond_attention_mask=uncond_attention_mask,
+                guidance_scale=guidance_scale,
+                clip_denoised=clip_denoised,
+            )
+
+            noise = torch.randn_like(z_t)
+
+            nonzero_mask = (t != 0).float()
+
+            while len(nonzero_mask.shape) < len(z_t.shape):
+                nonzero_mask = nonzero_mask[..., None]
+
+            z_t = (
+                out["mean"]
+                + nonzero_mask
+                * torch.exp(0.5 * out["log_variance"])
+                * noise
+            )
+
+            if return_trajectory:
+                trajectory.append(z_t.detach().cpu())
+
+        return DDPMSamplerOutput(
+            latents=z_t,
+            trajectory=trajectory,
+        )

src/losses/diffusion_loss.py

@@ -0,0 +1,252 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def extract(a: torch.Tensor, t: torch.Tensor, x_shape: torch.Size):
+    """
+    Extract coefficients at timestep t
+    a: [T]
+    t: [B]
+    returns: [B, 1, 1, 1]
+    """
+    b = t.shape[0]
+    out = a.gather(-1, t)
+    return out.view(b, *((1,) * (len(x_shape) - 1)))
+
+
+class DiffusionLoss(nn.Module):
+    """
+    Diffusion loss supporting:
+        - epsilon prediction
+        - v prediction 
+
+    v-prediction:
+        v = alpha_t * epsilon - sigma_t * x0
+    """
+
+    def __init__(
+        self,
+        prediction_type: str = "v",  # "epsilon" or "v"
+        loss_type: str = "mse",
+        snr_gamma: float | None = None,
+        snr_weighting: str = "none",
+        normalize_snr_weights: bool = False,
+        eps: float = 1e-8,
+    ):
+        super().__init__()
+
+        prediction_type = prediction_type.lower()
+        loss_type = loss_type.lower()
+        snr_weighting = snr_weighting.lower()
+
+        if prediction_type in {"eps", "epsilon"}:
+            prediction_type = "epsilon"
+
+        elif prediction_type in {"v", "v_prediction"}:
+            prediction_type = "v"
+
+        elif prediction_type in {"x0", "sample"}:
+            prediction_type = "x0"
+
+        else:
+            raise ValueError(
+                "prediction_type must be 'epsilon', 'v', or 'x0'"
+            )
+
+        if loss_type not in {"mse", "l1", "huber"}:
+            raise ValueError(
+                "loss_type must be 'mse', 'l1', or 'huber'"
+            )
+
+        if snr_weighting not in {"none", "min_snr"}:
+            raise ValueError(
+                "snr_weighting must be 'none' or 'min_snr'"
+            )
+
+        if snr_weighting == "min_snr" and snr_gamma is None:
+            raise ValueError(
+                "snr_gamma must be set when snr_weighting='min_snr'"
+            )
+
+        self.prediction_type = prediction_type
+        self.loss_type = loss_type
+        self.snr_gamma = snr_gamma
+        self.snr_weighting = snr_weighting
+        self.normalize_snr_weights = normalize_snr_weights
+        self.eps = eps
+
+    def v_target(self, x0, noise, alpha, sigma):
+        return alpha * noise - sigma * x0
+
+    def epsilon_target(self, x0, noise):
+        return noise
+
+    def x0_target(self, x0):
+        return x0
+
+    def get_target(
+        self,
+        x0: torch.Tensor,
+        noise: torch.Tensor,
+        alpha_t: torch.Tensor,
+        sigma_t: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.prediction_type == "epsilon":
+            return self.epsilon_target(x0, noise)
+
+        if self.prediction_type == "v":
+            return self.v_target(x0, noise, alpha_t, sigma_t)
+
+        if self.prediction_type == "x0":
+            return self.x0_target(x0)
+
+        raise RuntimeError("Invalid prediction type.")
+
+    def elementwise_loss(
+        self,
+        model_output: torch.Tensor,
+        target: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.loss_type == "mse":
+            return F.mse_loss(
+                model_output,
+                target,
+                reduction="none",
+            )
+
+        if self.loss_type == "l1":
+            return F.l1_loss(
+                model_output,
+                target,
+                reduction="none",
+            )
+
+        if self.loss_type == "huber":
+            return F.smooth_l1_loss(
+                model_output,
+                target,
+                reduction="none",
+            )
+
+        raise RuntimeError("Invalid loss type.")
+
+    def get_snr_weights(
+        self,
+        snr: torch.Tensor,
+    ) -> torch.Tensor | None:
+        """
+        Returns per-sample SNR weights.
+
+        snr:
+            [B]
+
+        For Min-SNR:
+            epsilon prediction:
+                weight = min(snr, gamma) / snr
+
+            v prediction:
+                weight = min(snr, gamma) / (snr + 1)
+
+            x0 prediction:
+                weight = min(snr, gamma)
+        """
+        if self.snr_weighting == "none":
+            return None
+
+        if self.snr_weighting == "min_snr":
+            if self.snr_gamma is None:
+                raise RuntimeError("snr_gamma is required for min_snr weighting.")
+
+            snr = snr.float().clamp(min=self.eps)
+
+            gamma = torch.full_like(
+                snr,
+                fill_value=float(self.snr_gamma),
+            )
+
+            clipped_snr = torch.minimum(
+                snr,
+                gamma,
+            )
+
+            if self.prediction_type == "epsilon":
+                weights = clipped_snr / snr
+
+            elif self.prediction_type == "v":
+                weights = clipped_snr / (snr + 1.0)
+
+            elif self.prediction_type == "x0":
+                weights = clipped_snr
+
+            else:
+                raise RuntimeError("Invalid prediction type.")
+
+            if self.normalize_snr_weights:
+                weights = weights / weights.mean().clamp(min=self.eps)
+
+            return weights
+
+        raise RuntimeError("Invalid SNR weighting type.")
+
+    def forward(
+        self,
+        model_output: torch.Tensor,
+        x0: torch.Tensor,
+        noise: torch.Tensor,
+        alpha_t: torch.Tensor,
+        sigma_t: torch.Tensor,
+        snr: torch.Tensor | None = None,
+        return_dict: bool = False,
+    ):
+
+        target = self.get_target(
+            x0=x0,
+            noise=noise,
+            alpha_t=alpha_t,
+            sigma_t=sigma_t,
+        )
+
+        loss = self.elementwise_loss(
+            model_output=model_output,
+            target=target,
+        )
+
+        # [B, C, H, W] -> [B]
+        per_sample_loss = loss.mean(
+            dim=tuple(range(1, loss.ndim)),
+        )
+
+        raw_loss = per_sample_loss.mean()
+
+        weights = None
+
+        if self.snr_weighting != "none":
+            if snr is None:
+                raise ValueError(
+                    "snr must be passed when SNR weighting is enabled."
+                )
+
+            weights = self.get_snr_weights(snr)
+
+            if weights is not None:
+                per_sample_loss = per_sample_loss * weights.to(per_sample_loss.device)
+
+        weighted_loss = per_sample_loss.mean()
+
+        if return_dict:
+            out = {
+                "loss": weighted_loss,
+                "raw_loss": raw_loss.detach(),
+            }
+
+            if weights is not None:
+                out["snr_weight_mean"] = weights.mean().detach()
+                out["snr_weight_min"] = weights.min().detach()
+                out["snr_weight_max"] = weights.max().detach()
+
+            return out
+
+        return weighted_loss

src/losses/vae_loss.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+@dataclass
+class VAELossOutput:
+    total_loss: torch.Tensor
+    recon_loss: torch.Tensor
+    kl_loss: torch.Tensor
+    perceptual_loss: torch.Tensor
+
+
+class VAELoss(nn.Module):
+    """
+    VAE loss:
+        total =
+            recon_weight * reconstruction_loss
+          + kl_weight * KL
+          + perceptual_weight * LPIPS
+
+    Inputs are expected to be in [-1, 1]
+    """
+
+    def __init__(
+        self,
+        recon_loss_type: str = "l1",
+        recon_weight: float = 1.0,
+        kl_weight: float = 1e-6,
+        perceptual_weight: float = 0.0,
+        use_lpips: bool = False,
+        lpips_net: str = "vgg",
+    ):
+        super().__init__()
+
+        if recon_loss_type not in {"l1", "mse"}:
+            raise ValueError(
+                f"Unknown recon_loss_type={recon_loss_type}. "
+                "Use 'l1' or 'mse'."
+            )
+
+        self.recon_loss_type = recon_loss_type
+        self.recon_weight = recon_weight
+        self.kl_weight = kl_weight
+        self.perceptual_weight = perceptual_weight
+        self.use_lpips = use_lpips
+
+        self.lpips_model = None
+
+        if use_lpips:
+            try:
+                import lpips
+            except ImportError as exc:
+                raise ImportError(
+                    "LPIPS is enabled but package 'lpips' is not installed. "
+                    "Install it with: pip install lpips"
+                ) from exc
+
+            self.lpips_model = lpips.LPIPS(net=lpips_net)
+            self.lpips_model.eval()
+
+            for p in self.lpips_model.parameters():
+                p.requires_grad = False
+
+    def reconstruction_loss(
+        self,
+        x_recon: torch.Tensor,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.recon_loss_type == "l1":
+            return F.l1_loss(x_recon, x)
+
+        if self.recon_loss_type == "mse":
+            return F.mse_loss(x_recon, x)
+
+        raise RuntimeError("Invalid reconstruction loss type.")
+
+    def perceptual_loss(
+        self,
+        x_recon: torch.Tensor,
+        x: torch.Tensor,
+    ) -> torch.Tensor:
+        if not self.use_lpips or self.lpips_model is None:
+            return torch.zeros((), device=x.device, dtype=x.dtype)
+
+        # LPIPS expects images in [-1, 1], which matches our transform.
+        with torch.cuda.amp.autocast(enabled=False):
+            loss = self.lpips_model(
+                x_recon.float(),
+                x.float(),
+            ).mean()
+
+        return loss.to(dtype=x.dtype)
+
+    def forward(
+        self,
+        x_recon: torch.Tensor,
+        x: torch.Tensor,
+        posterior,
+        kl_weight: float | None = None,
+    ) -> VAELossOutput:
+        """
+        Args:
+            x_recon:
+                Reconstructed image [B, 3, H, W], in [-1, 1].
+
+            x:
+                Target image [B, 3, H, W], in [-1, 1].
+
+            posterior:
+                DiagonalGaussianDistribution from vae.encode(x).
+
+            kl_weight:
+                Optional current KL weight. Useful for KL warmup.
+
+        Returns:
+            VAELossOutput.
+        """
+        current_kl_weight = self.kl_weight if kl_weight is None else kl_weight
+
+        recon = self.reconstruction_loss(x_recon, x)
+
+        # posterior.kl() returns [B], already summed over latent dimensions.
+        kl = posterior.kl().mean()
+
+        perceptual = self.perceptual_loss(x_recon, x)
+
+        total = (
+            self.recon_weight * recon
+            + current_kl_weight * kl
+            + self.perceptual_weight * perceptual
+        )
+
+        return VAELossOutput(
+            total_loss=total,
+            recon_loss=recon.detach(),
+            kl_loss=kl.detach(),
+            perceptual_loss=perceptual.detach(),
+        )

src/models/autoencoder/blocks.py

@@ -0,0 +1,352 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+def init_conv_kaiming(module: nn.Module) -> None:
+    """
+    Kaiming initialization for convolutional layers used with SiLU activations.
+    """
+    if isinstance(module, nn.Conv2d):
+        nn.init.kaiming_normal_(
+            module.weight,
+            mode="fan_out",
+            nonlinearity="relu",
+        )
+
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+
+def init_linear_xavier(module: nn.Module) -> None:
+    """
+    Xavier initialization for attention-style projection layers.
+    """
+    if isinstance(module, nn.Conv2d):
+        nn.init.xavier_uniform_(module.weight)
+
+        if module.bias is not None:
+            nn.init.zeros_(module.bias)
+
+
+def normalization(num_channels: int, num_groups: int = 32):
+    """
+    GroupNorm used in VAE blocks
+    """
+    num_groups = min(num_groups, num_channels)
+
+    while num_channels % num_groups != 0:
+        num_groups -= 1
+
+    return nn.GroupNorm(
+        num_groups=num_groups,
+        num_channels=num_channels,
+        eps=1e-6,
+        affine=True,
+    )
+
+
+class ResBlock(nn.Module):
+    """
+    Simple residual block:
+
+        x -> GroupNorm -> SiLU -> Conv
+          -> GroupNorm -> SiLU -> Conv
+          + shortcut
+
+    Used both in encoder and decoder.
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int | None = None,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+
+        if out_channels is None:
+            out_channels = in_channels
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        self.norm1 = normalization(in_channels)
+        self.conv1 = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+        self.norm2 = normalization(out_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.conv2 = nn.Conv2d(
+            out_channels,
+            out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+        if in_channels != out_channels:
+            self.shortcut = nn.Conv2d(
+                in_channels,
+                out_channels,
+                kernel_size=1,
+                stride=1,
+                padding=0,
+            )
+        else:
+            self.shortcut = nn.Identity()
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        init_conv_kaiming(self.conv1)
+        init_conv_kaiming(self.conv2)
+
+        if isinstance(self.shortcut, nn.Conv2d):
+            init_conv_kaiming(self.shortcut)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = self.shortcut(x)
+
+        h = self.norm1(x)
+        h = F.silu(h)
+        h = self.conv1(h)
+
+        h = self.norm2(h)
+        h = F.silu(h)
+        h = self.dropout(h)
+        h = self.conv2(h)
+
+        return h + residual
+
+
+class Downsample(nn.Module):
+    """
+    Downsample by factor 2 using strided convolution.
+    """
+
+    def __init__(self, channels: int):
+        super().__init__()
+
+        self.conv = nn.Conv2d(
+            channels,
+            channels,
+            kernel_size=3,
+            stride=2,
+            padding=1,
+        )
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        init_conv_kaiming(self.conv)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.conv(x)
+
+
+class Upsample(nn.Module):
+    """
+    Upsample by factor 2 using nearest-neighbor interpolation + convolution instead of ConvTranspose2d.
+    """
+
+    def __init__(self, channels: int):
+        super().__init__()
+
+        self.conv = nn.Conv2d(
+            channels,
+            channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        init_conv_kaiming(self.conv)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.interpolate(x, scale_factor=2.0, mode="nearest")
+        x = self.conv(x)
+        return x
+
+
+class SelfAttentionBlock(nn.Module):
+    """
+    Spatial self-attention block for feature maps.
+
+    Input:
+        x: [B, C, H, W]
+
+    then get:
+        [B, C, H, W] -> [B, H*W, C]
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int = 1,
+    ):
+        super().__init__()
+
+        if channels % num_heads != 0:
+            raise ValueError(
+                f"channels={channels} must be divisible by num_heads={num_heads}"
+            )
+
+        self.channels = channels
+        self.num_heads = num_heads
+        self.head_dim = channels // num_heads
+
+        self.norm = normalization(channels)
+
+        self.qkv = nn.Conv2d(
+            channels,
+            channels * 3,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+
+        self.proj_out = nn.Conv2d(
+            channels,
+            channels,
+            kernel_size=1,
+            stride=1,
+            padding=0,
+        )
+
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        init_linear_xavier(self.qkv)
+        init_linear_xavier(self.proj_out)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        b, c, h, w = x.shape
+        residual = x
+
+        x = self.norm(x)
+
+        qkv = self.qkv(x)
+        q, k, v = torch.chunk(qkv, chunks=3, dim=1)
+
+        # [B, C, H, W] -> [B, num_heads, H*W, head_dim]
+        q = q.view(b, self.num_heads, self.head_dim, h * w)
+        k = k.view(b, self.num_heads, self.head_dim, h * w)
+        v = v.view(b, self.num_heads, self.head_dim, h * w)
+
+        q = q.permute(0, 1, 3, 2)
+        k = k.permute(0, 1, 3, 2)
+        v = v.permute(0, 1, 3, 2)
+
+        # Output: [B, num_heads, H*W, head_dim]
+        out = F.scaled_dot_product_attention(q, k, v)
+
+        # [B, num_heads, H*W, head_dim] -> [B, C, H, W]
+        out = out.permute(0, 1, 3, 2).contiguous()
+        out = out.view(b, c, h, w)
+
+        out = self.proj_out(out)
+
+        return residual + out
+
+
+class AttnResBlock(nn.Module):
+    """
+    Optional attention block:
+
+        ResBlock -> SelfAttentionBlock
+    """
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int | None = None,
+        dropout: float = 0.0,
+        use_attention: bool = False,
+        num_heads: int = 1,
+    ):
+        super().__init__()
+
+        if out_channels is None:
+            out_channels = in_channels
+
+        self.resblock = ResBlock(
+            in_channels=in_channels,
+            out_channels=out_channels,
+            dropout=dropout,
+        )
+
+        if use_attention:
+            self.attn = SelfAttentionBlock(
+                channels=out_channels,
+                num_heads=num_heads,
+            )
+        else:
+            self.attn = nn.Identity()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.resblock(x)
+        x = self.attn(x)
+        return x
+
+
+class MidBlock(nn.Module):
+    """
+    Bottleneck block:
+
+        ResBlock -> SelfAttentionBlock -> ResBlock
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        dropout: float = 0.0,
+        use_attention: bool = True,
+        num_heads: int = 1,
+    ):
+        super().__init__()
+
+        self.block1 = ResBlock(
+            in_channels=channels,
+            out_channels=channels,
+            dropout=dropout,
+        )
+
+        if use_attention:
+            self.attn = SelfAttentionBlock(
+                channels=channels,
+                num_heads=num_heads,
+            )
+        else:
+            self.attn = nn.Identity()
+
+        self.block2 = ResBlock(
+            in_channels=channels,
+            out_channels=channels,
+            dropout=dropout,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.block1(x)
+        x = self.attn(x)
+        x = self.block2(x)
+        return x
+
+
+def zero_module(module: nn.Module) -> nn.Module:
+    """
+    Zero-initialize a module.
+    """
+    for p in module.parameters():
+        nn.init.zeros_(p)
+    return module

src/models/autoencoder/decoder.py

@@ -0,0 +1,162 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from src.models.autoencoder.blocks import (
+    ResBlock,
+    Upsample,
+    MidBlock,
+    normalization,
+)
+
+
+class Decoder(nn.Module):
+    """
+    VAE decoder.
+
+    Converts latent tensor back into image space.
+    Shape path:
+
+        [B, 4,   32,  32]
+            -> conv_in
+        [B, 512, 32,  32]
+            -> mid block
+        [B, 512, 32,  32]
+            -> upsample
+        [B, 512, 64,  64]
+            -> upsample
+        [B, 256, 128, 128]
+            -> upsample
+        [B, 128, 256, 256]
+            -> conv_out
+        [B, 3,   256, 256]
+
+    Output is in [-1, 1] because training images are normalized to [-1, 1].
+    """
+
+    def __init__(
+        self,
+        out_channels: int = 3,
+        latent_channels: int = 4,
+        base_channels: int = 128,
+        channel_multipliers: list[int] | tuple[int, ...] = (1, 2, 4, 4),
+        num_res_blocks: int = 2,
+        dropout: float = 0.0,
+        use_attention: bool = True,
+        attention_heads: int = 1
+    ):
+        super().__init__()
+
+        if len(channel_multipliers) < 2:
+            raise ValueError("channel_multipliers must contain at least 2 levels.")
+
+        self.out_channels = out_channels
+        self.latent_channels = latent_channels
+        self.base_channels = base_channels
+        self.channel_multipliers = list(channel_multipliers)
+        self.num_res_blocks = num_res_blocks
+
+        # Number of spatial upsampling operations
+        # Example:
+        #   [1, 2, 4, 4] has 4 levels, so decoder upsamples 3 times:
+        #   32 -> 64 -> 128 -> 256
+        self.num_upsamples = len(self.channel_multipliers) - 1
+
+        # Start from the deepest encoder channel count
+        current_channels = base_channels * self.channel_multipliers[-1]
+
+        self.conv_in = nn.Conv2d(
+            latent_channels,
+            current_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+        # Bottleneck block at the lowest spatial resolution.
+        self.mid = MidBlock(
+            channels=current_channels,
+            dropout=dropout,
+            use_attention=use_attention,
+            num_heads=attention_heads,
+        )
+
+        self.up_blocks = nn.ModuleList()
+
+        reversed_multipliers = list(reversed(self.channel_multipliers))
+
+        for level, multiplier in enumerate(reversed_multipliers):
+            out_stage_channels = base_channels * multiplier
+
+            resblocks = nn.ModuleList()
+
+            # one extra ResBlock per level
+            for _ in range(num_res_blocks + 1):
+                resblocks.append(
+                    ResBlock(
+                        in_channels=current_channels,
+                        out_channels=out_stage_channels,
+                        dropout=dropout,
+                    )
+                )
+                current_channels = out_stage_channels
+
+            # Upsample after every stage except the full-resolution
+            if level < len(reversed_multipliers) - 1:
+                upsample = Upsample(
+                    channels=current_channels
+                )
+            else:
+                upsample = nn.Identity()
+
+            self.up_blocks.append(
+                nn.ModuleDict(
+                    {
+                        "resblocks": resblocks,
+                        "upsample": upsample,
+                    }
+                )
+            )
+
+        self.norm_out = normalization(current_channels)
+
+        self.conv_out = nn.Conv2d(
+            current_channels,
+            out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+    def forward(self, z: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            z:
+                Latent tensor with shape [B, latent_channels, H/8, W/8].
+                For 256x256 images and downsample factor 8:
+                    [B, latent_channels, 32, 32]
+
+        Returns:
+            x_recon:
+                Reconstructed image tensor with shape [B, 3, H, W].
+                Values are in [-1, 1].
+        """
+        h = self.conv_in(z)
+
+        h = self.mid(h)
+
+        for stage in self.up_blocks:
+            for block in stage["resblocks"]:
+                h = block(h)
+
+            h = stage["upsample"](h)
+
+        h = self.norm_out(h)
+        h = F.silu(h)
+        h = self.conv_out(h)
+
+        x_recon = torch.tanh(h)
+
+        return x_recon

src/models/autoencoder/distributions.py

@@ -0,0 +1,113 @@
+from __future__ import annotations
+
+import torch
+
+
+class DiagonalGaussianDistribution:
+    """
+    Diagonal Gaussian posterior used by the VAE.
+
+    The encoder predicts a tensor called `moments`:
+
+        moments: [B, 2 * latent_channels, H, W]
+
+    We split it into:
+
+        mean:   [B, latent_channels, H, W]
+        logvar: [B, latent_channels, H, W]
+
+    Then sample:
+
+        z = mean + std * eps
+
+    where:
+
+        std = exp(0.5 * logvar)
+        eps ~ N(0, I)
+    """
+
+    def __init__(
+        self,
+        moments: torch.Tensor,
+        deterministic: bool = False,
+        logvar_min: float = -30.0,
+        logvar_max: float = 20.0,
+    ):
+        self.moments = moments
+        self.deterministic = deterministic
+
+        self.mean, self.logvar = torch.chunk(moments, chunks=2, dim=1)
+
+        # Clamp log-variance for numerical stability
+        self.logvar = torch.clamp(self.logvar, min=logvar_min, max=logvar_max)
+
+        self.var = torch.exp(self.logvar)
+        self.std = torch.exp(0.5 * self.logvar)
+
+        if self.deterministic:
+            self.var = torch.zeros_like(self.mean)
+            self.std = torch.zeros_like(self.mean)
+
+    def sample(self) -> torch.Tensor:
+        """
+        Reparameterized sampling:
+
+            z = mean + std * eps
+        """
+        eps = torch.randn_like(self.mean)
+        return self.mean + self.std * eps
+
+    def mode(self) -> torch.Tensor:
+        """
+        Most likely latent value
+        """
+        return self.mean
+
+    def kl(self) -> torch.Tensor:
+        """
+        KL divergence from posterior q(z|x) to standard normal N(0, I).
+
+        For diagonal Gaussian:
+
+            KL(q || N(0,I))
+            = 0.5 * (mean^2 + var - 1 - logvar)
+
+        Returns:
+            Per-sample KL with shape [B].
+        """
+        if self.deterministic:
+            return torch.zeros(self.mean.shape[0], device=self.mean.device)
+
+        kl = 0.5 * (
+            torch.pow(self.mean, 2)
+            + self.var
+            - 1.0
+            - self.logvar
+        )
+
+        # Sum over latent channels and spatial dimensions.
+        return torch.sum(kl, dim=[1, 2, 3])
+
+    def nll(self, sample: torch.Tensor) -> torch.Tensor:
+        """
+        Negative log likelihood of `sample` under this posterior.
+
+        Mostly useful for debugging, not essential for our VAE training loop.
+
+        Returns:
+            Per-sample NLL with shape [B].
+        """
+        if self.deterministic:
+            return torch.zeros(self.mean.shape[0], device=self.mean.device)
+
+        log_two_pi = torch.log(
+            torch.tensor(2.0 * torch.pi, device=sample.device, dtype=sample.dtype)
+        )
+
+        nll = 0.5 * (
+            log_two_pi
+            + self.logvar
+            + torch.pow(sample - self.mean, 2) / self.var
+        )
+
+        return torch.sum(nll, dim=[1, 2, 3])

src/models/autoencoder/encoder.py

@@ -0,0 +1,150 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from src.models.autoencoder.blocks import (
+    ResBlock,
+    Downsample,
+    MidBlock,
+    normalization,
+    SelfAttentionBlock,
+)
+
+class Encoder(nn.Module):
+    """
+    VAE encoder.
+    channel_multipliers=[1, 2, 4]: this controls the multiplier of number of feature maps
+
+        [B, 3, 256, 256]
+            -> [B, 128, 256, 256]
+            -> [B, 128, 128, 128]
+            -> [B, 256, 64, 64]
+            -> [B, 512, 32, 32]
+            -> [B, 2 * latent_channels, 32, 32]
+
+    Output channels are 2 * latent_channels because we predict:
+        mu
+        logvar
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        latent_channels: int = 8,
+        base_channels: int = 128,
+        channel_multipliers: list[int] | tuple[int, ...] = (1, 2, 4, 4),
+        num_res_blocks: int = 3,
+        dropout: float = 0.0,
+        use_attention: bool = True,
+        attention_heads: int = 4,
+        attention_resolutions: tuple[int, ...] = (32,),
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.latent_channels = latent_channels
+        self.base_channels = base_channels
+        self.channel_multipliers = list(channel_multipliers)
+        self.num_res_blocks = num_res_blocks
+        self.attention_resolutions = set(attention_resolutions)
+
+        # Initial projection
+        self.conv_in = nn.Conv2d(
+            in_channels,
+            base_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+        # Downsampling
+        self.down_blocks = nn.ModuleList()
+
+        current_channels = base_channels
+        current_resolution = 256
+
+        for level, multiplier in enumerate(self.channel_multipliers):
+            out_channels = base_channels * multiplier
+
+            stage = nn.ModuleDict()
+            stage["resblocks"] = nn.ModuleList()
+
+            for _ in range(num_res_blocks):
+                stage["resblocks"].append(
+                    ResBlock(
+                        in_channels=current_channels,
+                        out_channels=out_channels,
+                        dropout=dropout,
+                    )
+                )
+                current_channels = out_channels
+
+
+            # This part also adds attention to 64x64 resolution along with bottleneck.
+            if use_attention and current_resolution in self.attention_resolutions:
+                stage["attention"] = SelfAttentionBlock(
+                    channels=current_channels,
+                    num_heads=attention_heads,
+                )
+            else:
+                stage["attention"] = nn.Identity()
+
+            # Downsample after each stage except the final one
+            if level != len(self.channel_multipliers) - 1:
+                stage["downsample"] = Downsample(current_channels)
+                next_resolution = current_resolution // 2
+            else:
+                stage["downsample"] = nn.Identity()
+                next_resolution = current_resolution
+
+            self.down_blocks.append(stage)
+            current_resolution = next_resolution
+
+        # Bottleneck
+        self.mid = MidBlock(
+            channels=current_channels,
+            dropout=dropout,
+            use_attention=use_attention,
+            num_heads=attention_heads,
+        )
+
+        # Output projection to posterior parameters
+        self.norm_out = normalization(current_channels)
+        self.conv_out = nn.Conv2d(
+            current_channels,
+            2 * latent_channels,
+            kernel_size=3,
+            stride=1,
+            padding=1,
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            x:
+                Image tensor with shape [B, 3, H, W]
+
+        Returns:
+            moments:
+                Tensor with shape [B, 2 * latent_channels, H/8, W/8]
+                The first half is mu.
+                The second half is logvar.
+        """
+        h = self.conv_in(x)
+
+        for stage in self.down_blocks:
+            for block in stage["resblocks"]:
+                h = block(h)
+
+            h = stage["attention"](h)
+            h = stage["downsample"](h)
+
+        h = self.mid(h)
+
+        h = self.norm_out(h)
+        h = F.silu(h)
+        moments = self.conv_out(h)
+
+        return moments

src/models/autoencoder/vae.py

@@ -0,0 +1,215 @@
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from src.models.autoencoder.encoder import Encoder
+from src.models.autoencoder.decoder import Decoder
+from src.models.autoencoder.distributions import DiagonalGaussianDistribution
+
+
+class AutoencoderKL(nn.Module):
+    """
+    VAE / AutoencoderKL wrapper
+        posterior = vae.encode(x)
+        z = posterior.sample()
+        x_recon = vae.decode(z)
+
+    Or directly:
+
+        x_recon, posterior, z = vae(x)
+
+    Input image range:
+        x in [-1, 1]
+
+    Output image range:
+        x_recon in [-1, 1]
+    """
+
+    def __init__(
+        self,
+        in_channels: int = 3,
+        out_channels: int = 3,
+        latent_channels: int = 8,
+        base_channels: int = 128,
+        channel_multipliers: list[int] | tuple[int, ...] = (1, 2, 4, 4),
+        num_res_blocks: int = 3,
+        dropout: float = 0.0,
+        use_attention: bool = True,
+        attention_heads: int = 4,
+        scaling_factor: float = 1.0,
+        attention_resolutions: tuple[int, ...] = (32,),
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.latent_channels = latent_channels
+        self.base_channels = base_channels
+        self.channel_multipliers = list(channel_multipliers)
+        self.num_res_blocks = num_res_blocks
+        self.scaling_factor = scaling_factor
+
+        self.encoder = Encoder(
+            in_channels=in_channels,
+            latent_channels=latent_channels,
+            base_channels=base_channels,
+            channel_multipliers=channel_multipliers,
+            num_res_blocks=num_res_blocks,
+            dropout=dropout,
+            use_attention=use_attention,
+            attention_heads=attention_heads,
+            attention_resolutions= attention_resolutions
+        )
+
+        self.decoder = Decoder(
+            out_channels=out_channels,
+            latent_channels=latent_channels,
+            base_channels=base_channels,
+            channel_multipliers=channel_multipliers,
+            num_res_blocks=num_res_blocks,
+            dropout=dropout,
+            use_attention=use_attention,
+            attention_heads=attention_heads,
+        )
+
+    def encode(
+        self,
+        x: torch.Tensor,
+        deterministic: bool = False,
+    ) -> DiagonalGaussianDistribution:
+        """
+        Encode image into posterior distribution.
+            deterministic:
+                If True, posterior.sample() will return mean only.
+        Returns:
+            DiagonalGaussianDistribution.
+        """
+        moments = self.encoder(x)
+        posterior = DiagonalGaussianDistribution(
+            moments=moments,
+            deterministic=deterministic,
+        )
+        return posterior
+
+    def decode(
+        self,
+        z: torch.Tensor,
+        unscale: bool = True,
+    ) -> torch.Tensor:
+        """
+        Decode latent into image
+            z:
+                Latent tensor [B, latent_channels, H/8, W/8].
+
+            unscale:
+                If True, divide by scaling_factor before decoding.
+
+        Returns:
+            Reconstructed image in [-1, 1].
+        """
+        if unscale:
+            z = z / self.scaling_factor
+
+        return self.decoder(z)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        sample_posterior: bool = True,
+    ) -> tuple[torch.Tensor, DiagonalGaussianDistribution, torch.Tensor]:
+        """
+        Full VAE forward pass.
+
+        Args:
+            x:
+                Image tensor [B, 3, H, W], normalized to [-1, 1].
+
+            sample_posterior:
+                If True:
+                    z = posterior.sample()
+                If False:
+                    z = posterior.mode()
+
+        Returns:
+            x_recon:
+                Reconstructed image [B, 3, H, W].
+
+            posterior:
+                DiagonalGaussianDistribution object.
+
+            z:
+                Latent tensor before scaling.
+        """
+        posterior = self.encode(x)
+
+        if sample_posterior:
+            z = posterior.sample()
+        else:
+            z = posterior.mode()
+
+        x_recon = self.decode(z, unscale=False)
+
+        return x_recon, posterior, z
+
+    @torch.no_grad()
+    def reconstruct(
+        self,
+        x: torch.Tensor,
+        sample_posterior: bool = False,
+    ) -> torch.Tensor:
+        """
+        Convenience method for inference reconstruction.
+
+        By default, uses posterior mode for stable reconstructions.
+        """
+        x_recon, _, _ = self.forward(
+            x,
+            sample_posterior=sample_posterior,
+        )
+        return x_recon
+
+    @torch.no_grad()
+    def encode_to_latent(
+        self,
+        x: torch.Tensor,
+        sample_posterior: bool = False,
+        scale: bool = True,
+    ) -> torch.Tensor:
+        """
+        Encode image into latent tensor for latent diffusion training.
+
+        Usually for latent caching, use:
+
+            sample_posterior=False
+
+        because the posterior mean is deterministic and stable.
+
+        If scale=True:
+
+            z_scaled = z * scaling_factor
+
+        Stable Diffusion-style LDMs often scale latents before diffusion.
+        """
+        posterior = self.encode(x)
+
+        if sample_posterior:
+            z = posterior.sample()
+        else:
+            z = posterior.mode()
+
+        if scale:
+            z = z * self.scaling_factor
+
+        return z
+
+    @torch.no_grad()
+    def decode_from_latent(
+        self,
+        z: torch.Tensor,
+        unscale: bool = True,
+    ) -> torch.Tensor:
+        """
+        Decode latent tensor produced by diffusion model.
+        """
+        return self.decode(z, unscale=unscale)

src/models/conditioning/clip_text.py

@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import nn
+from transformers import CLIPTextModel, CLIPTokenizer
+
+
+@dataclass
+class TextConditioningOutput:
+    """
+    Output of the CLIP text encoder.
+
+    hidden_states:
+        Token-level CLIP embeddings.
+        Shape: [B, seq_len, hidden_dim]
+
+    attention_mask:
+        Token attention mask.
+        Shape: [B, seq_len]
+
+    pooled:
+        Optional pooled text embedding.
+        Shape: [B, hidden_dim]
+    """
+
+    hidden_states: torch.Tensor
+    attention_mask: torch.Tensor
+    pooled: torch.Tensor | None = None
+
+
+class FrozenCLIPTextEncoder(nn.Module):
+    """
+    Frozen CLIP text encoder for latent diffusion conditioning
+    model:
+        openai/clip-vit-large-patch14
+
+    This gives:
+        context_dim = 768
+        max_length = 77
+
+    """
+
+    def __init__(
+        self,
+        model_name: str = "openai/clip-vit-large-patch14",
+        max_length: int = 77,
+        freeze: bool = True,
+        use_last_hidden_state: bool = True,
+        cache_dir: str | None = None,
+        local_files_only: bool = False,
+    ):
+        super().__init__()
+
+        self.model_name = model_name
+        self.max_length = max_length
+        self.freeze = freeze
+        self.use_last_hidden_state = use_last_hidden_state
+        self.cache_dir = cache_dir
+        self.local_files_only = local_files_only
+
+        self.tokenizer = CLIPTokenizer.from_pretrained(
+            model_name,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+
+        self.text_model = CLIPTextModel.from_pretrained(
+            model_name,
+            cache_dir=cache_dir,
+            local_files_only=local_files_only,
+        )
+
+        if self.freeze:
+            self.text_model.eval()
+            for p in self.text_model.parameters():
+                p.requires_grad = False
+
+    @property
+    def context_dim(self) -> int:
+        return int(self.text_model.config.hidden_size)
+
+    @property
+    def vocab_size(self) -> int:
+        return int(self.tokenizer.vocab_size)
+
+    @property
+    def pad_token_id(self) -> int:
+        return int(self.tokenizer.pad_token_id)
+
+    def train(self, mode: bool = True):
+        """
+        Keep CLIP frozen/eval even if parent model calls .train().
+        """
+        super().train(mode)
+
+        if self.freeze:
+            self.text_model.eval()
+
+        return self
+
+    def tokenize(
+        self,
+        captions: list[str] | tuple[str, ...],
+        device: torch.device | str | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """
+        Tokenize captions into CLIP input tensors.
+        """
+        tokens = self.tokenizer(
+            list(captions),
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_length,
+            return_tensors="pt",
+        )
+
+        if device is not None:
+            tokens = {
+                key: value.to(device)
+                for key, value in tokens.items()
+            }
+
+        return tokens
+
+    def forward(
+        self,
+        captions: list[str] | tuple[str, ...],
+        device: torch.device | str | None = None,
+    ) -> TextConditioningOutput:
+        """
+        Produces CLIP textual embeddings as diffusion condition.
+        """
+        if device is None:
+            device = next(self.text_model.parameters()).device
+
+        tokens = self.tokenize(
+            captions=captions,
+            device=device,
+        )
+
+        with torch.no_grad() if self.freeze else torch.enable_grad():
+            outputs = self.text_model(
+                input_ids=tokens["input_ids"],
+                attention_mask=tokens["attention_mask"],
+                output_hidden_states=not self.use_last_hidden_state,
+                return_dict=True,
+            )
+
+        if self.use_last_hidden_state:
+            hidden_states = outputs.last_hidden_state
+        else:
+            # Penultimate layer is sometimes used in diffusion models.
+            hidden_states = outputs.hidden_states[-2]
+
+        pooled = outputs.pooler_output
+
+        return TextConditioningOutput(
+            hidden_states=hidden_states,
+            attention_mask=tokens["attention_mask"],
+            pooled=pooled,
+        )
+
+    @torch.no_grad()
+    def encode(
+        self,
+        captions: list[str] | tuple[str, ...],
+        device: torch.device | str | None = None,
+    ) -> torch.Tensor:
+        """
+        Convenience function.
+
+        Returns only token-level context:
+
+            [B, seq_len, context_dim]
+        """
+        return self.forward(
+            captions=captions,
+            device=device,
+        ).hidden_states

src/models/conditioning/null_conditioning.py

@@ -0,0 +1,144 @@
+from __future__ import annotations
+
+import random
+
+import torch
+from torch import nn
+
+from src.models.conditioning.clip_text import FrozenCLIPTextEncoder
+
+
+class ClassifierFreeGuidanceConditioner(nn.Module):
+    """
+    During training, with probability `cond_drop_prob`, captions are replaced
+    with the empty string:
+
+        "a dog on grass" -> ""
+
+    Later at sampling time, CFG uses:
+
+        pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond)
+
+    """
+
+    def __init__(
+        self,
+        text_encoder: FrozenCLIPTextEncoder,
+        cond_drop_prob: float = 0.1,
+        empty_text: str = "",
+    ):
+        super().__init__()
+
+        if cond_drop_prob < 0.0 or cond_drop_prob > 1.0:
+            raise ValueError("cond_drop_prob must be between 0 and 1.")
+
+        self.text_encoder = text_encoder
+        self.cond_drop_prob = cond_drop_prob
+        self.empty_text = empty_text
+
+    @property
+    def context_dim(self) -> int:
+        return self.text_encoder.context_dim
+
+    @property
+    def max_length(self) -> int:
+        return self.text_encoder.max_length
+
+    def apply_conditioning_dropout(
+        self,
+        captions: list[str] | tuple[str, ...],
+        force_drop_ids: torch.Tensor | None = None,
+    ) -> tuple[list[str], torch.Tensor]:
+        """
+        Replace some captions with empty text
+        """
+        captions = list(captions)
+        batch_size = len(captions)
+
+        if force_drop_ids is not None:
+            drop_mask = force_drop_ids.bool().cpu()
+        else:
+            drop_mask = torch.zeros(batch_size, dtype=torch.bool)
+
+            for i in range(batch_size):
+                if random.random() < self.cond_drop_prob:
+                    drop_mask[i] = True
+
+        dropped_captions = []
+
+        for caption, drop in zip(captions, drop_mask):
+            if bool(drop):
+                dropped_captions.append(self.empty_text)
+            else:
+                dropped_captions.append(caption)
+
+        return dropped_captions, drop_mask
+
+    def forward(
+        self,
+        captions: list[str] | tuple[str, ...],
+        device: torch.device | str | None = None,
+        apply_dropout: bool = True,
+        force_drop_ids: torch.Tensor | None = None,
+    ):
+        """
+        Encode captions with optional CFG dropout
+        """
+        if apply_dropout:
+            captions, drop_mask = self.apply_conditioning_dropout(
+                captions=captions,
+                force_drop_ids=force_drop_ids,
+            )
+        else:
+            captions = list(captions)
+            drop_mask = torch.zeros(
+                len(captions),
+                dtype=torch.bool,
+            )
+
+        output = self.text_encoder(
+            captions=captions,
+            device=device,
+        )
+
+        if device is not None:
+            drop_mask = drop_mask.to(device)
+
+        return {
+            "context": output.hidden_states,
+            "attention_mask": output.attention_mask,
+            "pooled": output.pooled,
+            "drop_mask": drop_mask,
+            "captions": captions,
+        }
+
+    @torch.no_grad()
+    def encode_cond_uncond(
+        self,
+        captions: list[str] | tuple[str, ...],
+        device: torch.device | str | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """
+        Encode both conditional and unconditional text.
+
+        Used during CFG sampling
+        """
+        captions = list(captions)
+        batch_size = len(captions)
+
+        cond_output = self.text_encoder(
+            captions=captions,
+            device=device,
+        )
+
+        uncond_output = self.text_encoder(
+            captions=[self.empty_text] * batch_size,
+            device=device,
+        )
+
+        return {
+            "cond_context": cond_output.hidden_states,
+            "cond_attention_mask": cond_output.attention_mask,
+            "uncond_context": uncond_output.hidden_states,
+            "uncond_attention_mask": uncond_output.attention_mask,
+        }