"""Private SEED path runtime: complete GPU sampling graphs and mixed precision. This preserves the official finite motion window. Motion K/V is recomputed: removing an old prefix changes deeper-layer history representations. Text K/V is safely cached because its projections depend only on the current genuine text features and fixed weights, with cross-text RoPE disabled in this release. No public/vendor files are changed; close restores original methods/FP32 weights. """ from __future__ import annotations from contextlib import contextmanager import copy import math import time import torch from fixed_text_runtime import FixedPromptBank, _feature_key from space.inference import FPS from window_runtime import WindowRuntime DEFAULT_TEXT_BUCKETS = (32, 64, 128, 256, 512, 1024) def _motion_attention_mask(indices, valid, target_start, history_cutoff): """Keep the partial mask, isolating new queries from an old prompt epoch. All arguments can stay on the GPU with fixed shapes during graph replay. Old queries keep their normal causal keys, so the cutoff introduces no fully masked valid rows. New queries cannot read them in any model layer. A zero cutoff is exactly the original mask. """ partial = (valid[:, None] & valid[None, :] & ((indices[:, None] >= target_start) | (indices[None, :] <= indices[:, None]))) epoch = ((indices[:, None] < history_cutoff) | (indices[None, :] >= history_cutoff)) return partial & epoch class _MixedPrecisionWeights: def __init__(self, model, precision): if precision not in ("bf16", "fp32"): raise ValueError("precision must be 'bf16' or 'fp32'") self.precision = precision self.originals = [] if precision == "fp32": return backbone = model.model modules = [backbone.patch_embedding, backbone.text_embedding] for block in backbone.blocks: modules.append(block.ffn) for attention in (block.self_attn, block.cross_attn): modules.extend((attention.q, attention.k, attention.v, attention.o)) seen = set() try: for module in modules: for parameter in module.parameters(): if id(parameter) in seen: continue seen.add(id(parameter)) original = parameter.data self.originals.append((parameter, original)) parameter.data = original.to(dtype=torch.bfloat16) except BaseException: # Construction can fail before the caller receives this handle. # Restore already converted storages here, including on OOM. self.restore() raise def restore(self): """Restore the actual original storages, without BF16 round-trip loss.""" for parameter, original in self.originals: parameter.data = original self.originals.clear() def configure_mixed_precision(model, precision="bf16"): """Match the old benchmark's BF16 projections; return a .restore() handle. Time embeddings, normalization, modulation, output head, latent state and integration remain FP32. Reference forwards must also use CUDA BF16 autocast with cache_enabled=False when precision is bf16. """ return _MixedPrecisionWeights(model, precision) class _SamplingCore: """One text-capacity graph sharing the runtime's actual ring and step id.""" def __init__(self, runtime, capacity, null_length): self.runtime, self.model = runtime, runtime.model self.capacity, self.null_length = int(capacity), int(null_length) self.history = int(runtime._history) self.precision = runtime._fast_precision self.device = self.model.generated[0].device self.state = self.model.generated[0] self.roots = self.model._stream_position self.step_id = runtime._fast_step_id self.history_cutoff = runtime._fast_history_cutoff self.time_table = runtime._fast_time_table self.rows = torch.arange(self.history, device=self.device, dtype=torch.long) self.text_length = max(self.capacity, self.null_length) width = self.model.text_module.dim dim = self.model.model.dim dtype = torch.bfloat16 if self.precision == "bf16" else torch.float32 self.context = torch.zeros(self.capacity, width, device=self.device) self.null_context = torch.zeros(self.null_length, width, device=self.device) self.cross_mask = torch.zeros(self.history, self.capacity, dtype=torch.bool, device=self.device) self.projected_text = torch.zeros(2, self.text_length, dim, device=self.device, dtype=dtype) # WanRMSNorm restores FP32 through its FP32 learned scale. self.text_k = [torch.zeros(2, self.text_length, dim, device=self.device) for _ in self.model.model.blocks] self.text_v = [torch.zeros(2, self.text_length, dim, device=self.device, dtype=dtype) for _ in self.model.model.blocks] self.output = torch.zeros(self.model.model_input_dim, device=self.device) self.graph = None self.text_revision = None self.text_refreshes = 0 def _autocast(self): return torch.autocast("cuda", enabled=self.precision == "bf16", dtype=torch.bfloat16, cache_enabled=False) @torch.inference_mode() def prepare_text(self, context, null_context, mask, revision=None): """Update only mutable inputs; no capture or model-state reset.""" self.cross_mask.zero_() self.cross_mask[:mask.shape[0]].copy_(mask) if revision is not None and revision == self.text_revision: return self.context.copy_(context) self.null_context.copy_(null_context) padded = self.context.new_zeros((2, self.text_length, self.context.shape[1])) padded[0, :self.capacity].copy_(self.context) padded[1, :self.null_length].copy_(self.null_context) # This work happens only for new/changed prompt features or a moved # bank span. It is independent of pose, timesteps and motion history. with self._autocast(): projected = self.model.model.text_embedding(padded) self.projected_text.copy_(projected) for block, keys, values in zip(self.model.model.blocks, self.text_k, self.text_v): attention = block.cross_attn keys.copy_(attention.norm_k(attention.k(projected))) values.copy_(attention.v(projected)) self.text_revision = revision self.text_refreshes += 1 @contextmanager def _cached_text_projections(self): """Bind cached tensors during eager/capture; replay needs no patching.""" saved = [] def replace(module, forward): saved.append((module, module.forward)) module.forward = forward replace(self.model.model.text_embedding, lambda _value: self.projected_text) for block, keys, values in zip(self.model.model.blocks, self.text_k, self.text_v): replace(block.cross_attn.k, lambda _value, cached=keys: cached) replace(block.cross_attn.norm_k, lambda value: value) replace(block.cross_attn.v, lambda _value, cached=values: cached) try: yield finally: for module, forward in reversed(saved): module.forward = forward @torch.inference_mode() def step(self): """GPU schedule, full-window denoising, CFG, Euler, output and tick.""" # CPU-generated table matches torch.tensor(current_step / 30) exactly. end = self.step_id + 1 window_start = (end - self.history).clamp(min=0) indices = window_start + self.rows valid = indices < end target_start = (self.step_id - 29).clamp(min=0) updating = valid & (indices >= target_start) latent = self.state.index_select(1, indices) raw_root = self.roots[0].index_select(0, indices) normalized_root = (raw_root - self.model.position_mean) / self.model.position_std conditioned = torch.cat((normalized_root.T[:, :, None, None], latent), dim=0) conditioned = torch.where(valid[None, :, None, None], conditioned, torch.zeros_like(conditioned)) level = (-indices.float() / 30 + self.time_table.index_select(0, self.step_id.reshape(1))[0]).clamp(0, 1) level = level * self.model.time_embedding_scale self_mask = _motion_attention_mask(indices, valid, target_start, self.history_cutoff) null_mask = valid[:, None].expand(self.history, self.null_length) with self._autocast(), self._cached_text_projections(): predicted = self.runtime._forward( [conditioned, conditioned], [level, level], [self.context, self.null_context], self.history, attn_mask=[self_mask, self_mask], text_k_rope_ids=None, text_q_rope_ids=None, cross_attn_mask=[self.cross_mask, null_mask], text_seq_len=self.text_length, y=None) # Keep scalar order equal to the official text_scale*pt + null_scale*pn. velocity = self.model.cfg_config["text_scale"] * predicted[0] + self.model.cfg_config["null_scale"] * predicted[1] updated = torch.where(updating[None, :, None, None], latent + velocity * (1 / 30), latent) self.state.index_copy_(1, indices, updated) committed_index = (end - 30).clamp(min=0).reshape(1) pose = self.state.index_select(1, committed_index)[:, 0, 0, 0] root = self.roots[0].index_select(0, committed_index)[0] self.output.copy_(torch.cat((root, pose * self.model.std + self.model.mean))) self.step_id.add_(1) def snapshot(self): """All mutable GPU data required to compare step() and graph.replay().""" return dict(state=self.state.clone(), roots=self.roots.clone(), step_id=self.step_id.clone(), history_cutoff=self.history_cutoff.clone(), context=self.context.clone(), null_context=self.null_context.clone(), cross_mask=self.cross_mask.clone(), projected_text=self.projected_text.clone(), text_k=[value.clone() for value in self.text_k], text_v=[value.clone() for value in self.text_v], output=self.output.clone()) def restore(self, snapshot): for name in ("state", "roots", "step_id", "history_cutoff", "context", "null_context", "cross_mask", "projected_text", "output"): getattr(self, name).copy_(snapshot[name]) for name in ("text_k", "text_v"): for destination, source in zip(getattr(self, name), snapshot[name]): destination.copy_(source) @torch.inference_mode() def capture(self): """Prepare once without consuming RNG or retaining probe motion.""" snapshot = self.snapshot() stream = torch.cuda.Stream() stream.wait_stream(torch.cuda.current_stream()) try: with torch.cuda.stream(stream): for _ in range(3): self.restore(snapshot) self.step_id.fill_(self.history - 1) self.step() torch.cuda.current_stream().wait_stream(stream) torch.cuda.synchronize() self.restore(snapshot) self.step_id.fill_(self.history - 1) self.step() reference = {"state": self.state.clone(), "output": self.output.clone(), "step_id": self.step_id.clone()} self.restore(snapshot) self.step_id.fill_(self.history - 1) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph, stream=stream): self.step() self.restore(snapshot) self.step_id.fill_(self.history - 1) graph.replay() torch.testing.assert_close(self.state, reference["state"], rtol=1e-5, atol=1e-5) torch.testing.assert_close(self.output, reference["output"], rtol=1e-5, atol=1e-5) torch.testing.assert_close(self.step_id, reference["step_id"], rtol=0, atol=0) self.graph = graph finally: self.restore(snapshot) class FastWindowRuntime(WindowRuntime): """Compatible local runtime using full-step graphs and a BF16 backbone. Native path controls, pending prompt/path replacement, 29-frame cold-start lookahead, finite 90/120/150-frame attention windows and ring rollback are kept. use_graph=False executes exactly the same optimized step eagerly. precision='fp32' keeps all original weights for precision comparisons. reset_history_on_prompt_change masks committed motion from earlier prompt epochs while preserving pending latents, scheduling, noise and recovery. All step APIs observe prompt changes; only step_replanned(replan_text=True) also replaces the text conditions of already pending frames. """ def __init__(self, model, recovery_type, metadata): super().__init__(model, recovery_type, dict(metadata)) self._fast_original_stream = None self._precision_handle = None self._fast_cores = {} self._overflow_core = None self._last_fast_core = None self._fast_text = None self._fast_step_id = self._fast_time_table = None self._fast_history_cutoff = None self._history_reset_enabled = False self._history_cutoff_absolute = None self._history_cutoff_local = 0 self._history_resets = 0 self._visible_history = 0 self._fast_precision = "bf16" self._fast_use_graph = False self._fast_graph_replays = self._fast_eager_calls = 0 self._fast_prewarm = None def _assert_fast_contract(self): model = self.model scheduler = model.time_scheduler if (model.batch_size != 1 or model.input_dim <= 0 or model.model_input_dim != model.input_dim + model.stream_condition_dim or tuple(model.spatial_shape) != (1, 1) or tuple(model.model.patch_size) != (1, 1, 1) or model.stream_condition_dim != 3 or model.attn_type != "partial" or model.prediction_type != "vel" or scheduler.steps != 30 or scheduler.chunk_size != 30 or scheduler.noise_type != "linear" or scheduler.sigma_type != "zero" or model.text_module.cross_rope): raise RuntimeError("FastWindowRuntime requires pose channels plus root3, 30/30 linear/zero, partial attention") def _start_fast_session(self, seed, history): """Base session initialization with local finite-history windows. The public Runtime intentionally still accepts only its released window sizes. Keep its seed, scheduler and recovery initialization here, while the local full-step runtime also supports 60/90 total attention frames. No forward-only graph wrapper is needed before full-step prewarming. """ history = int(history) if history not in (60, 90, 120, 150): raise ValueError("history must be 60, 90, 120 or 150 total window frames") if next(self.model.parameters()).device.type != "cuda": raise RuntimeError("start must run inside the ZeroGPU CUDA allocation") self.close() torch.manual_seed(int(seed)) self.model.schedule_config.clear() self.model.schedule_config.update(copy.deepcopy(self._schedule)) self.model.cfg_config.update(text_scale=4.0, null_scale=-3.0) self.model.init_generated(history, batch_size=1, schedule_config=copy.deepcopy(self._schedule)) self._recovery = self._recovery_type(smoothing_alpha=1.0, fps=getattr(self, "fps", FPS)) self._steps = self._frames = self._recovered = 0 self._action = self._prompt = None self._history = history self._ready = True @torch.inference_mode() def start(self, seed=0, history=120, use_graph=True, *, precision="bf16", text_token_buckets=DEFAULT_TEXT_BUCKETS, reset_history_on_prompt_change=False, cfg_scale=4.0): if precision not in ("bf16", "fp32"): raise ValueError("precision must be 'bf16' or 'fp32'") if not isinstance(reset_history_on_prompt_change, bool): raise TypeError("reset_history_on_prompt_change must be a bool") cfg_scale = float(cfg_scale) if not math.isfinite(cfg_scale) or cfg_scale < 0: raise ValueError("cfg_scale must be a finite nonnegative number") buckets = tuple(text_token_buckets) if (not buckets or any(isinstance(n, bool) or not isinstance(n, int) or n <= 0 for n in buckets) or tuple(sorted(set(buckets))) != buckets): raise ValueError("Text token buckets must be increasing positive integers") began = time.perf_counter() try: self._start_fast_session(seed, history) self.model.cfg_config.update(text_scale=cfg_scale, null_scale=1.0-cfg_scale) self._assert_fast_contract() self.model.model.forward = self._forward self._fast_precision, self._fast_use_graph = precision, bool(use_graph) self._precision_handle = configure_mixed_precision(self.model, precision) device = self.model.generated[0].device self._fast_step_id = torch.zeros((), dtype=torch.long, device=device) # Every capacity graph reads this same mutable scalar. Prompt # switches update its value, never the graph shape or capture. self._fast_history_cutoff = torch.zeros((), dtype=torch.long, device=device) self._history_reset_enabled = reset_history_on_prompt_change self._history_cutoff_absolute = None self._history_cutoff_local = 0 self._history_resets = self._visible_history = 0 self._fast_time_table = torch.tensor([index * (1 / 30) for index in range(self.model.buf_len + 1)], device=device, dtype=torch.float32) text = self.model.text_module self._fast_text = FixedPromptBank(text, text.get_stream_context, buckets) null = text.text_cache[""].to(device=device, dtype=torch.float32) for capacity in buckets: core = _SamplingCore(self, capacity, len(null)) context = torch.zeros_like(core.context) copied = min(len(null), capacity) context[:copied].copy_(null[:copied]) mask = torch.zeros_like(core.cross_mask) mask[:, :copied] = True core.prepare_text(context, null, mask) if use_graph: core.capture() self._fast_cores[capacity] = core self._fast_original_stream = self.model.stream_generate_step self.model.stream_generate_step = self._stream_generate_step self._fast_graph_replays = self._fast_eager_calls = 0 self._fast_prewarm = dict(complete=True, seconds=time.perf_counter() - began, captures=len(buckets) if use_graph else 0, capacities=list(buckets)) self.metadata.update(cfg_scale=cfg_scale, precision=("BF16 projections; FP32 normalization/time/head/latent/Euler" if precision == "bf16" else "FP32 parameters and sampling; original BF16 SDPA"), sampling_graph="full GPU step", motion_kv=False, text_kv=True) return self except BaseException: self.close() raise def _update_prompt_history(self, prompt): """Run after input validation/condition commit, before the GPU step. Runtime counters remain absolute while the model counters roll back. The earliest uncommitted frame is self._frames, not self._steps. Keep that epoch boundary fixed as new committed history accumulates. """ if (self._history_reset_enabled and self._prompt is not None and prompt != self._prompt): self._history_cutoff_absolute = self._frames self._history_resets += 1 offset = self._steps - self.model.current_step cutoff = max(0, (self._history_cutoff_absolute or 0) - offset) if cutoff != self._history_cutoff_local: self._fast_history_cutoff.fill_(cutoff) self._history_cutoff_local = cutoff end = self.model.current_step + 1 target_start = max(0, end - self.window_frames) window_start = max(0, end - self._history) self._visible_history = max(0, target_start - max(window_start, cutoff)) @torch.inference_mode() def _stream_generate_step(self, inputs): model = self.model inputs = model._extract_inputs(inputs) device = model.generated[0].device model._update_stream_condition(inputs, device) model.text_module.update_stream(inputs, device, model.param_dtype) model.condition_frames += 1 if model.condition_frames > model.buf_len: model._rollback() self._fast_step_id.sub_(model.seq_len) model._commit_stream_condition(model.condition_frames - 1) # Public step methods have already validated prompt features and root # rows. In particular, a rejected replan cannot advance the cutoff. self._update_prompt_history(inputs["text"][0]) step = model.current_step end = step + 1 start = max(0, end - self._history) contexts, meta = self._fast_text(start, end, device, model.param_dtype) null_feature = model.text_module.text_cache[""] null = null_feature.to(device=device, dtype=model.param_dtype) context, mask = contexts[0], meta["cross_attn_mask"][0] compatible = self._fast_text.graph_compatible and len(null) == self._fast_cores[self._fast_text.capacity].null_length if compatible: core = self._fast_cores[self._fast_text.capacity] null_key = _feature_key(null_feature) revision = (id(self._fast_text.bank), self._fast_text.uploads, null_key) if null_key[1] is None: revision = None else: contract = (len(context), len(null)) if self._overflow_core is None or (self._overflow_core.capacity, self._overflow_core.null_length) != contract: self._overflow_core = _SamplingCore(self, *contract) core, revision = self._overflow_core, None core.prepare_text(context, null, mask, revision) self._last_fast_core = core if self._fast_use_graph and compatible: core.graph.replay() self._fast_graph_replays += 1 else: core.step() self._fast_eager_calls += 1 # The official sigma=0 implementation still draws this exact shape. # Retain its RNG advancement so newly filled rollback noise matches. torch.randn_like(model.generated[0][:, max(0, step - 29):end]) model.current_step += 1 committable = max(0, model.condition_frames - 29) if committable > model.current_commit: model.current_commit = committable return {"generated": [core.output.unsqueeze(0)]} return {"generated": [core.output.new_empty((0, model.model_input_dim))]} def status(self): status = super().status() status.update(cfg_scale=float(self.model.cfg_config['text_scale']), history_reset_enabled=self._history_reset_enabled, history_cutoff_absolute=self._history_cutoff_absolute, history_resets=self._history_resets, visible_history=self._visible_history) status["graph"] = dict(enabled=self._fast_use_graph, cached_shapes=len(self._fast_cores) if self._fast_use_graph else 0, captures=0, replays=self._fast_graph_replays, eager_calls=self._fast_eager_calls) status["prewarm"] = dict(self._fast_prewarm) if self._fast_prewarm else None status["fixed_text"] = self._fast_text.status() if self._fast_text else None status["fast"] = dict(precision=self._fast_precision, complete_step_graph=True, motion_kv=False, text_kv=True, finite_window=self._history, text_refreshes=sum(core.text_refreshes for core in self._fast_cores.values())) return status def close(self): if self._fast_original_stream is not None: self.model.stream_generate_step = self._fast_original_stream self._fast_original_stream = None if self._fast_cores or self._overflow_core is not None: torch.cuda.synchronize() for core in self._fast_cores.values(): core.graph = None self._fast_cores.clear() self._overflow_core = None self._last_fast_core = None self._fast_text = None self._fast_step_id = self._fast_time_table = None self._fast_history_cutoff = None if self._precision_handle is not None: self._precision_handle.restore() self._precision_handle = None super().close()