| """Small pieces shared by the streaming session and the server. |
| |
| StageSampling / resolve_sampling which timesteps a body frame is denoised at |
| wrapper_ctor_kwargs_from_ckpt_config checkpoint config -> model constructor arguments |
| cond_dict, frame_tokens conditioning dict, tokens per frame |
| warm_z0_clean the clean write of the seed frame into memory |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
|
|
| import torch |
|
|
| from gta_streaming_kv_cache import SummaryKVCache |
| from noise_schedule import SampleGrid, sample_grid |
|
|
|
|
| |
| |
| |
| @dataclass(frozen=True) |
| class StageSampling: |
| """Where a body frame is denoised: a grid of timesteps and the order of model calls on it. |
| |
| Between two calls the current estimate of the clean frame is re-noised with FRESH noise to the next |
| timestep (``summary_streaming_inference.renoise_body_x0``). |
| """ |
|
|
| grid: SampleGrid |
| |
| |
| step_plan: tuple[int, ...] | None = None |
|
|
| def __post_init__(self): |
| n = len(self.grid) |
| if n < 2: |
| raise ValueError(f"sampling grid too short: {self.grid.timesteps}") |
| if self.step_plan is not None: |
| plan = tuple(int(i) for i in self.step_plan) |
| if not plan or any(i < 1 or i >= n for i in plan): |
| raise ValueError(f"step_plan must be non-empty grid indices in 1..{n - 1}; got {plan}") |
| if any(plan[k] < plan[k + 1] for k in range(len(plan) - 1)): |
| raise ValueError(f"step_plan must be non-increasing (high -> low t); got {plan}") |
|
|
| @property |
| def indices(self) -> list[int]: |
| """Grid indices of the model calls, in order.""" |
| if self.step_plan is not None: |
| return [int(i) for i in self.step_plan] |
| return list(range(len(self.grid) - 1, 0, -1)) |
|
|
| @property |
| def timesteps(self) -> list[int]: |
| """The timesteps (0..999) of those calls, e.g. ``[999, 878, 669]``.""" |
| return [int(self.grid.timesteps[i]) for i in self.indices] |
|
|
| @property |
| def label(self) -> str: |
| ts = self.timesteps |
| return f"{len(ts)} [{','.join(str(t) for t in ts)}]" if len(ts) <= 6 else f"{len(ts)}" |
|
|
|
|
| def resolve_sampling(config: dict, *, steps: int | None = None, tag: str = "sampling") -> StageSampling: |
| """The sampling plan a checkpoint config asks for. |
| |
| ``steps`` (default: the config's own ``len(timesteps)``) is the number of model calls per frame. The grid is |
| rebuilt from the config's schedule (``shifted_cosine_ratio``) and ``timestep_shift``; the caller checks it against |
| the config's ``timesteps`` list. |
| """ |
| if not isinstance(config, dict): |
| raise TypeError(f"ckpt config must be a dict; got {type(config)}") |
| if str(config["noise_schedule"]) != "shifted_cosine": |
| raise ValueError(f"noise_schedule must be 'shifted_cosine'; got {config['noise_schedule']!r}") |
| n = len(config["timesteps"]) if steps is None else int(steps) |
| grid = sample_grid(steps=n, shift=float(config["timestep_shift"]), ratio=float(config["shifted_cosine_ratio"])) |
| samp = StageSampling(grid=grid) |
| print(f"[{tag}] steps={n} shift={float(config['timestep_shift']):g} t=[{','.join(str(t) for t in samp.timesteps)}]") |
| return samp |
|
|
|
|
| |
| |
| |
| _ARCH_KEYS = ("dim", "ffn_dim", "num_heads", "num_layers", "summary_tokens", "raw_window_frames", "summary_window_frames") |
|
|
|
|
| def num_layers_from_ckpt_config(cfg: dict) -> int: |
| """Number of trunk blocks.""" |
| n = int(cfg["num_layers"]) |
| if n < 1: |
| raise ValueError(f"num_layers must be >= 1; got {n}") |
| return n |
|
|
|
|
| def wrapper_ctor_kwargs_from_ckpt_config(cfg: dict) -> dict: |
| """``SummaryformerWrapper`` kwargs from a checkpoint config. |
| |
| Memory geometry is X / Y / W: ``summary_tokens`` / ``raw_window_frames`` / ``summary_window_frames``. |
| """ |
| if not isinstance(cfg, dict): |
| raise TypeError(f"ckpt config must be a dict; got {type(cfg)}") |
| missing = [k for k in _ARCH_KEYS if k not in cfg] |
| if missing: |
| raise KeyError(f"ckpt config missing {missing}.") |
| return dict( |
| dim=int(cfg["dim"]), |
| ffn_dim=int(cfg["ffn_dim"]), |
| num_heads=int(cfg["num_heads"]), |
| num_layers=num_layers_from_ckpt_config(cfg), |
| summary_tokens=int(cfg["summary_tokens"]), |
| raw_window_frames=int(cfg["raw_window_frames"]), |
| summary_window_frames=int(cfg["summary_window_frames"]), |
| ) |
|
|
|
|
| |
| |
| |
| def frame_tokens(model, h: int, w: int) -> int: |
| """Tokens of one frame on an ``h x w`` latent grid: one token per latent cell.""" |
| return int(h) * int(w) |
|
|
|
|
| def cond_dict(*, fps=None, action_context=None, z0_focus_region=None) -> dict: |
| """``z0_focus_region`` is a seed-frame knob (Bool [B,H,W]).""" |
| out = { |
| "fps": fps, |
| "action_context": action_context, |
| } |
| if z0_focus_region is not None: |
| out["z0_focus_region"] = z0_focus_region |
| return out |
|
|
|
|
| @torch.no_grad() |
| def warm_z0_clean(wrapper, kv: SummaryKVCache, cond_pos: dict, xhat0: torch.Tensor) -> None: |
| """Write the clean seed frame's record into the Z0 slot (one pass at t = 0), then ``kv.note_z0_written()``. |
| |
| ``xhat0 [B, 1, C, H, W]``. No ``begin_frame`` and no cursor commit: the seed frame is the sink, not a RAW slot. |
| """ |
| if not (isinstance(xhat0, torch.Tensor) and xhat0.dim() == 5 and int(xhat0.size(1)) == 1): |
| raise TypeError(f"xhat0 must be Tensor[B,1,C,H,W]; got {getattr(xhat0, 'shape', type(xhat0))}") |
| t0 = torch.zeros((int(xhat0.size(0)),), device=xhat0.device, dtype=torch.float32) |
| wrapper.forward_z0(x_z0=xhat0, conditional_dict=cond_pos, t_z0=t0, kv_cache=kv) |
| kv.note_z0_written() |
|
|
