ddevin2002/tiny-self-forcing-doom
0
1"""Samplers.2 3- sample_clip(): bidirectional Euler sampling of a whole clip (teacher mode).4- CausalSampler: KV-cached chunk-wise autoregressive sampler — interface stubbed now,5 implemented with CausalSFT (see ../CLAUDE.md §4).6 7CLI: uv run python sample/rollout.py --ckpt ckpt_smoke/last.pt --out samples/8"""9 10from __future__ import annotations11 12import argparse13from pathlib import Path14 15import numpy as np16import torch17from PIL import Image18 19import sys20sys.path.insert(0, str(Path(__file__).resolve().parent.parent))21 22from model.dit import DiTConfig, VideoDiT23from model.masks import chunk_causal24 25 26@torch.no_grad()27def sample_clip(model: VideoDiT, actions: torch.Tensor, steps: int = 50,28 seed: int = 0) -> torch.Tensor:29 """actions: [B, T] long -> frames [B, T, 3, H, W] in [-1, 1]. Euler on v-pred flow."""30 cfg = model.cfg31 dev = next(model.parameters()).device32 b, tf = actions.shape33 g = torch.Generator(dev.type).manual_seed(seed)34 x = torch.randn(b, tf, cfg.channels, cfg.image_size, cfg.image_size,35 generator=g, device=dev)36 ts = torch.linspace(1.0, 0.0, steps + 1, device=dev)37 for i in range(steps):38 t = ts[i].expand(b, tf)39 v = model(x, t, actions)40 x = x - (ts[i] - ts[i + 1]) * v41 return x.clamp(-1, 1)42 43 44class CausalSampler:45 """Chunk-wise autoregressive sampler over a sliding window (../CLAUDE.md §4).46 47 Each chunk is denoised while attending to the last `context_chunks` chunks of48 already-decided frames. Those frames are clean and fixed for the duration of the49 chunk, so their keys/values are cached once and reused across all denoising50 steps -- the saving that makes real-time playback possible, since each step then51 costs one chunk of attention instead of a whole window.52 53 The cache is rebuilt per chunk rather than carried forward: pos_frame is a54 learned absolute embedding, so when the window slides the context occupies55 different positions and its old keys no longer describe it (pinned by56 tests/test_core.py::test_kv_cache_context_is_not_reused_across_positions).57 Rebuilding costs one forward per chunk against `steps_per_chunk` saved.58 """59 60 def __init__(self, model: VideoDiT, chunk: int = 4, steps_per_chunk: int = 50,61 context_chunks: int = 3, t_max: float = 1.0,62 noise_device: str = "cpu"):63 cfg = model.cfg64 if cfg.frames % chunk:65 raise ValueError(f"frames {cfg.frames} not divisible by chunk {chunk}")66 if (context_chunks + 1) * chunk > cfg.frames:67 raise ValueError(68 f"{context_chunks} context chunks + 1 current chunk of {chunk} exceeds "69 f"the model's {cfg.frames}-frame window")70 self.model = model71 self.chunk = chunk72 self.steps = steps_per_chunk73 self.context_chunks = context_chunks74 self.t_max = t_max75 # Draw noise on the CPU by default: torch.Generator("mps") and76 # torch.Generator("cpu") give different draws for the same seed, so a figure77 # regenerated on another machine would not reproduce.78 self.noise_device = noise_device79 self._decided: torch.Tensor | None = None80 self._actions: list[int] | None = None81 82 @property83 def ctx_len(self) -> int:84 return self.context_chunks * self.chunk85 86 def _t_grid(self, device) -> torch.Tensor:87 return torch.linspace(self.t_max, 0.0, self.steps + 1, device=device)88 89 def _noise(self, b: int, n: int, seed: int | None, g: torch.Generator | None,90 device) -> torch.Tensor:91 cfg = self.model.cfg92 shape = (b, n, cfg.channels, cfg.image_size, cfg.image_size)93 x = torch.randn(shape, generator=g, device=self.noise_device)94 return x.to(device)95 96 def _context_cache(self, ctx: torch.Tensor, ctx_actions: torch.Tensor):97 """Keys/values for already-decided frames, built one chunk at a time.98 99 Equivalent to one chunk-causal masked forward over the context, but never100 materialises the [T*S, T*S] mask and keeps every call on SDPA's fused path.101 """102 if ctx.shape[1] == 0:103 return None104 model, b = self.model, ctx.shape[0]105 cache = None106 for s in range(0, ctx.shape[1], self.chunk):107 e = min(s + self.chunk, ctx.shape[1])108 nxt: list = []109 model.forward_with_cache(110 ctx[:, s:e],111 torch.zeros(b, e - s, device=ctx.device), # context is clean: t = 0112 ctx_actions[:, s:e],113 cache, frame_start=s, collect=nxt,114 )115 cache = nxt116 return cache117 118 # ---- streaming primitives -------------------------------------------------119 120 def reset(self, batch: int = 1, init_frames: torch.Tensor | None = None,121 device=None) -> None:122 """Begin a rollout. `init_frames` seeds it with real frames (a whole number123 of chunks); without them the first chunk is generated cold."""124 cfg = self.model.cfg125 dev = device or next(self.model.parameters()).device126 if init_frames is None:127 self._decided = torch.empty(batch, 0, cfg.channels, cfg.image_size,128 cfg.image_size, device=dev)129 else:130 if init_frames.shape[1] % self.chunk:131 raise ValueError("init_frames must be a whole number of chunks")132 self._decided = init_frames.to(dev)133 self._actions = []134 135 def push(self, frames: torch.Tensor, actions: torch.Tensor) -> None:136 """Accept frames as decided history, trimming to the context window."""137 if self._decided is None:138 raise RuntimeError("call reset() first")139 self._decided = torch.cat([self._decided, frames.to(self._decided.device)], dim=1)140 keep = self.ctx_len or None141 if keep and self._decided.shape[1] > keep:142 self._decided = self._decided[:, -keep:]143 self._actions.extend(int(a) for a in actions.reshape(-1)[: frames.shape[1]])144 145 @torch.no_grad()146 def next_chunk(self, actions_chunk: torch.Tensor, ctx_actions: torch.Tensor,147 noise: torch.Tensor | None = None, steps: int | None = None,148 record_at: list[int] | None = None):149 """Denoise one chunk against the current history. Does NOT push the result.150 151 `record_at` additionally returns the ODE state at those step indices, which is152 what trajectory distillation matches against.153 """154 model = self.model155 b, n = actions_chunk.shape156 dev = self._decided.device157 ctx = self._decided[:, -self.ctx_len:] if self.ctx_len else self._decided[:, :0]158 cache = self._context_cache(ctx, ctx_actions)159 steps = steps or self.steps160 ts = self._t_grid(dev) if steps == self.steps else torch.linspace(self.t_max, 0.0, steps + 1, device=dev)161 162 x = noise.to(dev) if noise is not None else self._noise(b, n, None, None, dev)163 recorded = []164 for i in range(steps):165 if record_at is not None and i in record_at:166 recorded.append(x.clone())167 v = model.forward_with_cache(x, ts[i].expand(b, n), actions_chunk, cache,168 frame_start=ctx.shape[1])169 x = x - (ts[i] - ts[i + 1]) * v170 x = x.clamp(-1, 1)171 return (x, recorded) if record_at is not None else x172 173 # ---- the four consumers, each a thin wrapper -------------------------------174 175 @torch.no_grad()176 def _drive(self, actions: torch.Tensor, init_frames: torch.Tensor | None,177 seed: int, gt_frames: torch.Tensor | None = None,178 max_chunks: int | None = None) -> torch.Tensor:179 """Shared loop. `gt_frames` pushes ground truth instead of the model's own180 output (teacher forcing); `max_chunks` stops early."""181 b, total = actions.shape182 if total % self.chunk:183 raise ValueError(f"{total} frames is not a whole number of {self.chunk}-frame chunks")184 dev = next(self.model.parameters()).device185 g = torch.Generator(self.noise_device).manual_seed(seed)186 187 self.reset(batch=b, init_frames=init_frames, device=dev)188 out = [self._decided.clone()] if self._decided.shape[1] else []189 produced = self._decided.shape[1]190 n_chunks = 0191 while produced < total and (max_chunks is None or n_chunks < max_chunks):192 n = min(self.chunk, total - produced)193 have = self._decided.shape[1]194 ctx_a = actions[:, produced - have:produced]195 x = self.next_chunk(actions[:, produced:produced + n], ctx_a,196 noise=self._noise(b, n, None, g, dev))197 keep = x if gt_frames is None else gt_frames[:, produced:produced + n].to(dev)198 out.append(x)199 self.push(keep, actions[:, produced:produced + n])200 produced += n201 n_chunks += 1202 return torch.cat(out, dim=1)[:, :produced]203 204 def rollout(self, actions: torch.Tensor, init_frames: torch.Tensor | None = None,205 seed: int = 0) -> torch.Tensor:206 """Free rollout. actions: [B, T] long -> frames [B, T, 3, H, W] in [-1, 1].207 208 `init_frames` [B, k, 3, H, W] seeds it with real frames (the replay protocol209 in ../CLAUDE.md §6 starts the model from the frames the engine did); k must be210 a whole number of chunks. Without it the first chunk is generated cold.211 """212 return self._drive(actions, init_frames, seed)213 214 def rollout_teacher_forced(self, gt_frames: torch.Tensor, actions: torch.Tensor,215 init_frames: torch.Tensor | None = None,216 seed: int = 0) -> torch.Tensor:217 """Every chunk is predicted from ground-truth history rather than the model's218 own. Distillation arm (a), and the honest short-horizon fidelity measurement:219 errors cannot compound, so this isolates per-step quality from drift."""220 if init_frames is None:221 init_frames = gt_frames[:, : self.ctx_len]222 return self._drive(actions, init_frames, seed, gt_frames=gt_frames)223 224 def rollout_prefix(self, init_frames: torch.Tensor, actions: torch.Tensor,225 k_chunks: int, seed: int = 0) -> torch.Tensor:226 """Roll k chunks forward on the model's own output and stop -- the227 self-generated context that distillation arm (b) trains against."""228 want = init_frames.shape[1] + k_chunks * self.chunk229 return self._drive(actions[:, :want], init_frames, seed, max_chunks=k_chunks)230 231 232def to_uint8(x: torch.Tensor) -> np.ndarray:233 return ((x.cpu().float() + 1) * 127.5).round().clamp(0, 255).to(torch.uint8) \234 .permute(0, 1, 3, 4, 2).numpy()235 236 237def save_strip_and_gif(frames_u8: np.ndarray, out_dir: Path, name: str, upscale: int = 4):238 """frames_u8: [T, H, W, 3] -> horizontal strip PNG + animated GIF."""239 out_dir.mkdir(parents=True, exist_ok=True)240 t, h, w, _ = frames_u8.shape241 strip = frames_u8.transpose(1, 0, 2, 3).reshape(h, t * w, 3)242 Image.fromarray(strip).resize((t * w * 2, h * 2), Image.NEAREST) \243 .save(out_dir / f"{name}_strip.png")244 imgs = [Image.fromarray(f).resize((w * upscale, h * upscale), Image.NEAREST)245 for f in frames_u8]246 imgs[0].save(out_dir / f"{name}.gif", save_all=True, append_images=imgs[1:],247 duration=115, loop=0) # ~8.75 fps, matching ACTION_REPEAT=4 at 35 tics/s248 249 250def main() -> None:251 p = argparse.ArgumentParser()252 p.add_argument("--ckpt", required=True)253 p.add_argument("--out", type=Path, default=Path("samples"))254 p.add_argument("--steps", type=int, default=50)255 p.add_argument("--actions", type=str, default=None,256 help="comma-separated action ids, e.g. 1,1,1,4,4 (default: forward)")257 p.add_argument("--ema", action=argparse.BooleanOptionalAction, default=True)258 args = p.parse_args()259 260 state = torch.load(args.ckpt, map_location="cpu")261 cfg = DiTConfig(**state["model_cfg"])262 model = VideoDiT(cfg)263 model.load_state_dict(state["ema" if args.ema else "model"])264 dev = ("cuda" if torch.cuda.is_available()265 else "mps" if torch.backends.mps.is_available() else "cpu")266 model = model.to(dev).eval()267 268 if args.actions:269 ids = [int(s) for s in args.actions.split(",")]270 ids = (ids * cfg.frames)[: cfg.frames]271 else:272 ids = [1] * cfg.frames273 actions = torch.tensor([ids], device=dev)274 clip = sample_clip(model, actions, steps=args.steps)275 save_strip_and_gif(to_uint8(clip)[0], args.out, f"clip_s{args.steps}")276 print(f"wrote {args.out}/clip_s{args.steps}_strip.png and .gif "277 f"(step {state['step']}, actions {ids[:8]}…)")278 279 280if __name__ == "__main__":281 main()282