diff --git a/.gitattributes b/.gitattributes index a6344aac8c09253b3b630fb776ae94478aa0275b..b4f1facfcc5f3bd41f109cc26b4ba23ce92f7c35 100644 --- a/.gitattributes +++ b/.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 diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..3bdbfd60cbacd04097db9d5736848b261b245aff --- /dev/null +++ b/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. diff --git a/checkpoints/ldm_model.pt b/checkpoints/ldm_model.pt new file mode 100644 index 0000000000000000000000000000000000000000..843efa525e874ac530f542e48488de1bc1b7119f --- /dev/null +++ b/checkpoints/ldm_model.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:79a241aa2d3f67e89d8d9227e5022efc36d603ddf07c426c599b78304dd668b7 +size 1798549690 diff --git a/checkpoints/vae_model.pt b/checkpoints/vae_model.pt new file mode 100644 index 0000000000000000000000000000000000000000..272d2c25cb7ea845b248e57b447da0c2cd07c24e --- /dev/null +++ b/checkpoints/vae_model.pt @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:179bfe367d29e1e726638b71a8a9dfc17651d9b368d4fade5029022c89df5b15 +size 426519034 diff --git a/configs/ldm_config.yaml b/configs/ldm_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..cc6fefb4e9065f56c9803d7be7b919e47ed841c3 --- /dev/null +++ b/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 diff --git a/configs/vae_config.yaml b/configs/vae_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c73a17151ddb4733d50d1fdc9aca8d00e64abaa1 --- /dev/null +++ b/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 diff --git a/generation_config.yaml b/generation_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..eaa84d5850fd598a93e51898e4bcb0a3bd795e72 --- /dev/null +++ b/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 diff --git a/inference.py b/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..438972c742a921ca0911a92065abb2bcf689f675 --- /dev/null +++ b/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() diff --git a/outputs/test_ddim/0000_a_small_dog_sitting_on_a_red_couch.png b/outputs/test_ddim/0000_a_small_dog_sitting_on_a_red_couch.png new file mode 100644 index 0000000000000000000000000000000000000000..014c7e9e955ca0702cf17a66acfd50c0158731d3 Binary files /dev/null and b/outputs/test_ddim/0000_a_small_dog_sitting_on_a_red_couch.png differ diff --git a/outputs/test_ddim/grid.png b/outputs/test_ddim/grid.png new file mode 100644 index 0000000000000000000000000000000000000000..014c7e9e955ca0702cf17a66acfd50c0158731d3 Binary files /dev/null and b/outputs/test_ddim/grid.png differ diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..c18df55cdeb552ad6e9c31982a21d4eb82a60070 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,8 @@ +torch +torchvision +transformers +huggingface_hub +python-dotenv +PyYAML +Pillow +tqdm diff --git a/src/__init__.py b/src/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/__pycache__/__init__.cpython-311.pyc b/src/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..66ff56204b7a86af181039b79597c816082fe475 Binary files /dev/null and b/src/__pycache__/__init__.cpython-311.pyc differ diff --git a/src/diffusion/__init__.py b/src/diffusion/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/diffusion/__pycache__/__init__.cpython-311.pyc b/src/diffusion/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c2c469bcc3391c889875959564653e3d7bf93be Binary files /dev/null and b/src/diffusion/__pycache__/__init__.cpython-311.pyc differ diff --git a/src/diffusion/__pycache__/gaussian_diffusion.cpython-311.pyc b/src/diffusion/__pycache__/gaussian_diffusion.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..180430271c90f9987b554d6d3fddd80d9531d6b1 Binary files /dev/null and b/src/diffusion/__pycache__/gaussian_diffusion.cpython-311.pyc differ diff --git a/src/diffusion/__pycache__/noise_schedule.cpython-311.pyc b/src/diffusion/__pycache__/noise_schedule.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9a0b8cfa59ae54a43bfd004789f9c8b6843e4672 Binary files /dev/null and b/src/diffusion/__pycache__/noise_schedule.cpython-311.pyc differ diff --git a/src/diffusion/__pycache__/prediction.cpython-311.pyc b/src/diffusion/__pycache__/prediction.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..df9a2d93becbaed964ab01f11f5e2ab7db55962d Binary files /dev/null and b/src/diffusion/__pycache__/prediction.cpython-311.pyc differ diff --git a/src/diffusion/gaussian_diffusion.py b/src/diffusion/gaussian_diffusion.py new file mode 100644 index 0000000000000000000000000000000000000000..885c997f18c0059810db1c606f3fda5d0177fece --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/diffusion/noise_schedule.py b/src/diffusion/noise_schedule.py new file mode 100644 index 0000000000000000000000000000000000000000..e6cdd3b53e5ac3318d7da78fee70797f2f45689d --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/diffusion/prediction.py b/src/diffusion/prediction.py new file mode 100644 index 0000000000000000000000000000000000000000..b9f8d8f64466739ce2fba7a0145814b8d101137b --- /dev/null +++ b/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'." + ) \ No newline at end of file diff --git a/src/diffusion/samplers/__init__.py b/src/diffusion/samplers/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a6847a68921675c6f5d1c886b0b483d152c7b4e7 --- /dev/null +++ b/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", +] \ No newline at end of file diff --git a/src/diffusion/samplers/__pycache__/__init__.cpython-311.pyc b/src/diffusion/samplers/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce63c09eff394f36d0f309a5dd696ee2da633e78 Binary files /dev/null and b/src/diffusion/samplers/__pycache__/__init__.cpython-311.pyc differ diff --git a/src/diffusion/samplers/__pycache__/ddim.cpython-311.pyc b/src/diffusion/samplers/__pycache__/ddim.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..354a40ad7e47469254db1bb389721fd4a3496eca Binary files /dev/null and b/src/diffusion/samplers/__pycache__/ddim.cpython-311.pyc differ diff --git a/src/diffusion/samplers/__pycache__/ddpm.cpython-311.pyc b/src/diffusion/samplers/__pycache__/ddpm.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..303417830088f2f085db3fdaf25aae2a8128fa3b Binary files /dev/null and b/src/diffusion/samplers/__pycache__/ddpm.cpython-311.pyc differ diff --git a/src/diffusion/samplers/ddim.py b/src/diffusion/samplers/ddim.py new file mode 100644 index 0000000000000000000000000000000000000000..71d539a17747090684f1fdcd326ff2035cfff3f0 --- /dev/null +++ b/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, + ) \ No newline at end of file diff --git a/src/diffusion/samplers/ddpm.py b/src/diffusion/samplers/ddpm.py new file mode 100644 index 0000000000000000000000000000000000000000..c970691742f3cc322a2fdd7bfe5f5437fdc60813 --- /dev/null +++ b/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, + ) \ No newline at end of file diff --git a/src/losses/__pycache__/diffusion_loss.cpython-311.pyc b/src/losses/__pycache__/diffusion_loss.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3c452e228360c1c8309d6714e531b3031f4c5125 Binary files /dev/null and b/src/losses/__pycache__/diffusion_loss.cpython-311.pyc differ diff --git a/src/losses/__pycache__/vae_loss.cpython-311.pyc b/src/losses/__pycache__/vae_loss.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7e602ba8ef26ef1916d3f43625ab8994d7544d42 Binary files /dev/null and b/src/losses/__pycache__/vae_loss.cpython-311.pyc differ diff --git a/src/losses/diffusion_loss.py b/src/losses/diffusion_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..0ed1c7b2482d47dc59e8477fe12f50b2a538c53f --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/losses/vae_loss.py b/src/losses/vae_loss.py new file mode 100644 index 0000000000000000000000000000000000000000..b71a110477181a2c59c67711327bcb8396e24045 --- /dev/null +++ b/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(), + ) \ No newline at end of file diff --git a/src/models/__init__.py b/src/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/models/__pycache__/__init__.cpython-311.pyc b/src/models/__pycache__/__init__.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dbd0203900e9c5aa36d7f62b69c13f21fcd5b1c4 Binary files /dev/null and b/src/models/__pycache__/__init__.cpython-311.pyc differ diff --git a/src/models/autoencoder/__pycache__/blocks.cpython-311.pyc b/src/models/autoencoder/__pycache__/blocks.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fc1a42b75a398a1152e58c4d5b47b02310711559 Binary files /dev/null and b/src/models/autoencoder/__pycache__/blocks.cpython-311.pyc differ diff --git a/src/models/autoencoder/__pycache__/decoder.cpython-311.pyc b/src/models/autoencoder/__pycache__/decoder.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7d2fd787276df2f5c6c5419f2c85baff9589706b Binary files /dev/null and b/src/models/autoencoder/__pycache__/decoder.cpython-311.pyc differ diff --git a/src/models/autoencoder/__pycache__/distributions.cpython-311.pyc b/src/models/autoencoder/__pycache__/distributions.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..fa550ed25e090f93a9a755f1a9281ad6ed238e34 Binary files /dev/null and b/src/models/autoencoder/__pycache__/distributions.cpython-311.pyc differ diff --git a/src/models/autoencoder/__pycache__/encoder.cpython-311.pyc b/src/models/autoencoder/__pycache__/encoder.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e914791d48e5eef909fe3a923cd39f54b16f0c5 Binary files /dev/null and b/src/models/autoencoder/__pycache__/encoder.cpython-311.pyc differ diff --git a/src/models/autoencoder/__pycache__/vae.cpython-311.pyc b/src/models/autoencoder/__pycache__/vae.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dc84366e3559cf7cb4448d0f5c452da09b670802 Binary files /dev/null and b/src/models/autoencoder/__pycache__/vae.cpython-311.pyc differ diff --git a/src/models/autoencoder/blocks.py b/src/models/autoencoder/blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..ee7dc7c127da59480fd5699e957c1d5c3e2fd596 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/models/autoencoder/decoder.py b/src/models/autoencoder/decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..8c537bf7f27bbee1e80944f72cca2a67269af838 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/models/autoencoder/distributions.py b/src/models/autoencoder/distributions.py new file mode 100644 index 0000000000000000000000000000000000000000..276ee25347830b6f28e08f300c1c79dfdba122de --- /dev/null +++ b/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]) \ No newline at end of file diff --git a/src/models/autoencoder/encoder.py b/src/models/autoencoder/encoder.py new file mode 100644 index 0000000000000000000000000000000000000000..3221999745e2e6f8bcf00a47212105952d7fa4db --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/models/autoencoder/vae.py b/src/models/autoencoder/vae.py new file mode 100644 index 0000000000000000000000000000000000000000..ca90152e03036f1abd3f6511f63995766fc21454 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/src/models/conditioning/__pycache__/clip_text.cpython-311.pyc b/src/models/conditioning/__pycache__/clip_text.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..18f42b0da66cf9940d70cfd2bfdf5a232f6328fa Binary files /dev/null and b/src/models/conditioning/__pycache__/clip_text.cpython-311.pyc differ diff --git a/src/models/conditioning/clip_text.py b/src/models/conditioning/clip_text.py new file mode 100644 index 0000000000000000000000000000000000000000..5fce0225076d0856c8c9320f9bcb9862e8142721 --- /dev/null +++ b/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 \ No newline at end of file diff --git a/src/models/conditioning/null_conditioning.py b/src/models/conditioning/null_conditioning.py new file mode 100644 index 0000000000000000000000000000000000000000..a8893ea3b6493935e0f0093d8a27a0925d1164f9 --- /dev/null +++ b/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, + } \ No newline at end of file diff --git a/src/models/diffusion/__pycache__/attention.cpython-311.pyc b/src/models/diffusion/__pycache__/attention.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..c352faeaa189f3504e15ffabf23fa55b8d33422c Binary files /dev/null and b/src/models/diffusion/__pycache__/attention.cpython-311.pyc differ diff --git a/src/models/diffusion/__pycache__/blocks.cpython-311.pyc b/src/models/diffusion/__pycache__/blocks.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aea8c7dad4206d76670d02a2649d49aa9309a224 Binary files /dev/null and b/src/models/diffusion/__pycache__/blocks.cpython-311.pyc differ diff --git a/src/models/diffusion/__pycache__/timestep.cpython-311.pyc b/src/models/diffusion/__pycache__/timestep.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e80d8e5d2931768faffac6cd9ea87dc42226f004 Binary files /dev/null and b/src/models/diffusion/__pycache__/timestep.cpython-311.pyc differ diff --git a/src/models/diffusion/__pycache__/unet.cpython-311.pyc b/src/models/diffusion/__pycache__/unet.cpython-311.pyc new file mode 100644 index 0000000000000000000000000000000000000000..aee6de77ee0c306ee214a0d195aa7a6d96f13c25 Binary files /dev/null and b/src/models/diffusion/__pycache__/unet.cpython-311.pyc differ diff --git a/src/models/diffusion/attention.py b/src/models/diffusion/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..16a82946dadfb67b89b5bf496c73cece64324f25 --- /dev/null +++ b/src/models/diffusion/attention.py @@ -0,0 +1,324 @@ +from __future__ import annotations + +import math + +import torch +from torch import nn +import torch.nn.functional as F + + +class GEGLU(nn.Module): + """ + Gated GELU feed-forward projection. + + Used in transformer-style attention blocks. + """ + + def __init__( + self, + dim_in: int, + dim_out: int, + ): + super().__init__() + + self.proj = nn.Linear( + dim_in, + dim_out * 2, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x, gate = self.proj(x).chunk(2, dim=-1) + return x * F.gelu(gate) + + +class FeedForward(nn.Module): + """ + Transformer feed-forward block + """ + + def __init__( + self, + dim: int, + mult: int = 4, + dropout: float = 0.0, + ): + super().__init__() + + inner_dim = dim * mult + + self.net = nn.Sequential( + GEGLU(dim, inner_dim), + nn.Dropout(dropout), + nn.Linear(inner_dim, dim), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.net(x) + + +class CrossAttention(nn.Module): + """ + Multi-head attention. + + If context is None: + self-attention + + If context is provided: + cross-attention from x to context. + """ + + def __init__( + self, + query_dim: int, + context_dim: int | None = None, + num_heads: int = 8, + head_dim: int = 64, + dropout: float = 0.0, + ): + super().__init__() + + inner_dim = num_heads * head_dim + context_dim = query_dim if context_dim is None else context_dim + + self.query_dim = query_dim + self.context_dim = context_dim + self.num_heads = num_heads + self.head_dim = head_dim + self.inner_dim = inner_dim + + self.to_q = nn.Linear( + query_dim, + inner_dim, + bias=False, + ) + + self.to_k = nn.Linear( + context_dim, + inner_dim, + bias=False, + ) + + self.to_v = nn.Linear( + context_dim, + inner_dim, + bias=False, + ) + + self.to_out = nn.Sequential( + nn.Linear(inner_dim, query_dim), + nn.Dropout(dropout), + ) + + def forward( + self, + x: torch.Tensor, + context: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + b, n, _ = x.shape + + if context is None: + context = x + + q = self.to_q(x) + k = self.to_k(context) + v = self.to_v(context) + + q = q.view(b, n, self.num_heads, self.head_dim).transpose(1, 2) + k = k.view(b, -1, self.num_heads, self.head_dim).transpose(1, 2) + v = v.view(b, -1, self.num_heads, self.head_dim).transpose(1, 2) + + # q: [B, heads, N, head_dim] + # k: [B, heads, M, head_dim] + # v: [B, heads, M, head_dim] + + attn_mask = None + + if attention_mask is not None: + # attention_mask: [B, M], 1 for valid tokens, 0 for padding. + # scaled_dot_product_attention expects True where attention is allowed + # or additive mask depending on dtype. Bool mask is okay. + attn_mask = attention_mask.bool() + attn_mask = attn_mask[:, None, None, :] # [B, 1, 1, M] + + out = F.scaled_dot_product_attention( + q, + k, + v, + attn_mask=attn_mask, + ) + + out = out.transpose(1, 2).contiguous() + out = out.view(b, n, self.inner_dim) + + out = self.to_out(out) + + return out + + +class BasicTransformerBlock(nn.Module): + """ + Transformer block used inside spatial U-Net feature maps + + it has: + + self-attention + cross-attention + feed-forward + """ + + def __init__( + self, + dim: int, + context_dim: int, + num_heads: int = 8, + head_dim: int = 64, + dropout: float = 0.0, + ): + super().__init__() + + self.norm1 = nn.LayerNorm(dim) + self.self_attn = CrossAttention( + query_dim=dim, + context_dim=None, + num_heads=num_heads, + head_dim=head_dim, + dropout=dropout, + ) + + self.norm2 = nn.LayerNorm(dim) + self.cross_attn = CrossAttention( + query_dim=dim, + context_dim=context_dim, + num_heads=num_heads, + head_dim=head_dim, + dropout=dropout, + ) + + self.norm3 = nn.LayerNorm(dim) + self.ff = FeedForward( + dim=dim, + mult=4, + dropout=dropout, + ) + + def forward( + self, + x: torch.Tensor, + context: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + x = x + self.self_attn( + self.norm1(x), + context=None, + ) + + x = x + self.cross_attn( + self.norm2(x), + context=context, + attention_mask=attention_mask, + ) + + x = x + self.ff( + self.norm3(x), + ) + + return x + + +class SpatialTransformer(nn.Module): + """ + Applies transformer attention on 2D feature maps. + This is where text conditioning enters the U-Net. + """ + + def __init__( + self, + channels: int, + context_dim: int, + num_heads: int = 8, + head_dim: int = 64, + depth: int = 1, + dropout: float = 0.0, + ): + super().__init__() + + self.channels = channels + self.context_dim = context_dim + self.num_heads = num_heads + self.head_dim = head_dim + self.depth = depth + + self.norm = nn.GroupNorm( + num_groups=32, + num_channels=channels, + eps=1e-6, + affine=True, + ) + + inner_dim = num_heads * head_dim + + self.proj_in = nn.Conv2d( + channels, + inner_dim, + kernel_size=1, + stride=1, + padding=0, + ) + + self.transformer_blocks = nn.ModuleList( + [ + BasicTransformerBlock( + dim=inner_dim, + context_dim=context_dim, + num_heads=num_heads, + head_dim=head_dim, + dropout=dropout, + ) + for _ in range(depth) + ] + ) + + self.proj_out = nn.Conv2d( + inner_dim, + channels, + kernel_size=1, + stride=1, + padding=0, + ) + + # Stable residual start. + nn.init.zeros_(self.proj_out.weight) + nn.init.zeros_(self.proj_out.bias) + + def forward( + self, + x: torch.Tensor, + context: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + b, c, h, w = x.shape + + residual = x + + x = self.norm(x) + x = self.proj_in(x) + + inner_dim = x.shape[1] + + x = x.permute(0, 2, 3, 1).contiguous() + x = x.view(b, h * w, inner_dim) + + for block in self.transformer_blocks: + x = block( + x, + context=context, + attention_mask=attention_mask, + ) + + x = x.view(b, h, w, inner_dim) + x = x.permute(0, 3, 1, 2).contiguous() + + x = self.proj_out(x) + + return x + residual \ No newline at end of file diff --git a/src/models/diffusion/blocks.py b/src/models/diffusion/blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..1b90644f6502f20dbdfb2729906e57d4c1344fa3 --- /dev/null +++ b/src/models/diffusion/blocks.py @@ -0,0 +1,479 @@ +from __future__ import annotations + +import torch +from torch import nn +import torch.nn.functional as F + +from src.models.diffusion.attention import SpatialTransformer + + +def normalization( + channels: int, + num_groups: int = 32, +) -> nn.GroupNorm: + """ + GroupNorm + """ + num_groups = min(num_groups, channels) + + while channels % num_groups != 0: + num_groups -= 1 + + return nn.GroupNorm( + num_groups=num_groups, + num_channels=channels, + eps=1e-6, + affine=True, + ) + + +class TimeResBlock(nn.Module): + """ + Residual block conditioned on timestep embedding. + Time embedding is projected and added after the first conv. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + time_embed_dim: int, + dropout: float = 0.0, + ): + super().__init__() + + self.in_channels = in_channels + self.out_channels = out_channels + self.time_embed_dim = time_embed_dim + + self.norm1 = normalization(in_channels) + self.conv1 = nn.Conv2d( + in_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + ) + + self.time_proj = nn.Linear( + time_embed_dim, + out_channels, + ) + + 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() + + def forward( + self, + x: torch.Tensor, + time_emb: torch.Tensor, + ) -> torch.Tensor: + residual = self.shortcut(x) + + h = self.norm1(x) + h = F.silu(h) + h = self.conv1(h) + + time_out = self.time_proj( + F.silu(time_emb), + ) + + h = h + time_out[:, :, None, None] + + 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, + ) + + def forward( + self, + x: torch.Tensor, + ) -> torch.Tensor: + return self.conv(x) + + +class Upsample(nn.Module): + """ + Upsample by factor 2 using nearest-neighbor + conv + """ + + def __init__( + self, + channels: int, + ): + super().__init__() + + self.conv = nn.Conv2d( + channels, + channels, + kernel_size=3, + stride=1, + padding=1, + ) + + 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 AttentionBlock(nn.Module): + """ + Optional text-conditioned attention block. + + If use_attention=True: + applies SpatialTransformer. + + If use_attention=False: + identity. + """ + + def __init__( + self, + channels: int, + context_dim: int, + num_heads: int, + head_dim: int, + transformer_depth: int = 1, + dropout: float = 0.0, + use_attention: bool = True, + ): + super().__init__() + + if use_attention: + self.block = SpatialTransformer( + channels=channels, + context_dim=context_dim, + num_heads=num_heads, + head_dim=head_dim, + depth=transformer_depth, + dropout=dropout, + ) + else: + self.block = nn.Identity() + + def forward( + self, + x: torch.Tensor, + context: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + if isinstance(self.block, nn.Identity): + return x + + if context is None: + raise ValueError("AttentionBlock requires context when use_attention=True.") + + return self.block( + x, + context=context, + attention_mask=attention_mask, + ) + + +class DownBlock(nn.Module): + """ + U-Net down block. + + Contains: + ResBlock(s) + optional SpatialTransformer(s) + optional downsample + + Returns: + x + skip features + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + time_embed_dim: int, + num_res_blocks: int, + context_dim: int, + num_heads: int, + head_dim: int, + transformer_depth: int = 1, + dropout: float = 0.0, + use_attention: bool = False, + add_downsample: bool = True, + ): + super().__init__() + + self.resblocks = nn.ModuleList() + self.attentions = nn.ModuleList() + + current_channels = in_channels + + for _ in range(num_res_blocks): + self.resblocks.append( + TimeResBlock( + in_channels=current_channels, + out_channels=out_channels, + time_embed_dim=time_embed_dim, + dropout=dropout, + ) + ) + + self.attentions.append( + AttentionBlock( + channels=out_channels, + context_dim=context_dim, + num_heads=num_heads, + head_dim=head_dim, + transformer_depth=transformer_depth, + dropout=dropout, + use_attention=use_attention, + ) + ) + + current_channels = out_channels + + if add_downsample: + self.downsample = Downsample(out_channels) + else: + self.downsample = nn.Identity() + + self.out_channels = out_channels + self.add_downsample = add_downsample + + def forward( + self, + x: torch.Tensor, + time_emb: torch.Tensor, + context: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + skips = [] + + for resblock, attention in zip(self.resblocks, self.attentions): + x = resblock( + x, + time_emb, + ) + + x = attention( + x, + context=context, + attention_mask=attention_mask, + ) + + skips.append(x) + + x = self.downsample(x) + + return x, skips + + +class UpBlock(nn.Module): + """ + U-Net up block. + + Takes skip features from encoder path. + """ + + def __init__( + self, + in_channels: int, + skip_channels: int, + out_channels: int, + time_embed_dim: int, + num_res_blocks: int, + context_dim: int, + num_heads: int, + head_dim: int, + transformer_depth: int = 1, + dropout: float = 0.0, + use_attention: bool = False, + add_upsample: bool = True, + ): + super().__init__() + + self.resblocks = nn.ModuleList() + self.attentions = nn.ModuleList() + + current_channels = in_channels + + for _ in range(num_res_blocks): + self.resblocks.append( + TimeResBlock( + in_channels=current_channels + skip_channels, + out_channels=out_channels, + time_embed_dim=time_embed_dim, + dropout=dropout, + ) + ) + + self.attentions.append( + AttentionBlock( + channels=out_channels, + context_dim=context_dim, + num_heads=num_heads, + head_dim=head_dim, + transformer_depth=transformer_depth, + dropout=dropout, + use_attention=use_attention, + ) + ) + + current_channels = out_channels + + if add_upsample: + self.upsample = Upsample(out_channels) + else: + self.upsample = nn.Identity() + + self.out_channels = out_channels + self.add_upsample = add_upsample + + def forward( + self, + x: torch.Tensor, + skips: list[torch.Tensor], + time_emb: torch.Tensor, + context: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + for resblock, attention in zip(self.resblocks, self.attentions): + if len(skips) == 0: + raise RuntimeError("Not enough skip connections for UpBlock.") + + skip = skips.pop() + + x = torch.cat( + [x, skip], + dim=1, + ) + + x = resblock( + x, + time_emb, + ) + + x = attention( + x, + context=context, + attention_mask=attention_mask, + ) + + x = self.upsample(x) + + return x + + +class MiddleBlock(nn.Module): + """ + U-Net bottleneck block + """ + + def __init__( + self, + channels: int, + time_embed_dim: int, + context_dim: int, + num_heads: int, + head_dim: int, + transformer_depth: int = 1, + dropout: float = 0.0, + use_attention: bool = True, + ): + super().__init__() + + self.res1 = TimeResBlock( + in_channels=channels, + out_channels=channels, + time_embed_dim=time_embed_dim, + dropout=dropout, + ) + + self.attn = AttentionBlock( + channels=channels, + context_dim=context_dim, + num_heads=num_heads, + head_dim=head_dim, + transformer_depth=transformer_depth, + dropout=dropout, + use_attention=use_attention, + ) + + self.res2 = TimeResBlock( + in_channels=channels, + out_channels=channels, + time_embed_dim=time_embed_dim, + dropout=dropout, + ) + + def forward( + self, + x: torch.Tensor, + time_emb: torch.Tensor, + context: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + x = self.res1( + x, + time_emb, + ) + + x = self.attn( + x, + context=context, + attention_mask=attention_mask, + ) + + x = self.res2( + x, + time_emb, + ) + + return x \ No newline at end of file diff --git a/src/models/diffusion/timestep.py b/src/models/diffusion/timestep.py new file mode 100644 index 0000000000000000000000000000000000000000..7f619134ecd5e4981e94bb67e731fc7af0594c01 --- /dev/null +++ b/src/models/diffusion/timestep.py @@ -0,0 +1,73 @@ +from __future__ import annotations + +import math + +import torch +from torch import nn +import torch.nn.functional as F + + +def sinusoidal_timestep_embedding( + timesteps: torch.Tensor, + dim: int, + max_period: int = 10_000, +) -> torch.Tensor: + """ + Create sinusoidal timestep embeddings + """ + half_dim = dim // 2 + + frequencies = torch.exp( + -math.log(max_period) + * torch.arange( + start=0, + end=half_dim, + dtype=torch.float32, + device=timesteps.device, + ) + / half_dim + ) + + args = timesteps.float()[:, None] * frequencies[None] + + embedding = torch.cat( + [ + torch.cos(args), + torch.sin(args), + ], + dim=-1, + ) + + if dim % 2 == 1: + embedding = F.pad(embedding, (0, 1)) + + return embedding + + +class TimestepEmbedding(nn.Module): + + def __init__( + self, + embedding_dim: int, + time_embed_dim: int, + ): + super().__init__() + + self.embedding_dim = embedding_dim + self.time_embed_dim = time_embed_dim + + self.mlp = nn.Sequential( + nn.Linear(embedding_dim, time_embed_dim), + nn.SiLU(), + nn.Linear(time_embed_dim, time_embed_dim), + ) + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + emb = sinusoidal_timestep_embedding( + timesteps=timesteps, + dim=self.embedding_dim, + ) + + emb = self.mlp(emb) + + return emb \ No newline at end of file diff --git a/src/models/diffusion/unet.py b/src/models/diffusion/unet.py new file mode 100644 index 0000000000000000000000000000000000000000..bf3b03c55acd4a1b51396ecfaf216fd739aee8c2 --- /dev/null +++ b/src/models/diffusion/unet.py @@ -0,0 +1,303 @@ +from __future__ import annotations + +import torch +from torch import nn +import torch.nn.functional as F + +from src.models.diffusion.timestep import TimestepEmbedding +from src.models.diffusion.blocks import ( + DownBlock, + MiddleBlock, + UpBlock, + normalization, +) + + +class LatentDiffusionUNet(nn.Module): + """ + Lightweight text-conditioned U-Net for latent diffusion. + + Input: + z_t: + noisy latent [B, in_channels, H, W] + + timesteps: + diffusion timestep [B] + + context: + CLIP token embeddings [B, seq_len, context_dim] + + Output: + prediction [B, out_channels, H, W] + """ + + def __init__( + self, + in_channels: int = 8, + out_channels: int = 8, + latent_size: int = 32, + base_channels: int = 128, + channel_multipliers: list[int] | tuple[int, ...] = (1, 2, 3), + num_res_blocks: int = 2, + attention_resolutions: list[int] | tuple[int, ...] = (16, 8), + context_dim: int = 768, + num_heads: int = 4, + head_dim: int = 32, + transformer_depth: int = 1, + dropout: float = 0.0, + time_embedding_dim: int | None = None, + use_middle_attention: bool = True, + ): + super().__init__() + + self.in_channels = in_channels + self.out_channels = out_channels + self.latent_size = latent_size + self.base_channels = base_channels + self.channel_multipliers = list(channel_multipliers) + self.num_res_blocks = num_res_blocks + self.attention_resolutions = set(attention_resolutions) + self.context_dim = context_dim + self.num_heads = num_heads + self.head_dim = head_dim + self.transformer_depth = transformer_depth + self.dropout = dropout + + if time_embedding_dim is None: + time_embedding_dim = base_channels * 4 + + self.time_embedding_dim = time_embedding_dim + + self.time_embed = TimestepEmbedding( + embedding_dim=base_channels, + time_embed_dim=time_embedding_dim, + ) + + self.conv_in = nn.Conv2d( + in_channels, + base_channels, + kernel_size=3, + stride=1, + padding=1, + ) + + channels_per_level = [ + base_channels * multiplier + for multiplier in self.channel_multipliers + ] + + # ------------------------- + # Down path + # ------------------------- + self.down_blocks = nn.ModuleList() + + current_channels = base_channels + current_resolution = latent_size + + self.down_block_channels: list[int] = [] + self.down_block_resolutions: list[int] = [] + + for level, out_channels_level in enumerate(channels_per_level): + is_last = level == len(channels_per_level) - 1 + + use_attention = current_resolution in self.attention_resolutions + + block = DownBlock( + in_channels=current_channels, + out_channels=out_channels_level, + time_embed_dim=time_embedding_dim, + num_res_blocks=num_res_blocks, + context_dim=context_dim, + num_heads=num_heads, + head_dim=head_dim, + transformer_depth=transformer_depth, + dropout=dropout, + use_attention=use_attention, + add_downsample=not is_last, + ) + + self.down_blocks.append(block) + + self.down_block_channels.append(out_channels_level) + self.down_block_resolutions.append(current_resolution) + + current_channels = out_channels_level + + if not is_last: + current_resolution //= 2 + + # ------------------------- + # Middle + # ------------------------- + self.middle = MiddleBlock( + channels=current_channels, + time_embed_dim=time_embedding_dim, + context_dim=context_dim, + num_heads=num_heads, + head_dim=head_dim, + transformer_depth=transformer_depth, + dropout=dropout, + use_attention=use_middle_attention, + ) + + # ------------------------- + # Up path + # ------------------------- + self.up_blocks = nn.ModuleList() + + reversed_channels = list(reversed(self.down_block_channels)) + reversed_resolutions = list(reversed(self.down_block_resolutions)) + + for level, (skip_channels, resolution) in enumerate( + zip(reversed_channels, reversed_resolutions) + ): + is_last = level == len(reversed_channels) - 1 + + out_channels_level = skip_channels + use_attention = resolution in self.attention_resolutions + + block = UpBlock( + in_channels=current_channels, + skip_channels=skip_channels, + out_channels=out_channels_level, + time_embed_dim=time_embedding_dim, + num_res_blocks=num_res_blocks, + context_dim=context_dim, + num_heads=num_heads, + head_dim=head_dim, + transformer_depth=transformer_depth, + dropout=dropout, + use_attention=use_attention, + add_upsample=not is_last, + ) + + self.up_blocks.append(block) + + current_channels = out_channels_level + + self.norm_out = normalization(current_channels) + + self.conv_out = nn.Conv2d( + current_channels, + out_channels, + kernel_size=3, + stride=1, + padding=1, + ) + + # Stable start: initially predicts near zero. + nn.init.zeros_(self.conv_out.weight) + nn.init.zeros_(self.conv_out.bias) + + def forward( + self, + z_t: torch.Tensor, + timesteps: torch.Tensor, + context: torch.Tensor, + attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """ + Args: + z_t: + Noisy latent [B, C, H, W] + + timesteps: + Diffusion timesteps [B] + + context: + CLIP token embeddings [B, seq_len, context_dim] + + attention_mask: + CLIP token mask [B, seq_len], optional + + Returns: + model_output: + [B, out_channels, H, W] + """ + if z_t.ndim != 4: + raise ValueError(f"z_t must be [B, C, H, W], got {z_t.shape}") + + if timesteps.ndim != 1: + raise ValueError(f"timesteps must be [B], got {timesteps.shape}") + + time_emb = self.time_embed(timesteps) + + x = self.conv_in(z_t) + + skips: list[torch.Tensor] = [] + + for block in self.down_blocks: + x, block_skips = block( + x=x, + time_emb=time_emb, + context=context, + attention_mask=attention_mask, + ) + + skips.extend(block_skips) + + x = self.middle( + x=x, + time_emb=time_emb, + context=context, + attention_mask=attention_mask, + ) + + for block in self.up_blocks: + x = block( + x=x, + skips=skips, + time_emb=time_emb, + context=context, + attention_mask=attention_mask, + ) + + if len(skips) != 0: + raise RuntimeError( + f"Unused skip connections remain: {len(skips)}. " + "Check U-Net block construction." + ) + + x = self.norm_out(x) + x = F.silu(x) + x = self.conv_out(x) + + return x + + +def count_parameters( + model: nn.Module, + trainable_only: bool = True, +) -> int: + if trainable_only: + return sum(p.numel() for p in model.parameters() if p.requires_grad) + + return sum(p.numel() for p in model.parameters()) + + +def build_latent_diffusion_unet_from_config(cfg: dict) -> LatentDiffusionUNet: + """ + Build U-Net from config dictionary. + + Expects: + + cfg["model"] + """ + m = cfg["model"] + + return LatentDiffusionUNet( + in_channels=int(m["in_channels"]), + out_channels=int(m["out_channels"]), + latent_size=int(m.get("latent_size", 32)), + base_channels=int(m["base_channels"]), + channel_multipliers=tuple(m["channel_multipliers"]), + num_res_blocks=int(m["num_res_blocks"]), + attention_resolutions=tuple(m.get("attention_resolutions", [16, 8])), + context_dim=int(m["context_dim"]), + num_heads=int(m["num_heads"]), + head_dim=int(m["head_dim"]), + transformer_depth=int(m.get("transformer_depth", 1)), + dropout=float(m.get("dropout", 0.0)), + time_embedding_dim=m.get("time_embedding_dim", None), + use_middle_attention=bool(m.get("use_middle_attention", True)), + ) \ No newline at end of file diff --git a/src/utils/__init__.py b/src/utils/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/utils/config.py b/src/utils/config.py new file mode 100644 index 0000000000000000000000000000000000000000..e1aef88602ae1081b80100dab1cf920b31e72325 --- /dev/null +++ b/src/utils/config.py @@ -0,0 +1,79 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import yaml + + +def load_yaml(path: str | Path) -> dict[str, Any]: + path = Path(path) + with open(path, "r", encoding="utf-8") as f: + return yaml.safe_load(f) + + +def deep_update(base: dict, update: dict) -> dict: + """ + Recursively merge update into base. + """ + for key, value in update.items(): + if ( + key in base + and isinstance(base[key], dict) + and isinstance(value, dict) + ): + deep_update(base[key], value) + else: + base[key] = value + return base + + +def load_config(config_path: str | Path) -> dict[str, Any]: + """ + Load a config file and recursively merge files listed under `defaults`. + + Example: + defaults: + - ../paths.yaml + - ../data/coco_256.yaml + """ + config_path = Path(config_path) + cfg = load_yaml(config_path) + + merged: dict[str, Any] = {} + + defaults = cfg.pop("defaults", []) + for default_path in defaults: + default_file = (config_path.parent / default_path).resolve() + default_cfg = load_config(default_file) + merged = deep_update(merged, default_cfg) + + merged = deep_update(merged, cfg) + return merged + + +def get_by_key(cfg: dict[str, Any], key: str) -> Any: + """ + Get nested config value using dot notation. + + Example: + get_by_key(cfg, "data.coco_root") + """ + value: Any = cfg + for part in key.split("."): + value = value[part] + return value + + +def resolve_path_key(cfg: dict[str, Any], key: str) -> Path: + """ + Resolve a path stored in config using dot notation. + """ + return Path(get_by_key(cfg, key)).expanduser().resolve() + +def save_yaml(data: dict, path: str | Path) -> None: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + + with open(path, "w", encoding="utf-8") as f: + yaml.safe_dump(data, f, sort_keys=False) \ No newline at end of file diff --git a/src/utils/gpu_monitor.py b/src/utils/gpu_monitor.py new file mode 100644 index 0000000000000000000000000000000000000000..a0fb2cc0420c0cbea85f537459af7bfea3347bb7 --- /dev/null +++ b/src/utils/gpu_monitor.py @@ -0,0 +1,157 @@ +from __future__ import annotations + +import json +import subprocess +import threading +import time +from contextlib import contextmanager +from pathlib import Path +from typing import Any + +import torch + + +class GPUMonitor: + + def __init__( + self, + enabled: bool = True, + device: torch.device | str | None = None, + sync_cuda: bool = True, + sample_interval_s: float = 0.25, + use_nvidia_smi: bool = True, + ): + self.enabled = bool(enabled) and torch.cuda.is_available() + self.device = torch.device(device or "cuda") + self.sync_cuda = bool(sync_cuda) + self.sample_interval_s = float(sample_interval_s) + self.use_nvidia_smi = bool(use_nvidia_smi) + self.records: dict[str, list[dict[str, Any]]] = {} + + if self.enabled and self.device.type == "cuda": + self.device_index = self.device.index + if self.device_index is None: + self.device_index = torch.cuda.current_device() + else: + self.device_index = None + + def _sync(self): + if self.enabled and self.sync_cuda: + torch.cuda.synchronize(self.device) + + def _mem_stats_mb(self) -> dict[str, float]: + if not self.enabled: + return {} + return { + "allocated_mb": torch.cuda.memory_allocated(self.device) / 1024**2, + "reserved_mb": torch.cuda.memory_reserved(self.device) / 1024**2, + "max_allocated_mb": torch.cuda.max_memory_allocated(self.device) / 1024**2, + "max_reserved_mb": torch.cuda.max_memory_reserved(self.device) / 1024**2, + } + + def _query_nvidia_smi(self) -> tuple[float | None, float | None]: + if not (self.enabled and self.use_nvidia_smi and self.device_index is not None): + return None, None + try: + out = subprocess.check_output( + [ + "nvidia-smi", + f"--id={self.device_index}", + "--query-gpu=utilization.gpu,memory.used", + "--format=csv,noheader,nounits", + ], + text=True, + stderr=subprocess.DEVNULL, + timeout=1.0, + ).strip() + first_line = out.splitlines()[0] + util_s, mem_s = [x.strip() for x in first_line.split(",")[:2]] + return float(util_s), float(mem_s) + except Exception: + return None, None + + def _sampler(self, stop: threading.Event, samples: list[dict[str, float]]): + while not stop.is_set(): + util, mem = self._query_nvidia_smi() + if util is not None or mem is not None: + samples.append({ + "t": time.perf_counter(), + "gpu_util_percent": util, + "memory_used_mb": mem, + }) + stop.wait(self.sample_interval_s) + + @contextmanager + def track(self, name: str): + if not self.enabled: + yield + return + + self._sync() + torch.cuda.reset_peak_memory_stats(self.device) + before = self._mem_stats_mb() + start = time.perf_counter() + + samples: list[dict[str, float]] = [] + stop = threading.Event() + thread = None + if self.use_nvidia_smi: + thread = threading.Thread(target=self._sampler, args=(stop, samples), daemon=True) + thread.start() + + try: + yield + finally: + self._sync() + end = time.perf_counter() + stop.set() + if thread is not None: + thread.join(timeout=1.0) + after = self._mem_stats_mb() + + gpu_utils = [s["gpu_util_percent"] for s in samples if s.get("gpu_util_percent") is not None] + smi_mem = [s["memory_used_mb"] for s in samples if s.get("memory_used_mb") is not None] + + record: dict[str, Any] = { + "seconds": end - start, + "torch_allocated_before_mb": before.get("allocated_mb"), + "torch_allocated_after_mb": after.get("allocated_mb"), + "torch_allocated_delta_mb": after.get("allocated_mb", 0.0) - before.get("allocated_mb", 0.0), + "torch_reserved_before_mb": before.get("reserved_mb"), + "torch_reserved_after_mb": after.get("reserved_mb"), + "torch_reserved_delta_mb": after.get("reserved_mb", 0.0) - before.get("reserved_mb", 0.0), + "torch_peak_allocated_mb": after.get("max_allocated_mb"), + "torch_peak_reserved_mb": after.get("max_reserved_mb"), + "nvidia_smi_num_samples": len(samples), + "nvidia_smi_gpu_util_avg_percent": sum(gpu_utils) / len(gpu_utils) if gpu_utils else None, + "nvidia_smi_gpu_util_max_percent": max(gpu_utils) if gpu_utils else None, + "nvidia_smi_memory_used_avg_mb": sum(smi_mem) / len(smi_mem) if smi_mem else None, + "nvidia_smi_memory_used_max_mb": max(smi_mem) if smi_mem else None, + } + self.records.setdefault(name, []).append(record) + + def summary(self) -> dict[str, Any]: + return self.records + + def save_json(self, path: str | Path): + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + json.dump(self.summary(), f, indent=2) + + def print_summary(self, title: str = "GPU usage"): + if not self.enabled: + print(f"{title}: CUDA unavailable or GPU monitor disabled.") + return + print(f"\n{title}") + for name, items in self.records.items(): + latest = items[-1] + util = latest.get("nvidia_smi_gpu_util_avg_percent") + util_s = "n/a" if util is None else f"{util:.1f}% avg" + print( + f" {name}: " + f"peak_alloc={latest['torch_peak_allocated_mb']:.1f} MB, " + f"peak_reserved={latest['torch_peak_reserved_mb']:.1f} MB, " + f"delta_alloc={latest['torch_allocated_delta_mb']:+.1f} MB, " + f"gpu_util={util_s}" + ) diff --git a/src/utils/timer.py b/src/utils/timer.py new file mode 100644 index 0000000000000000000000000000000000000000..c39856adfe14a9556751fb69314b85fc501dc0db --- /dev/null +++ b/src/utils/timer.py @@ -0,0 +1,147 @@ +from __future__ import annotations + +import json +import time +from contextlib import ContextDecorator +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + + +@dataclass +class TimerRecord: + name: str + total: float = 0.0 + count: int = 0 + + @property + def avg(self) -> float: + return self.total / max(1, self.count) + + +class Timer(ContextDecorator): + """ + It can be used as: + timer = Timer(sync_cuda=True) + + with timer("sample"): + images = sample() + + timer.print_summary() + timer.save_json("timing.json") + """ + + def __init__( + self, + name: Optional[str] = None, + records: Optional[dict[str, TimerRecord]] = None, + enabled: bool = True, + sync_cuda: bool = True, + verbose: bool = False, + ): + self.name = name + self.records = records if records is not None else {} + self.enabled = enabled + self.sync_cuda = sync_cuda + self.verbose = verbose + self._start_time: Optional[float] = None + + def __call__(self, name: str): + return Timer( + name=name, + records=self.records, + enabled=self.enabled, + sync_cuda=self.sync_cuda, + verbose=self.verbose, + ) + + def __enter__(self): + if not self.enabled: + return self + self._sync() + self._start_time = time.perf_counter() + return self + + def __exit__(self, exc_type, exc, tb): + if not self.enabled: + return False + self._sync() + if self._start_time is None: + elapsed = 0.0 + else: + elapsed = time.perf_counter() - self._start_time + + name = self.name or "unnamed" + if name not in self.records: + self.records[name] = TimerRecord(name=name) + + record = self.records[name] + record.total += elapsed + record.count += 1 + + if self.verbose: + print(f"[timer] {name}: {elapsed:.4f}s", flush=True) + + return False + + def reset(self) -> None: + self.records.clear() + + def get(self, name: str) -> TimerRecord: + return self.records[name] + + def summary(self) -> dict[str, dict[str, float | int]]: + return { + name: { + "total_sec": record.total, + "count": record.count, + "avg_sec": record.avg, + } + for name, record in self.records.items() + } + + def save_json(self, path: str | Path) -> None: + path = Path(path) + path.parent.mkdir(parents=True, exist_ok=True) + with open(path, "w", encoding="utf-8") as f: + json.dump(self.summary(), f, indent=2) + + def print_summary(self, sort_by: str = "total", title: str = "Timing summary") -> None: + if not self.records: + print(f"{title}: no records", flush=True) + return + + if sort_by == "total": + key_fn = lambda item: item[1].total + elif sort_by == "avg": + key_fn = lambda item: item[1].avg + elif sort_by == "count": + key_fn = lambda item: item[1].count + else: + raise ValueError("sort_by must be 'total', 'avg', or 'count'.") + + items = sorted(self.records.items(), key=key_fn, reverse=True) + + print("=" * 72, flush=True) + print(title, flush=True) + print("=" * 72, flush=True) + print(f"{'operation':36s} {'count':>8s} {'total(s)':>12s} {'avg(s)':>12s}", flush=True) + print("-" * 72, flush=True) + + for name, record in items: + print( + f"{name:36s} {record.count:8d} {record.total:12.4f} {record.avg:12.4f}", + flush=True, + ) + + print("=" * 72, flush=True) + + def _sync(self) -> None: + if not self.sync_cuda: + return + try: + import torch + if torch.cuda.is_available(): + torch.cuda.synchronize() + except ImportError: + pass diff --git a/src/utils/upload_best_pt_to_HF.py b/src/utils/upload_best_pt_to_HF.py new file mode 100644 index 0000000000000000000000000000000000000000..217878d165ac5da24254bdd303e97d175bf15623 --- /dev/null +++ b/src/utils/upload_best_pt_to_HF.py @@ -0,0 +1,24 @@ +from dotenv import load_dotenv +from huggingface_hub import HfApi +import os + +load_dotenv() + +token = os.getenv("HF_TOKEN") +if token is None: + raise RuntimeError("HF_TOKEN not found. Check your .env file.") + +repo_id = "anilegin/lightweight-diffusion-vae" + +api = HfApi(token=token) + +api.upload_file( + path_or_fileobj="/leonardo_scratch/large/userexternal/aegin000/models/vae/best.pt", + path_in_repo="best.pt", + repo_id=repo_id, + repo_type="model", + token=token, + commit_message="Replace best.pt with inference-only VAE weights", +) + +print("Uploaded clean best.pt to Hugging Face.")