File size: 8,987 Bytes
ae4e95d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 | 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
|