inference.py

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 clear_cuda_cache():
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        try:
            torch.cuda.ipc_collect()
        except Exception:
            pass


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):
    if not torch.isfinite(image).all():
        finite = image[torch.isfinite(image)]
        if finite.numel() > 0:
            stats = (
                f"finite_min={finite.min().item():.6f}, "
                f"finite_max={finite.max().item():.6f}, "
                f"finite_mean={finite.mean().item():.6f}"
            )
        else:
            stats = "no finite values"

        raise RuntimeError(
            "Non-finite values found in image tensor before saving. "
            f"{stats}. "
            "This usually means fp16 sampling or VAE decoding became unstable. "
            "Try --precision fp32, lower --guidance-scale, or decode the VAE in fp32."
        )

    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_from_paths(
    image_paths: list[Path],
    path: str | Path,
    nrow: int | None = None,
    padding: int = 2,
):
    if not image_paths:
        return

    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)

    images = [Image.open(p).convert("RGB") for p in image_paths]

    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, (x, y))

    grid.save(path)

    for img in images:
        img.close()


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):

    if not torch.isfinite(latents).all():
        finite = latents[torch.isfinite(latents)]
        if finite.numel() > 0:
            stats = (
                f"finite_min={finite.min().item():.6f}, "
                f"finite_max={finite.max().item():.6f}, "
                f"finite_mean={finite.mean().item():.6f}"
            )
        else:
            stats = "no finite values"

        raise RuntimeError(
            "Non-finite values found in sampled latents before VAE decode. "
            f"{stats}. "
            "Try --precision fp32 or lower --guidance-scale."
        )

    # Explicitly unscale here, using generation_config.yaml's vae.scaling_factor.
    z = latents.float() / float(scaling_factor)

    # Decode VAE in fp32. This avoids fp16 overflow/NaNs on some Colab GPUs.
    if z.is_cuda:
        with torch.autocast("cuda", enabled=False):
            images = vae.decode(z, unscale=False)
    else:
        images = vae.decode(z, unscale=False)

    if hasattr(images, "sample"):
        images = images.sample

    images = images.float()

    if not torch.isfinite(images).all():
        finite = images[torch.isfinite(images)]
        if finite.numel() > 0:
            stats = (
                f"finite_min={finite.min().item():.6f}, "
                f"finite_max={finite.max().item():.6f}, "
                f"finite_mean={finite.mean().item():.6f}"
            )
        else:
            stats = "no finite values"

        raise RuntimeError(
            "Non-finite values produced by VAE decode. "
            f"{stats}. "
            "Try --precision fp32. If this still happens, inspect the sampled latents."
        )

    return ((images + 1.0) / 2.0).clamp(0.0, 1.0)


@torch.no_grad()
def decode_and_save_latents(
    vae,
    latents: torch.Tensor,
    prompts: list[str],
    output_dir: Path,
    global_start_index: int,
    scaling_factor: float,
    dtype: torch.dtype,
    decode_batch_size: int,
) -> list[Path]:
    saved_paths: list[Path] = []
    decode_batch_size = max(1, int(decode_batch_size))

    for local_start in range(0, latents.shape[0], decode_batch_size):
        local_end = min(local_start + decode_batch_size, latents.shape[0])
        latent_chunk = latents[local_start:local_end]

        images = decode_latents(
            vae=vae,
            latents=latent_chunk,
            scaling_factor=scaling_factor,
            dtype=dtype,
        )

        for j, image in enumerate(images):
            local_idx = local_start + j
            prompt = prompts[local_idx]
            global_idx = global_start_index + local_idx
            out_path = output_dir / f"{global_idx:04d}_{sanitize_filename(prompt)}.png"
            save_image_tensor(image, out_path)
            saved_paths.append(out_path)

        del latent_chunk, images
        clear_cuda_cache()

    return saved_paths


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.")

    # Memory-saving options.
    parser.add_argument("--low-vram", action="store_true", help="Offload CLIP/UNet/VAE between stages to reduce peak GPU memory.")
    parser.add_argument("--decode-batch-size", type=int, default=1, help="Number of latents to VAE-decode at once.")
    parser.add_argument("--no-grid", action="store_true", help="Skip grid.png creation to reduce CPU RAM usage.")
    parser.add_argument("--grid-max-images", type=int, default=64, help="Maximum saved images to include in grid.png.")

    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)
    low_vram = bool(args.low_vram and device.type == "cuda")

    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="cpu" if low_vram else 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("Batch size:", batch_size)
    print("Low VRAM:", low_vram)
    print("Decode batch size:", max(1, int(args.decode_batch_size)))
    print("Output dir:", output_dir)
    print("=============================================")

    saved_paths: list[Path] = []

    for batch_start in tqdm(range(0, len(prompts), batch_size), desc="batches"):
        batch_prompts = prompts[batch_start: batch_start + batch_size]

        if low_vram:
            text_encoder.to(device)
            clear_cuda_cache()

        cond_context, uncond_context = encode_contexts(
            text_encoder=text_encoder,
            prompts=batch_prompts,
            empty_text=empty_text,
            device=device,
            dtype=dtype,
        )

        # CLIP is no longer needed once context embeddings are produced.
        if low_vram:
            text_encoder.to("cpu")
            clear_cuda_cache()

        shape = (len(batch_prompts), latent_channels, latent_size, latent_size)

        if low_vram:
            unet.to(device)
            clear_cuda_cache()

        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
        del sample_out, cond_context, uncond_context

        # UNet is no longer needed once latents are sampled.
        if low_vram:
            unet.to("cpu")
            clear_cuda_cache()
            vae.to(device)
            clear_cuda_cache()

        batch_saved = decode_and_save_latents(
            vae=vae,
            latents=latents,
            prompts=batch_prompts,
            output_dir=output_dir,
            global_start_index=batch_start,
            scaling_factor=scaling_factor,
            dtype=dtype,
            decode_batch_size=args.decode_batch_size,
        )
        saved_paths.extend(batch_saved)

        del latents

        if low_vram:
            vae.to("cpu")
            clear_cuda_cache()

        clear_cuda_cache()

    if not args.no_grid:
        grid_paths = saved_paths[: max(0, int(args.grid_max_images))]
        save_image_grid_from_paths(grid_paths, output_dir / "grid.png")

    print("Saved", len(saved_paths), "images to", output_dir)


if __name__ == "__main__":
    main()