Add sampling.py — Euler and Heun ODE samplers

liquidflow/sampling.py

@@ -0,0 +1,91 @@
+"""
+Sampling / inference for LiquidFlow.
+
+Uses ODE integration (Euler or Heun's method) to solve:
+    x_{t+dt} = x_t + v_θ(x_t, t) * dt
+
+Starting from x_0 ~ N(0, I) and integrating to x_1 (clean image).
+"""
+
+import torch
+import torch.nn as nn
+from tqdm import tqdm
+
+
+@torch.no_grad()
+def euler_sample(model, shape, num_steps=50, device='cpu', class_label=None, cfg_scale=0.0):
+    """
+    Generate images using Euler method ODE integration.
+    
+    Flow matching: integrate dx/dt = v_θ(x_t, t) from t=0 to t=1
+    x_0 ~ N(0, I) → x_1 = generated image
+    """
+    B = shape[0]
+    x = torch.randn(shape, device=device)
+    dt = 1.0 / num_steps
+    
+    for i in range(num_steps):
+        t = torch.full((B,), i * dt, device=device)
+        v = model(x, t, class_label)
+        if cfg_scale > 0 and class_label is not None:
+            v_uncond = model(x, t, None)
+            v = v_uncond + cfg_scale * (v - v_uncond)
+        x = x + v * dt
+    
+    return x
+
+
+@torch.no_grad()
+def heun_sample(model, shape, num_steps=25, device='cpu', class_label=None, cfg_scale=0.0):
+    """
+    Generate images using Heun's method (2nd order) ODE integration.
+    More accurate than Euler. Each step costs 2 model evaluations.
+    """
+    B = shape[0]
+    x = torch.randn(shape, device=device)
+    dt = 1.0 / num_steps
+    
+    def get_v(x_in, t_in):
+        v = model(x_in, t_in, class_label)
+        if cfg_scale > 0 and class_label is not None:
+            v_uncond = model(x_in, t_in, None)
+            v = v_uncond + cfg_scale * (v - v_uncond)
+        return v
+    
+    for i in range(num_steps):
+        t = torch.full((B,), i * dt, device=device)
+        t_next = torch.full((B,), min((i + 1) * dt, 1.0), device=device)
+        k1 = get_v(x, t)
+        x_hat = x + dt * k1
+        if i < num_steps - 1:
+            k2 = get_v(x_hat, t_next)
+            x = x + dt * 0.5 * (k1 + k2)
+        else:
+            x = x + dt * k1
+    
+    return x
+
+
+@torch.no_grad()
+def generate_grid(model, num_images=16, num_steps=50, img_size=128,
+                  device='cpu', class_label=None, cfg_scale=0.0, method='euler'):
+    """Generate a grid of images. Returns (B, C, H, W) tensor in [0, 1]."""
+    shape = (num_images, 3, img_size, img_size)
+    if method == 'euler':
+        images = euler_sample(model, shape, num_steps, device, class_label, cfg_scale)
+    elif method == 'heun':
+        images = heun_sample(model, shape, num_steps, device, class_label, cfg_scale)
+    else:
+        raise ValueError(f"Unknown method: {method}")
+    return images.clamp(-1, 1) * 0.5 + 0.5
+
+
+def make_grid_image(images, nrow=4, padding=2):
+    """Arrange images into a grid. Returns a PIL Image."""
+    from torchvision.utils import make_grid
+    from PIL import Image
+    import numpy as np
+    grid = make_grid(images, nrow=nrow, padding=padding, normalize=False)
+    grid = grid.permute(1, 2, 0).cpu().numpy()
+    grid = (grid * 255).clip(0, 255).astype(np.uint8)
+    return Image.fromarray(grid)