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): # end_dev is a graph input of the compiled stream (one prefix length). 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