| from __future__ import annotations |
|
|
| import torch |
| import torch.nn as nn |
| from gta_streaming_kv_cache import SummaryKVCache |
| from summaryformer import Summaryformer |
|
|
|
|
| def _unpack_kv_conditional_dict(conditional_dict: dict) -> dict: |
| """Model kwargs from ``conditional_dict``. |
| |
| All keys are optional: ``action_context``, ``fps``, ``z0_focus_region``. |
| ``z0_focus_region`` (Bool [B,hs,ws], True = subject) ADDS the focus embedding to those z0 tokens; |
| missing / None = no-op. Wrong type/shape is a RuntimeError at the model. |
| """ |
| if not isinstance(conditional_dict, dict): |
| raise TypeError("conditional_dict must be a dict.") |
| out: dict = {} |
| action_context = conditional_dict.get("action_context", None) |
| if action_context is not None: |
| out["action_context"] = action_context |
| fps = conditional_dict.get("fps", None) |
| if fps is not None: |
| out["fps"] = fps |
| z0_focus_region = conditional_dict.get("z0_focus_region", None) |
| if z0_focus_region is not None: |
| if not isinstance(z0_focus_region, torch.Tensor): |
| raise TypeError( |
| f"conditional_dict['z0_focus_region'] must be a Tensor or omitted; " |
| f"got {type(z0_focus_region).__name__}" |
| ) |
| out["z0_focus_region"] = z0_focus_region |
| return out |
|
|
|
|
| class SummaryformerWrapper(nn.Module): |
| """Thin wrapper around ``Summaryformer``. |
| |
| One full-res stream, one trunk, KV-cached inference: ``forward_z0`` for the seed |
| frame, ``forward_body`` for a body frame. |
| |
| Layout contract at this boundary (the model is BCFHW everywhere): |
| z_rest BFCHW ``[B,F,C,H,W]`` — renderer I/O layout, checked, never guessed |
| z0 ``[B,C,H,W]`` or F=1 in either 5-D layout |
| Outputs are BFCHW. |
| """ |
|
|
| def __init__( |
| self, |
| dim=1280, |
| ffn_dim=5120, |
| num_heads=20, |
| num_layers: int = 28, |
| *, |
| summary_tokens: int, |
| raw_window_frames: int, |
| summary_window_frames: int, |
| ): |
| super().__init__() |
| self.model = Summaryformer( |
| dim=dim, |
| ffn_dim=ffn_dim, |
| num_heads=num_heads, |
| num_layers=int(num_layers), |
| summary_tokens=int(summary_tokens), |
| raw_window_frames=int(raw_window_frames), |
| summary_window_frames=int(summary_window_frames), |
| ) |
|
|
| def new_visual_kv_cache(self) -> SummaryKVCache: |
| """Unbuilt ``SummaryKVCache`` (one buffer) with this model's X / Y / W / depth / heads; call |
| ``.build(batch=, frame_tokens=Tf, dtype=, device=)`` per session.""" |
| return SummaryKVCache(self.model) |
|
|
| def _as_bcfhw_rest(self, x: torch.Tensor, *, name: str) -> torch.Tensor: |
| """BFCHW ``[B,F,C,H,W]`` → BCFHW. C must sit at dim 2; a BCFHW input is an error.""" |
| c = int(self.model.LATENT_CHANNELS) |
| if not (isinstance(x, torch.Tensor) and x.dim() == 5 and int(x.size(2)) == c): |
| raise TypeError( |
| f"{name} must be Tensor[B,F,C={c},H,W] (BFCHW); got " |
| f"shape={getattr(x, 'shape', None)}" |
| ) |
| return x.permute(0, 2, 1, 3, 4).contiguous() |
|
|
| def _as_bcfhw_by_c(self, x: torch.Tensor, *, name: str, frames: int | None) -> torch.Tensor: |
| """4-D ``[B,C,H,W]`` or 5-D with C at dim 1 (BCFHW) / dim 2 (BFCHW) → BCFHW. |
| |
| Both 5-D layouts are accepted because the model side is BCFHW while the |
| caller side is BFCHW; C decides, and when both dims equal C the |
| model-side layout (BCFHW) wins. ``frames`` pins F. |
| """ |
| c = int(self.model.LATENT_CHANNELS) |
| if not isinstance(x, torch.Tensor): |
| raise TypeError(f"{name} must be a Tensor; got {type(x)}") |
| if x.dim() == 4 and int(x.size(1)) == c: |
| out = x.unsqueeze(2) |
| elif x.dim() == 5 and int(x.size(1)) == c: |
| out = x |
| elif x.dim() == 5 and int(x.size(2)) == c: |
| out = x.permute(0, 2, 1, 3, 4) |
| else: |
| raise TypeError( |
| f"{name} expects [B,C={c},H,W], [B,C,F,H,W] or [B,F,C,H,W]; got shape={tuple(x.shape)}" |
| ) |
| if frames is not None and int(out.size(2)) != int(frames): |
| raise ValueError(f"{name} expects F={int(frames)}; got shape={tuple(x.shape)}") |
| return out.contiguous() |
|
|
| def _as_bcfhw_z0(self, x: torch.Tensor | None, *, name: str) -> torch.Tensor | None: |
| return None if x is None else self._as_bcfhw_by_c(x, name=name, frames=1) |
|
|
| def forward(self, *args, **kwargs): |
| raise TypeError("Use forward_z0 (seed frame) / forward_body (one body frame); both need a kv_cache.") |
|
|
| def forward_body( |
| self, |
| conditional_dict: dict, |
| x_rest: torch.Tensor, |
| t_rest: torch.Tensor, |
| *, |
| kv_cache, |
| absolute_start_frames: int | torch.Tensor, |
| ) -> torch.Tensor: |
| """KV-cached denoise of ONE body frame. The first pass of a frame donates (``begin_frame``); every pass |
| writes that RAW slot and reads ``[0 : end]``. |
| |
| - ``x_rest``: BFCHW ``[B,1,C,H,W]`` |
| - ``t_rest``: ``[B,1]`` |
| - ``absolute_start_frames``: this frame's absolute id (>= 1; 0 is z0); a 0-d int64 device |
| tensor is the static-compile streaming form |
| - Returns BFCHW ``[B,1,C,H,W]``. Call ``kv_cache.commit_current_frame()`` after the clean recache |
| (cursor advance, no copy). |
| """ |
| if not isinstance(kv_cache, SummaryKVCache): |
| raise RuntimeError( |
| "kv_cache must be a SummaryKVCache from wrapper.new_visual_kv_cache()." |
| ) |
| x_rest_bcfhw = self._as_bcfhw_rest(x_rest, name="x_rest") |
| pred_rest_bcfhw = self.model._forward_kv_cached_z_rest( |
| x_rest=x_rest_bcfhw, |
| t_rest=t_rest, |
| kv_cache=kv_cache, |
| absolute_start_frames=absolute_start_frames if isinstance(absolute_start_frames, torch.Tensor) else int(absolute_start_frames), |
| **_unpack_kv_conditional_dict(conditional_dict), |
| ) |
| return pred_rest_bcfhw.permute(0, 2, 1, 3, 4).contiguous() |
|
|
| def forward_z0( |
| self, |
| x_z0: torch.Tensor, |
| conditional_dict: dict, |
| t_z0: torch.Tensor, |
| kv_cache, |
| ) -> torch.Tensor: |
| """ |
| KV-cached **z0** denoise (seed / image stream). Writes record(0) into the Z0 slot on every pass and |
| reads ``[0 : REC_z0]``; call ``kv_cache.note_z0_written()`` after the clean recache. |
| |
| Image path still carries ``sum(0)`` (within-frame only). |
| |
| - ``x_z0``: ``[B,C,H,W]`` / ``[B,1,C,H,W]`` / ``[B,C,1,H,W]`` |
| - ``t_z0``: ``[B]`` |
| - Returns ``[B,1,C,H,W]`` (BFCHW, F=1). |
| """ |
| if not (isinstance(t_z0, torch.Tensor) and t_z0.dim() == 1): |
| raise TypeError( |
| f"t_z0 must be Tensor[B]; got " |
| f"type={type(t_z0)} shape={getattr(t_z0, 'shape', None)}" |
| ) |
| x_bcfhw = self._as_bcfhw_z0(x_z0, name="x_z0") |
| if x_bcfhw is None: |
| raise TypeError("forward_z0 requires x_z0.") |
| if int(t_z0.size(0)) != int(x_bcfhw.size(0)): |
| raise RuntimeError( |
| f"t_z0 batch mismatch: got B={int(t_z0.size(0))}, " |
| f"expected {int(x_bcfhw.size(0))}" |
| ) |
| if not isinstance(kv_cache, SummaryKVCache): |
| raise RuntimeError( |
| "forward_z0 requires a SummaryKVCache from wrapper.new_visual_kv_cache()." |
| ) |
|
|
| model_kwargs = _unpack_kv_conditional_dict(conditional_dict) |
| pred_bcfhw = self.model._forward_kv_cached_z0( |
| x_z0=x_bcfhw, |
| t_z0=t_z0, |
| kv_cache=kv_cache, |
| **model_kwargs, |
| ) |
| return pred_bcfhw.permute(0, 2, 1, 3, 4).contiguous() |
|
|
| @staticmethod |
| def mark_streaming_caches_static(*caches) -> int: |
| """Pin leaf addresses of the visual KV cache for cudagraphs. |
| |
| Call once after caches are finalized, before the first compiled forward. |
| Accepts a visual KV cache or ``None``. |
| """ |
| marked = 0 |
| for cache in caches: |
| if cache is None: |
| continue |
| if isinstance(cache, SummaryKVCache): |
| |
| for scalar in cache.device_scalars(): |
| torch._dynamo.mark_static_address(scalar) |
| marked += 1 |
| layers = getattr(cache, "layers", cache) |
| try: |
| layer_iter = list(layers) |
| except TypeError: |
| continue |
| for layer in layer_iter: |
| if not isinstance(layer, dict): |
| continue |
| for t in layer.values(): |
| if isinstance(t, torch.Tensor): |
| torch._dynamo.mark_static_address(t) |
| marked += 1 |
| return marked |
|
|
