Spaces:
Running on Zero
Running on Zero
| """Euler flow-matching sampler with interval classifier-free guidance.""" | |
| from __future__ import annotations | |
| from typing import Optional | |
| import torch | |
| 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() | |