rti/sampling.py

"""Euler flow-matching sampler with interval classifier-free guidance."""
from __future__ import annotations

from typing import Optional

import torch


@torch.no_grad()
def sample(model, ctx, mask, *, budget: int, steps: int = 100, cfg_scale: float = 5.0,
           image_size: int = 512, noise_scale: float = 2.0, t_eps: float = 0.001,
           cfg_interval: tuple[float, float] = (0.0, 1.0), generator=None,
           device: Optional[str] = None, dtype=torch.bfloat16, progress: bool = False):
    """model: an RTIMiniT2I. ctx, mask: encoded prompts [B,L,D] and [B,L].
    budget: region-token budget R; R >= N runs dense.
    Returns images in [0,1], [B,3,image_size,image_size], on CPU."""
    device = device or next(model.parameters()).device
    B = ctx.shape[0]
    ctx = ctx.to(device, dtype)
    mask = mask.to(device, dtype)
    null = torch.zeros_like(mask)
    x = torch.randn(B, 3, image_size, image_size, generator=generator,
                    device=device, dtype=dtype) * noise_scale
    ts = torch.linspace(0.0, 1.0, steps + 1, device=device, dtype=dtype)
    lo, hi = cfg_interval
    it = range(steps)
    if progress:
        from tqdm.auto import tqdm
        it = tqdm(it)
    with torch.autocast(device.type if hasattr(device, "type") else "cuda", dtype=dtype):
        for i in it:
            t0 = float(ts[i])
            tt = torch.full((2 * B,), t0, device=device, dtype=x.dtype)
            x2 = torch.cat([x, x])
            x0 = model.forward_xpred(x2, tt, torch.cat([ctx, ctx]),
                                     torch.cat([mask, null]), budget=budget)
            v = (x0 - x2) / max(1.0 - t0, t_eps)
            v_cond, v_uncond = v.chunk(2, 0)
            s = cfg_scale if lo <= t0 <= hi else 1.0
            v = v_uncond + (v_cond - v_uncond) * s
            x = x + (ts[i + 1] - ts[i]) * v
    return ((x.float().clamp(-1, 1) + 1) / 2).cpu()