novelkathor/DDIM_Image_Generation

1"""Visualization helpers: image grids, denoising-trajectory GIFs, and2latent-interpolation grids.3 4All functions accept tensors in the [-1, 1] range (model output convention)5unless otherwise stated, and write/return uint8 arrays in [0, 255].6"""7from __future__ import annotations8 9import math10import os11from typing import List, Optional, Sequence12 13import numpy as np14import torch15from PIL import Image16 17 18# ---------------------------------------------------------------------------19# Small primitives20# ---------------------------------------------------------------------------21def to_uint8(x: torch.Tensor) -> np.ndarray:22    """Tensor in [-1, 1] (B,3,H,W) or (3,H,W) -> uint8 numpy (H,W,3) or (B,H,W,3)."""23    x = x.detach().to(torch.float32).cpu()24    x = (x.clamp(-1.0, 1.0) + 1.0) * 127.525    x = x.round().clamp(0, 255).to(torch.uint8)26    if x.ndim == 4:27        return x.permute(0, 2, 3, 1).numpy()                # (B,H,W,3)28    if x.ndim == 3:29        return x.permute(1, 2, 0).numpy()                   # (H,W,3)30    raise ValueError(f"unsupported shape {x.shape}")31 32 33def make_grid(images: torch.Tensor, nrow: Optional[int] = None, pad: int = 2,34              pad_value: float = 1.0) -> np.ndarray:35    """Lay a batch of images out as a grid. Inputs in [-1, 1].36 37    Returns uint8 (H, W, 3).38    """39    if images.ndim != 4:40        raise ValueError(f"expected (B,C,H,W), got {images.shape}")41    B, C, H, W = images.shape42    if nrow is None:43        nrow = int(math.ceil(math.sqrt(B)))44    ncol = int(math.ceil(B / nrow))45 46    grid_h = ncol * H + (ncol + 1) * pad47    grid_w = nrow * W + (nrow + 1) * pad48 49    grid = torch.full((C, grid_h, grid_w), pad_value, dtype=images.dtype)50    for i in range(B):51        r, c = divmod(i, nrow)52        y = pad + r * (H + pad)53        x = pad + c * (W + pad)54        grid[:, y:y + H, x:x + W] = images[i]55 56    return to_uint8(grid)57 58 59def save_image_grid(images: torch.Tensor, path: str, nrow: Optional[int] = None) -> str:60    arr = make_grid(images, nrow=nrow)61    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)62    Image.fromarray(arr).save(path)63    return path64 65 66# ---------------------------------------------------------------------------67# Denoising trajectory GIF68# ---------------------------------------------------------------------------69def trajectory_to_gif(70    trajectory: Sequence[torch.Tensor],71    path: str,72    fps: int = 10,73    nrow: Optional[int] = None,74) -> str:75    """Save a list of tensors (each (B,C,H,W) in [-1,1]) as an animated GIF.76 77    Each frame is laid out as a grid of all batch items.78    """79    import imageio.v2 as imageio       # local import; heavy dep80 81    frames = []82    for x in trajectory:83        if x.ndim == 3:84            x = x.unsqueeze(0)85        frames.append(make_grid(x, nrow=nrow))86 87    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)88    duration = 1.0 / max(fps, 1)89    imageio.mimsave(path, frames, format="GIF", duration=duration, loop=0)90    return path91 92 93# ---------------------------------------------------------------------------94# Latent interpolation95# ---------------------------------------------------------------------------96def slerp(z1: torch.Tensor, z2: torch.Tensor, t: float) -> torch.Tensor:97    """Spherical linear interpolation between two same-shape latents.98 99    Falls back to lerp if vectors are nearly colinear (avoids div-by-zero).100    """101    flat1 = z1.flatten(start_dim=0)102    flat2 = z2.flatten(start_dim=0)103    dot = (flat1 * flat2).sum() / (flat1.norm() * flat2.norm() + 1e-12)104    dot = dot.clamp(-1.0, 1.0)105    omega = torch.acos(dot)106    sin_omega = torch.sin(omega)107    if sin_omega.abs() < 1e-6:108        return (1 - t) * z1 + t * z2109    a = torch.sin((1 - t) * omega) / sin_omega110    b = torch.sin(t * omega) / sin_omega111    return a * z1 + b * z2112 113 114def interpolate_latents(z1: torch.Tensor, z2: torch.Tensor, num_steps: int = 8,115                        method: str = "slerp") -> torch.Tensor:116    """Return a tensor of shape (num_steps, *z1.shape) of interpolated latents."""117    ts = torch.linspace(0.0, 1.0, num_steps)118    out = []119    for t in ts:120        if method == "slerp":121            out.append(slerp(z1, z2, t.item()))122        elif method == "lerp":123            out.append((1 - t) * z1 + t * z2)124        else:125            raise ValueError(method)126    return torch.stack(out, dim=0)127 128 129# ---------------------------------------------------------------------------130# Self-test131# ---------------------------------------------------------------------------132if __name__ == "__main__":133    import tempfile134 135    torch.manual_seed(0)136    imgs = torch.randn(8, 3, 32, 32).clamp(-1, 1)137 138    grid = make_grid(imgs, nrow=4)139    assert grid.dtype == np.uint8 and grid.ndim == 3 and grid.shape[2] == 3140 141    with tempfile.TemporaryDirectory() as td:142        p1 = save_image_grid(imgs, os.path.join(td, "g.png"))143        assert os.path.exists(p1)144 145        traj = [torch.randn(4, 3, 16, 16).clamp(-1, 1) for _ in range(6)]146        p2 = trajectory_to_gif(traj, os.path.join(td, "t.gif"), fps=8, nrow=2)147        assert os.path.exists(p2) and os.path.getsize(p2) > 0148 149    z1 = torch.randn(1, 3, 16, 16)150    z2 = torch.randn(1, 3, 16, 16)151    interps = interpolate_latents(z1, z2, num_steps=5, method="slerp")152    assert interps.shape == (5, 1, 3, 16, 16)153    # endpoints recovered154    assert torch.allclose(interps[0], z1, atol=1e-5)155    assert torch.allclose(interps[-1], z2, atol=1e-5)156 157    print("visualize.py: all tests passed")158