# Copyright 2024 NVIDIA CORPORATION & AFFILIATES # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # # SPDX-License-Identifier: Apache-2.0 """Diffusers-backed LTX-2 refiner used by Sana-WM inference. The Sana-WM refiner checkpoint is a standard LTX-2 transformer plus text connectors. Diffusers already owns those modules, but its public transformer forward always runs the audio stream and does not expose the streaming sink/current video self-attention mask that this refiner was trained with. This wrapper keeps the custom surface narrow: load diffusers components, encode the prompt through Gemma + ``LTX2TextConnectors``, and run a video-only forward through the diffusers transformer blocks. The only local attention code is the streaming sink/current split, implemented with diffusers attention modules without materializing the full sequence-by-sequence mask. """ from __future__ import annotations import gc import hashlib import json import os import re import time import types from contextlib import nullcontext from pathlib import Path import torch import torch.nn.functional as F from torch import nn STAGE_2_DISTILLED_SIGMA_VALUES: tuple[float, ...] = (0.909375, 0.725, 0.421875, 0.0) class DiffusersLTX2Refiner(nn.Module): """Small Sana-WM adapter around diffusers LTX-2 modules.""" def __init__( self, refiner_root: str | Path, gemma_root: str | Path, *, dtype: torch.dtype, device: torch.device | str, text_max_sequence_length: int = 1024, ) -> None: super().__init__() self.refiner_root = Path(refiner_root) self.gemma_root = Path(gemma_root) self.dtype = dtype self.device = torch.device(device) self.text_max_sequence_length = int(text_max_sequence_length) self._te_nvfp4_requested = _env_flag("SANA_WM_REFINER_NVFP4") self._te_nvfp4_recipe = None self._te_nvfp4_converted = False self._self_qkv_fused = False self._attention_backend = os.environ.get("SANA_WM_REFINER_ATTN_BACKEND", "").strip() self._uniform_timestep_cache: dict[tuple[int, int, float, str], tuple[torch.Tensor, torch.Tensor]] = {} self.transformer, self.connectors = self._load_diffusers_components() def _load_diffusers_components(self) -> tuple[nn.Module, nn.Module]: from diffusers.models.transformers.transformer_ltx2 import LTX2VideoTransformer3DModel from diffusers.pipelines.ltx2 import LTX2TextConnectors cache_path = self._prepared_transformer_cache_path() if cache_path is not None and cache_path.is_file(): t0 = time.perf_counter() print(f"[refiner-cache] loading prepared transformer from {cache_path}", flush=True) try: transformer = torch.load(cache_path, map_location=self.device, weights_only=False).eval() self._te_nvfp4_converted = bool(self._te_nvfp4_requested) self._self_qkv_fused = _env_flag("SANA_WM_REFINER_FUSE_SELF_QKV") self._te_nvfp4_recipe = self._make_nvfp4_recipe() if self._te_nvfp4_converted else None print(f"[refiner-cache] loaded prepared transformer in {time.perf_counter() - t0:.1f}s", flush=True) except Exception as exc: print(f"[refiner-cache] failed to load {cache_path}: {exc}; rebuilding", flush=True) transformer = LTX2VideoTransformer3DModel.from_pretrained( self.refiner_root, subfolder="transformer", torch_dtype=self.dtype, ).eval() else: transformer = LTX2VideoTransformer3DModel.from_pretrained( self.refiner_root, subfolder="transformer", torch_dtype=self.dtype, ).eval() if not self._te_nvfp4_requested and os.environ.get("SANA_WM_REFINER_FP8_STORAGE", "").lower() in { "1", "true", "yes", "on", }: skip_patterns = None extra_skip_patterns = _env_tuple("SANA_WM_REFINER_FP8_SKIP_PATTERNS") if extra_skip_patterns: from diffusers.hooks.layerwise_casting import DEFAULT_SKIP_MODULES_PATTERN skip_patterns = tuple(dict.fromkeys((*DEFAULT_SKIP_MODULES_PATTERN, *extra_skip_patterns))) transformer.enable_layerwise_casting( storage_dtype=torch.float8_e4m3fn, compute_dtype=self.dtype, skip_modules_pattern=skip_patterns, ) connectors = LTX2TextConnectors.from_pretrained( self.refiner_root, subfolder="connectors", torch_dtype=self.dtype, ).eval() return transformer, connectors def _make_nvfp4_recipe(self): """Build the refiner quant recipe. Default NVFP4 (W4A4, Blackwell); ``SANA_WM_REFINER_QUANT=fp8block`` selects FP8 block scaling (W8A8), which also runs on Hopper.""" import transformer_engine.common.recipe as te_recipe quant = os.environ.get("SANA_WM_REFINER_QUANT", "nvfp4").strip().lower() if quant in {"fp8block", "fp8", "float8block"}: return te_recipe.Float8BlockScaling() return te_recipe.NVFP4BlockScaling( disable_rht=True, disable_stochastic_rounding=True, ) def _prepared_transformer_cache_path(self) -> Path | None: root = _prepared_module_cache_root() if root is None or not self._te_nvfp4_requested: return None payload = { "kind": "refiner_transformer_prepared_v2", "refiner_root": _path_fingerprint(self.refiner_root / "transformer"), "dtype": str(self.dtype), "torch": torch.__version__, "refiner_nvfp4": os.environ.get("SANA_WM_REFINER_NVFP4", ""), "refiner_quant": os.environ.get("SANA_WM_REFINER_QUANT", ""), "refiner_nvfp4_skip_patterns": os.environ.get("SANA_WM_REFINER_NVFP4_SKIP_PATTERNS", ""), "refiner_fuse_self_qkv": os.environ.get("SANA_WM_REFINER_FUSE_SELF_QKV", ""), "te_cpu_staging": os.environ.get("SANA_WM_TE_NVFP4_CPU_STAGING", ""), } try: import transformer_engine payload["transformer_engine"] = getattr(transformer_engine, "__version__", "unknown") except Exception: payload["transformer_engine"] = "unavailable" return root / "refiner" / f"{_prepared_module_cache_hash(payload)}.pt" def _save_prepared_transformer_cache(self) -> None: if os.environ.get("SANA_WM_PREPARED_MODULE_CACHE_SAVE", "1").strip().lower() in { "", "0", "false", "no", "off", }: return cache_path = self._prepared_transformer_cache_path() if cache_path is None or cache_path.is_file(): return cache_path.parent.mkdir(parents=True, exist_ok=True) tmp_path = cache_path.with_suffix(cache_path.suffix + f".tmp.{os.getpid()}") t0 = time.perf_counter() print(f"[refiner-cache] saving prepared transformer to {cache_path}", flush=True) restore = _strip_local_callables_for_pickle(self.transformer) if restore: print(f"[refiner-cache] stripped {len(restore)} init-only callables before save", flush=True) try: torch.save(self.transformer, tmp_path) os.replace(tmp_path, cache_path) except Exception as exc: print(f"[refiner-cache] failed to save {cache_path}: {exc}", flush=True) finally: _restore_stripped_pickle_values(restore) try: if tmp_path.exists(): tmp_path.unlink() except OSError: pass if cache_path.is_file(): print(f"[refiner-cache] saved prepared transformer in {time.perf_counter() - t0:.1f}s", flush=True) def prepare_transformer_nvfp4(self) -> None: """Lazily replace eligible refiner Linear layers with TE NVFP4 Linear modules.""" self._prepare_self_qkv_fusion() if not self._te_nvfp4_requested or self._te_nvfp4_converted: return recipe = self._make_nvfp4_recipe() converted, skipped = _replace_linear_with_te_nvfp4( self.transformer, recipe=recipe, params_dtype=self.dtype, skip_patterns=tuple( dict.fromkeys( ( "^proj_in$", "^proj_out$", "(^|\\.)audio_", "audio_to_video", "video_to_audio", "av_cross_attn", "caption_projection", "time_embed", *_env_tuple("SANA_WM_REFINER_NVFP4_SKIP_PATTERNS"), ) ) ), ) if converted <= 0: raise RuntimeError(f"SANA_WM_REFINER_NVFP4=1 converted no Linear layers; skipped={skipped}.") self._te_nvfp4_recipe = recipe self._te_nvfp4_converted = True _empty_cuda_cache() self._save_prepared_transformer_cache() def _prepare_self_qkv_fusion(self) -> None: if self._self_qkv_fused or not _env_flag("SANA_WM_REFINER_FUSE_SELF_QKV"): return converted = _fuse_refiner_self_qkv(self.transformer) if converted <= 0: raise RuntimeError("SANA_WM_REFINER_FUSE_SELF_QKV=1 fused no self-attention QKV modules.") self._self_qkv_fused = True print(f"[refiner-fuse-qkv] fused {converted} self-attention QKV groups", flush=True) def offload_video_unused_audio_modules(self, device: torch.device | str = "cpu") -> None: """Keep LTX-2 audio-only branches off GPU for this wrapper's video-only forward.""" _offload_video_unused_audio_modules(self.transformer, device) _empty_cuda_cache() def move_video_modules(self, device: torch.device | str) -> None: """Move only the modules and direct parameters used by the video-only forward.""" _move_ltx2_video_modules_to(self.transformer, device) _empty_cuda_cache() def _nvfp4_autocast(self): if not self._te_nvfp4_converted: return nullcontext() import transformer_engine.pytorch as te return te.fp8_autocast(enabled=True, fp8_recipe=self._te_nvfp4_recipe) def _attention_backend_context(self): if not self._attention_backend: return nullcontext() from diffusers.models.attention_dispatch import attention_backend return attention_backend(self._attention_backend) def _uniform_timestep_tensors( self, *, batch_size: int, seq_len: int, sigma: float, ) -> tuple[torch.Tensor, torch.Tensor]: sigma_value = float(sigma) if not _env_flag("SANA_WM_REFINER_TIMESTEP_CACHE"): raw_sigma = torch.full( (int(batch_size), int(seq_len), 1), sigma_value, dtype=torch.float32, device=self.device ) model_timestep = raw_sigma.squeeze(-1) * float(self.transformer.config.timestep_scale_multiplier) return model_timestep, raw_sigma key = (int(batch_size), int(seq_len), sigma_value, str(self.device)) cached = self._uniform_timestep_cache.get(key) if cached is not None: return cached raw_sigma = torch.full((int(batch_size), int(seq_len), 1), sigma_value, dtype=torch.float32, device=self.device) model_timestep = raw_sigma.squeeze(-1) * float(self.transformer.config.timestep_scale_multiplier) cached = (model_timestep, raw_sigma) self._uniform_timestep_cache[key] = cached return cached @torch.inference_mode() def refine_latents( self, sana_latent: torch.Tensor, prompt: str, *, fps: float, sink_size: int = 1, seed: int = 42, progress: bool = True, block_size: int | None = None, kv_max_frames: int = 11, sigmas: tuple[float, ...] = STAGE_2_DISTILLED_SIGMA_VALUES, ) -> torch.Tensor: """Run the LTX-2 refiner and return refined VAE latents. When ``block_size`` is ``None`` (default), uses the legacy single-shot path that denoises all current frames jointly. When ``block_size`` is set (canonical: 3), runs the chunk-causal AR recipe with sliding-window attention over ``[source_sink + recent_history + active_block]``, matching tian's ``run_reforcing_inference`` contract — the model was trained to refine ``block_size`` frames at a time with clean prior context, and feeding the full sequence at once is out-of-distribution. Args: sana_latent: ``(B, C, F, H, W)`` stage-1 latent. prompt: text prompt. fps: video frame rate (drives LTX-2 RoPE temporal scaling). sink_size: how many leading raw ``z_sana`` frames to anchor as the attention sink (canonical: 1). seed: noise seed for the FM endpoint. progress: show a tqdm bar. block_size: latent frames per AR block (canonical: 3). ``None`` disables AR mode. kv_max_frames: maximum context+active frames retained in the sliding window when AR mode is active (canonical: 11 = 1 sink + 10 recent). sigmas: descending Euler schedule terminating at 0.0 (canonical 3-step distilled: ``(0.909375, 0.725, 0.421875, 0.0)``). """ if sana_latent.shape[2] <= sink_size: raise ValueError(f"Stage-1 latent has {sana_latent.shape[2]} frames but sink_size={sink_size}.") self.transformer.to("cpu") _empty_cuda_cache() prompt_embeds, prompt_attention_mask = self._encode_prompt(prompt) self.move_video_modules(self.device) self.offload_video_unused_audio_modules("cpu") self.prepare_transformer_nvfp4() z = sana_latent.to(device=self.device, dtype=self.dtype) sigmas_t = torch.tensor(sigmas, dtype=torch.float32, device=self.device) start_sigma = float(sigmas_t[0]) if block_size is not None: return self._refine_latents_ar( z=z, prompt_embeds=prompt_embeds, prompt_attention_mask=prompt_attention_mask, fps=fps, sigmas=sigmas_t, source_sink_frames=int(sink_size), block_size=int(block_size), kv_max_frames=int(kv_max_frames), seed=int(seed), progress=bool(progress), ) sink = z[:, :, :sink_size].contiguous() current = z[:, :, sink_size:].contiguous() generator = torch.Generator(device=self.device).manual_seed(int(seed)) eps = torch.randn(current.shape, generator=generator, device=self.device, dtype=self.dtype) noisy = (1.0 - start_sigma) * current + start_sigma * eps iterator = range(len(sigmas_t) - 1) if progress: from tqdm.auto import tqdm iterator = tqdm(iterator, desc="refiner", unit="step") for step_index in iterator: sigma = sigmas_t[step_index] denoised = self._predict_current_x0( sink=sink, noisy_current=noisy, prompt_embeds=prompt_embeds, prompt_attention_mask=prompt_attention_mask, sigma=sigma, fps=fps, ) noisy_tokens = _pack_latents( noisy, patch_size=self.transformer.config.patch_size, patch_size_t=self.transformer.config.patch_size_t, ) velocity = (noisy_tokens.float() - denoised.float()) / sigma.float() next_tokens = noisy_tokens.float() + velocity * (sigmas_t[step_index + 1] - sigma).float() noisy = _unpack_latents( next_tokens.to(self.dtype), num_frames=noisy.shape[2], height=noisy.shape[3], width=noisy.shape[4], patch_size=self.transformer.config.patch_size, patch_size_t=self.transformer.config.patch_size_t, ) return torch.cat([sink, noisy], dim=2) @torch.inference_mode() def _refine_latents_ar( self, *, z: torch.Tensor, prompt_embeds: torch.Tensor, prompt_attention_mask: torch.Tensor, fps: float, sigmas: torch.Tensor, source_sink_frames: int, block_size: int, kv_max_frames: int, seed: int, progress: bool, ) -> torch.Tensor: """Chunk-causal AR refinement — thin wrapper around ``RefinerChunkRunner``. Implements the canonical ``rf_shifted_sink`` KV-cache contract end-to-end: 1. Pre-capture **pre-RoPE** sink K/V from raw ``z_sana[:source_sink_frames]`` at σ=0 (``_kv_cache_capture`` hook). The sink frames themselves are **never refined** — they sit unchanged in the output volume. 2. AR blocks cover frames ``[source_sink_frames, T_full)`` in ``block_size``-frame chunks. For each block: - Initialize ``x_t = (1-σ₀)·z_sana_block + σ₀·ε`` (single eps per block). - 3-step deterministic Euler. Each step injects the per-layer prefix ``{sink_k_pre, sink_v, sink_pe, history_k, history_v}`` where ``sink_pe`` is rebuilt at ``sink_rope_offset = active_start - history_frames - source_sink_frames`` so the sink slides to sit immediately before the bounded working cache (official RF layout). - Capture **post-RoPE** K/V from the refined block under the same prefix (``_tf_capture_kv`` hook); append to ``history_kv_post`` and trim to ``kv_max_frames - source_sink_frames``. For the chunk-pipelined interactive path, build a ``RefinerChunkRunner`` directly and feed one block at a time as stage-1 yields it. The returned tensor has the same shape ``(B, C, T_full, H, W)`` as ``z``; the first ``source_sink_frames`` slots carry the raw sink latents unchanged, the rest carry the refined output. """ runner = RefinerChunkRunner( self, prompt_embeds=prompt_embeds, prompt_attention_mask=prompt_attention_mask, fps=fps, sigmas=sigmas, source_sink_frames=int(source_sink_frames), block_size=int(block_size), kv_max_frames=int(kv_max_frames), seed=int(seed), spatial_shape=(int(z.shape[3]), int(z.shape[4])), ) T_full = z.shape[2] sink_size = int(source_sink_frames) # Output keeps the raw sink prefix verbatim; AR blocks fill frames # [sink_size, T_full). output = z.clone() n_active = max(T_full - sink_size, 0) n_blocks = (n_active + block_size - 1) // block_size if n_active > 0 else 0 iterator = range(n_blocks) if progress: from tqdm.auto import tqdm iterator = tqdm(iterator, desc="refiner-ar", unit="block") for block_idx in iterator: block_start = sink_size + block_idx * block_size block_end = min(block_start + block_size, T_full) clean_block = z[:, :, block_start:block_end] refined = runner.refine_block( block_idx=block_idx, clean_block=clean_block, block_start=block_start, block_end=block_end, sink_seed_frames=(z[:, :, :sink_size] if block_idx == 0 else None), ) output[:, :, block_start:block_end] = refined return output def _predict_x0_active_block( self, *, active: torch.Tensor, # (B, C, N_active, H, W) at σ_cur active_positions: list[int], sigma_cur: float, prompt_embeds: torch.Tensor, prompt_attention_mask: torch.Tensor, fps: float, kv_prefix_per_layer: list[dict[str, object]] | None, active_video_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, capture_post_kv: bool = False, capture_layer_mask: list[bool] | None = None, ) -> torch.Tensor | tuple[torch.Tensor, list[tuple[torch.Tensor, torch.Tensor] | None]]: """Forward through the transformer on the ACTIVE BLOCK ONLY and return x0. The active block's Q attends to ``[prefix, current]`` K/V via the ``_tf_kv_prefix`` hook on every self-attention block. All active tokens carry the same ``sigma_cur`` (matching tian's per-block uniform σ). """ latent_tokens = _pack_latents( active, patch_size=self.transformer.config.patch_size, patch_size_t=self.transformer.config.patch_size_t, ) batch_size, seq_len, _ = latent_tokens.shape # Use a per-token uniform sigma for the active block. model_timestep, raw_sigma = self._uniform_timestep_tensors( batch_size=int(batch_size), seq_len=int(seq_len), sigma=float(sigma_cur), ) video_rotary_emb = active_video_rotary_emb if video_rotary_emb is None: video_rotary_emb = _build_rotary_emb_for_absolute_positions( transformer=self.transformer, batch_size=batch_size, frame_positions=active_positions, height=int(active.shape[3]), width=int(active.shape[4]), device=self.device, fps=float(fps), ) # Replace the per-frame uniform-σ adaLN time embedding with the active # block's mean sigma (= sigma_cur here), mirroring tian's prompt_sigma # `mean_active` mode. _set_kv_prefix_on_blocks(self.transformer, kv_prefix_per_layer) if capture_post_kv: _set_capture_flag_on_blocks(self.transformer, "post_rope", enable=True, layer_mask=capture_layer_mask) try: velocity = self._forward_video_only_with_rope( hidden_states=latent_tokens, encoder_hidden_states=prompt_embeds, timestep=model_timestep, encoder_attention_mask=prompt_attention_mask, video_rotary_emb=video_rotary_emb, n_context_tokens=0, ) finally: if capture_post_kv: _set_capture_flag_on_blocks(self.transformer, "post_rope", enable=False) _clear_kv_prefix_on_blocks(self.transformer) captured_kv = ( _collect_captured_kv_from_blocks(self.transformer, "post_rope", layer_mask=capture_layer_mask) if capture_post_kv else None ) # FM x0 prediction: x_t - σ_cur · v. denoised_tokens = latent_tokens.float() - velocity.float() * raw_sigma denoised = _unpack_latents( denoised_tokens.to(self.dtype), num_frames=int(active.shape[2]), height=int(active.shape[3]), width=int(active.shape[4]), patch_size=self.transformer.config.patch_size, patch_size_t=self.transformer.config.patch_size_t, ) if captured_kv is not None: return denoised, captured_kv return denoised @torch.inference_mode() def _capture_block_kv( self, *, clean_block: torch.Tensor, # (B, C, N, H, W) treated as σ=0 (clean) input frame_positions: list[int], prompt_embeds: torch.Tensor, prompt_attention_mask: torch.Tensor, fps: float, capture_mode: str, # "pre_rope" or "post_rope" kv_prefix_per_layer: list[dict[str, object]] | None, capture_layer_mask: list[bool] | None = None, video_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, ) -> list[tuple[torch.Tensor, torch.Tensor] | None]: """Run one forward at σ=0 with capture hooks; return per-layer (K, V). ``capture_mode='pre_rope'`` uses the ``_kv_cache_capture`` hook (stored before RoPE so a future window can re-RoPE the sink to its shifted offset). ``capture_mode='post_rope'`` uses ``_tf_capture_kv`` (stored with RoPE already baked at the block's absolute positions, ready to concatenate into the next window's prefix). """ latent_tokens = _pack_latents( clean_block, patch_size=self.transformer.config.patch_size, patch_size_t=self.transformer.config.patch_size_t, ) batch_size, seq_len, _ = latent_tokens.shape model_timestep = torch.zeros(batch_size, seq_len, dtype=torch.float32, device=self.device) if video_rotary_emb is None: video_rotary_emb = _build_rotary_emb_for_absolute_positions( transformer=self.transformer, batch_size=batch_size, frame_positions=frame_positions, height=int(clean_block.shape[3]), width=int(clean_block.shape[4]), device=self.device, fps=float(fps), ) stop_after_layer = None stop_after_capture_kv_layer = None if capture_layer_mask is not None and not all(capture_layer_mask): stop_after_layer = max(idx for idx, keep in enumerate(capture_layer_mask) if keep) if _env_flag("SANA_WM_REFINER_CAPTURE_KV_ONLY_LAST"): if capture_layer_mask is None: stop_after_capture_kv_layer = len(self.transformer.transformer_blocks) - 1 else: stop_after_capture_kv_layer = max(idx for idx, keep in enumerate(capture_layer_mask) if keep) stop_after_layer = None _set_kv_prefix_on_blocks(self.transformer, kv_prefix_per_layer) _set_capture_flag_on_blocks(self.transformer, capture_mode, enable=True, layer_mask=capture_layer_mask) try: _ = self._forward_video_only_with_rope( hidden_states=latent_tokens, encoder_hidden_states=prompt_embeds, timestep=model_timestep, encoder_attention_mask=prompt_attention_mask, video_rotary_emb=video_rotary_emb, n_context_tokens=0, skip_output_projection=True, stop_after_layer=stop_after_layer, stop_after_capture_kv_layer=stop_after_capture_kv_layer, ) finally: _set_capture_flag_on_blocks(self.transformer, capture_mode, enable=False) _clear_kv_prefix_on_blocks(self.transformer) return _collect_captured_kv_from_blocks(self.transformer, capture_mode, layer_mask=capture_layer_mask) @torch.inference_mode() def _encode_prompt(self, prompt: str) -> tuple[torch.Tensor, torch.Tensor]: from transformers import AutoTokenizer, Gemma3ForConditionalGeneration tokenizer = AutoTokenizer.from_pretrained(self.gemma_root) tokenizer.padding_side = "left" if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token text_inputs = tokenizer( [prompt.strip()], padding="max_length", max_length=self.text_max_sequence_length, truncation=True, add_special_tokens=True, return_tensors="pt", ) input_ids = text_inputs.input_ids.to(self.device) attention_mask = text_inputs.attention_mask.to(self.device) text_encoder = Gemma3ForConditionalGeneration.from_pretrained( self.gemma_root, torch_dtype=self.dtype, low_cpu_mem_usage=True, ).eval() text_encoder.to(self.device) text_backbone = getattr(text_encoder, "model", text_encoder) outputs = text_backbone(input_ids=input_ids, attention_mask=attention_mask, output_hidden_states=True) hidden_states = torch.stack(outputs.hidden_states, dim=-1) sequence_lengths = attention_mask.sum(dim=-1) del text_encoder, text_backbone, outputs _empty_cuda_cache() prompt_embeds = _pack_text_embeds( hidden_states, sequence_lengths, device=self.device, padding_side=tokenizer.padding_side, ).to(dtype=self.dtype) del hidden_states _empty_cuda_cache() self.connectors.to(self.device) connector_prompt_embeds, _, connector_attention_mask = self.connectors(prompt_embeds, attention_mask) self.connectors.to("cpu") del prompt_embeds, attention_mask _empty_cuda_cache() return connector_prompt_embeds.to(device=self.device, dtype=self.dtype), connector_attention_mask.to( device=self.device ) def _predict_current_x0( self, *, sink: torch.Tensor, noisy_current: torch.Tensor, prompt_embeds: torch.Tensor, prompt_attention_mask: torch.Tensor, sigma: torch.Tensor, fps: float, ) -> torch.Tensor: full_latent = torch.cat([sink, noisy_current], dim=2) batch_size, _, num_frames, height, width = full_latent.shape latent_tokens = _pack_latents( full_latent, patch_size=self.transformer.config.patch_size, patch_size_t=self.transformer.config.patch_size_t, ) n_context_tokens = _pack_latents( sink, patch_size=self.transformer.config.patch_size, patch_size_t=self.transformer.config.patch_size_t, ).shape[1] raw_timestep = torch.zeros(batch_size, latent_tokens.shape[1], 1, dtype=torch.float32, device=self.device) raw_timestep[:, n_context_tokens:, 0] = sigma.float() model_timestep = raw_timestep.squeeze(-1) * float(self.transformer.config.timestep_scale_multiplier) velocity = self._forward_video_only( hidden_states=latent_tokens, encoder_hidden_states=prompt_embeds, timestep=model_timestep, encoder_attention_mask=prompt_attention_mask, num_frames=num_frames, height=height, width=width, fps=fps, n_context_tokens=n_context_tokens, ) denoised = latent_tokens.float() - velocity.float() * raw_timestep return denoised[:, n_context_tokens:, :].to(self.dtype) def _forward_video_only_with_rope( self, *, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, timestep: torch.Tensor, encoder_attention_mask: torch.Tensor | None, video_rotary_emb: tuple[torch.Tensor, torch.Tensor], n_context_tokens: int, skip_output_projection: bool = False, stop_after_layer: int | None = None, stop_after_capture_kv_layer: int | None = None, ) -> torch.Tensor: """Shared body of ``_forward_video_only`` that takes a pre-built RoPE. Used by the AR refinement path where each block forward needs custom per-frame absolute positions in the source video. """ transformer = self.transformer batch_size = hidden_states.size(0) seq_len = int(hidden_states.shape[1]) profiler = None if _refiner_layer_profile_enabled(): forward_kind = "capture" if skip_output_projection else "predict" prefix_tokens = _current_refiner_prefix_tokens(transformer) profiler = _RefinerLayerCudaProfiler( enabled=True, device=self.device, label=f"{forward_kind} seq={seq_len} prefix={prefix_tokens}", ) with _profile_section(profiler, "mask_prepare"): encoder_attention_mask = _prepare_encoder_attention_mask(encoder_attention_mask, hidden_states.dtype) with _profile_section(profiler, "proj_in"): hidden_states = transformer.proj_in(hidden_states) with _profile_section(profiler, "time_embed"): temb, embedded_timestep = transformer.time_embed( timestep.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype, ) temb = temb.view(batch_size, -1, temb.size(-1)) embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1)) with _profile_section(profiler, "caption_projection"): if _has_cross_attention_kv_cache(transformer): encoder_hidden_states = None else: encoder_hidden_states = transformer.caption_projection(encoder_hidden_states) encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1)) with self._attention_backend_context(), self._nvfp4_autocast(): for layer_idx, block in enumerate(transformer.transformer_blocks): capture_kv_only = stop_after_capture_kv_layer is not None and layer_idx >= int( stop_after_capture_kv_layer ) hidden_states = _forward_video_block( block=block, hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb, video_rotary_emb=video_rotary_emb, encoder_attention_mask=encoder_attention_mask, n_context_tokens=n_context_tokens, profiler=profiler, capture_kv_only=capture_kv_only, ) if capture_kv_only: break if stop_after_layer is not None and layer_idx >= int(stop_after_layer): break if skip_output_projection: if profiler is not None: profiler.finish() return hidden_states with _profile_section(profiler, "proj_out"): scale_shift_values = transformer.scale_shift_table[None, None] + embedded_timestep[:, :, None] shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] hidden_states = transformer.norm_out(hidden_states) hidden_states = hidden_states * (1 + scale) + shift hidden_states = transformer.proj_out(hidden_states) if profiler is not None: profiler.finish() return hidden_states def _forward_video_only( self, *, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, timestep: torch.Tensor, encoder_attention_mask: torch.Tensor | None, num_frames: int, height: int, width: int, fps: float, n_context_tokens: int, ) -> torch.Tensor: transformer = self.transformer batch_size = hidden_states.size(0) seq_len = int(hidden_states.shape[1]) profiler = None if _refiner_layer_profile_enabled(): profiler = _RefinerLayerCudaProfiler( enabled=True, device=self.device, label=f"legacy seq={seq_len} prefix={int(n_context_tokens)}", ) with _profile_section(profiler, "mask_prepare"): encoder_attention_mask = _prepare_encoder_attention_mask(encoder_attention_mask, hidden_states.dtype) with _profile_section(profiler, "rope"): video_coords = transformer.rope.prepare_video_coords( batch_size, num_frames, height, width, hidden_states.device, fps=fps ) video_rotary_emb = transformer.rope(video_coords, device=hidden_states.device) with _profile_section(profiler, "proj_in"): hidden_states = transformer.proj_in(hidden_states) with _profile_section(profiler, "time_embed"): temb, embedded_timestep = transformer.time_embed( timestep.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype, ) temb = temb.view(batch_size, -1, temb.size(-1)) embedded_timestep = embedded_timestep.view(batch_size, -1, embedded_timestep.size(-1)) with _profile_section(profiler, "caption_projection"): if _has_cross_attention_kv_cache(transformer): encoder_hidden_states = None else: encoder_hidden_states = transformer.caption_projection(encoder_hidden_states) encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_states.size(-1)) with self._attention_backend_context(), self._nvfp4_autocast(): for block in transformer.transformer_blocks: hidden_states = _forward_video_block( block=block, hidden_states=hidden_states, encoder_hidden_states=encoder_hidden_states, temb=temb, video_rotary_emb=video_rotary_emb, encoder_attention_mask=encoder_attention_mask, n_context_tokens=n_context_tokens, profiler=profiler, ) with _profile_section(profiler, "proj_out"): scale_shift_values = transformer.scale_shift_table[None, None] + embedded_timestep[:, :, None] shift, scale = scale_shift_values[:, :, 0], scale_shift_values[:, :, 1] hidden_states = transformer.norm_out(hidden_states) hidden_states = hidden_states * (1 + scale) + shift hidden_states = transformer.proj_out(hidden_states) if profiler is not None: profiler.finish() return hidden_states class RefinerChunkRunner: """Stateful per-AR-block driver for ``DiffusersLTX2Refiner``. Owns the rolling KV state that the chunk-causal AR recipe accumulates as refiner blocks complete: * ``_sink_kv_pre``: per-layer pre-RoPE K/V captured from the first ``source_sink_frames`` raw stage-1 latents at σ=0. Lazily filled on the first call to :meth:`refine_block` (the orchestrator only has the first stage-1 chunk in hand by then). * ``_history_kv_post``: per-layer post-RoPE K/V of every refined block already produced, trimmed to ``kv_max_frames - source_sink_frames`` frames so the sliding window stays bounded. * ``_history_frames``: number of frames currently in ``_history_kv_post`` (drives token-level trim). The numerical contract is identical to a single in-place call to ``_refine_latents_ar``: same RNG-seeded epsilon stream consumed block-by-block, same ``rf_shifted_sink`` per-window prefix dict, same 3-step deterministic Euler, same post-RoPE capture under that prefix. The orchestrator can therefore call :meth:`refine_block` once per stage-1 chunk without changing inference semantics, and concurrently launch the downstream causal-VAE decode on a separate CUDA stream while the next block's refinement runs on the refiner stream. """ def __init__( self, refiner: DiffusersLTX2Refiner, *, prompt_embeds: torch.Tensor, prompt_attention_mask: torch.Tensor, fps: float, sigmas: torch.Tensor, source_sink_frames: int, block_size: int, kv_max_frames: int, seed: int, spatial_shape: tuple[int, int], n_active_frames: int | None = None, latent_channels: int | None = None, batch_size: int = 1, ) -> None: self._refiner = refiner self._prompt_embeds = prompt_embeds self._prompt_attention_mask = prompt_attention_mask self._fps = float(fps) self._sigmas = sigmas self._sigma_values = [float(v) for v in sigmas.detach().float().cpu()] self._sigma_pairs = list(zip(self._sigma_values[:-1], self._sigma_values[1:])) self._sigma_max = self._sigma_values[0] self._n_steps = int(sigmas.numel() - 1) self._source_sink_frames = int(source_sink_frames) self._block_size = int(block_size) self._kv_max_frames = int(kv_max_frames) self._max_history_frames = int(kv_max_frames) - int(source_sink_frames) self._device = refiner.device self._dtype = refiner.dtype self._generator = torch.Generator(device=self._device).manual_seed(int(seed)) self._kv_cache_storage_dtype = _resolve_kv_cache_storage_dtype() transformer = refiner.transformer self._n_layers = len(transformer.transformer_blocks) H, W = spatial_shape self._H, self._W = int(H), int(W) self._tokens_per_frame = ( int(H // transformer.config.patch_size) * int(W // transformer.config.patch_size) * int(transformer.config.patch_size_t) ) self._precomputed_eps_blocks: list[torch.Tensor] | None = None if ( _env_flag("SANA_WM_REFINER_PREGENERATE_NOISE") and n_active_frames is not None and latent_channels is not None ): n_active = int(n_active_frames) channels = int(latent_channels) batch = int(batch_size) n_blocks = (n_active + self._block_size - 1) // self._block_size if n_active > 0 else 0 self._precomputed_eps_blocks = [] for block_idx in range(n_blocks): active_len = min(self._block_size, n_active - block_idx * self._block_size) self._precomputed_eps_blocks.append( torch.randn( (batch, channels, active_len, self._H, self._W), generator=self._generator, device=self._device, dtype=self._dtype, ) ) print(f"[refiner-noise] precomputed {len(self._precomputed_eps_blocks)} eps blocks", flush=True) if _env_flag("SANA_WM_REFINER_CROSS_ATTN_KV_CACHE"): with refiner._nvfp4_autocast(): _set_cross_attention_kv_cache(refiner.transformer, prompt_embeds, prompt_attention_mask) else: _clear_cross_attention_kv_cache(refiner.transformer) self._sink_kv_pre: list[tuple[torch.Tensor, torch.Tensor]] | None = None self._history_kv_post: list[tuple[torch.Tensor, torch.Tensor] | None] = [None] * self._n_layers self._history_frames: int = 0 self._history_layer_mask = _refiner_history_layer_mask(self._n_layers) self._exact_capture_layer_mask = _refiner_exact_capture_layer_mask( self._n_layers, default_mask=self._history_layer_mask, ) if not all(self._history_layer_mask): kept = sum(1 for keep in self._history_layer_mask if keep) print( f"[refiner-history] recent history enabled on {kept}/{self._n_layers} layers", flush=True, ) if self._exact_capture_layer_mask != self._history_layer_mask: kept = sum(1 for keep in self._exact_capture_layer_mask if keep) print( f"[refiner-history] exact post-capture on {kept}/{self._n_layers} layers", flush=True, ) @torch.inference_mode() def pre_capture_sink(self, sink_seed_frames: torch.Tensor) -> None: """Capture the source-sink K/V before the first active refiner block. The sink is just the conditioning latent frame and does not depend on stage-1 sampling. Scheduling this on the refiner stream lets it overlap with stage-1 chunk 0 while preserving the exact same cached K/V that ``refine_block`` would have produced lazily. """ if self._sink_kv_pre is not None: return if sink_seed_frames is None: raise ValueError("pre_capture_sink requires sink_seed_frames.") if sink_seed_frames.shape[2] != self._source_sink_frames: raise ValueError( f"sink_seed_frames has {sink_seed_frames.shape[2]} frames " f"but source_sink_frames={self._source_sink_frames}." ) source_sink = sink_seed_frames.contiguous() self._sink_kv_pre = [ _store_kv_pair(pair, self._kv_cache_storage_dtype) for pair in self._refiner._capture_block_kv( clean_block=source_sink, frame_positions=list(range(self._source_sink_frames)), prompt_embeds=self._prompt_embeds, prompt_attention_mask=self._prompt_attention_mask, fps=self._fps, capture_mode="pre_rope", kv_prefix_per_layer=None, ) ] @torch.inference_mode() def refine_block( self, *, block_idx: int, clean_block: torch.Tensor, block_start: int, block_end: int, sink_seed_frames: torch.Tensor | None = None, ) -> torch.Tensor: """Refine one AR block; advance internal KV state. Args: block_idx: 0-based block index in the AR schedule. Used only for bookkeeping; positional state derives from ``block_start``. clean_block: ``(B, C, active_len, H, W)`` clean stage-1 latents covering frames ``[block_start, block_end)``. The active block is what actually gets refined; sink frames live outside the active range and are passed via ``sink_seed_frames`` on the first call. block_start: absolute latent-frame index of the active block's first frame (drives the ``rf_shifted_sink`` RoPE offset). Must be >= ``source_sink_frames`` so the sink doesn't overlap the active region. block_end: absolute latent-frame index just past the active block. sink_seed_frames: ``(B, C, source_sink_frames, H, W)`` raw sink latents used once on the very first ``refine_block`` call to pre-capture the pre-RoPE sink K/V at ``sigma=0`` with frame positions ``[0, source_sink_frames)``. Required on the first call; ignored thereafter. The orchestrator owns these — they are typically the first ``source_sink_frames`` of stage-1's first chunk. Returns: ``(B, C, active_len, H, W)`` refined latents for this block. """ refiner = self._refiner device = self._device profiler = _RefinerCudaProfiler(enabled=_refiner_profile_enabled(), device=device, block_idx=int(block_idx)) B = int(clean_block.shape[0]) active_len = block_end - block_start if block_start < self._source_sink_frames: raise ValueError( f"block_start={block_start} overlaps the source sink " f"(source_sink_frames={self._source_sink_frames})." ) # 1) On the first call: pre-capture PRE-RoPE sink K/V from the supplied # raw sink latents at sigma=0 with absolute positions [0, sink_size). if self._sink_kv_pre is None: with profiler.section("sink_capture"): if sink_seed_frames is None: raise ValueError("First refine_block call requires sink_seed_frames " "(raw stage-1 sink latents).") if sink_seed_frames.shape[2] != self._source_sink_frames: raise ValueError( f"sink_seed_frames has {sink_seed_frames.shape[2]} frames " f"but source_sink_frames={self._source_sink_frames}." ) self.pre_capture_sink(sink_seed_frames) # 2) Build per-window kv_prefix dict per layer. with profiler.section("prefix_build"): sink_rope_offset_history = block_start - self._history_frames - self._source_sink_frames sink_rope_offset_no_history = block_start - self._source_sink_frames sink_pe_history = _build_rotary_emb_for_absolute_positions( transformer=refiner.transformer, batch_size=B, frame_positions=list( range(sink_rope_offset_history, sink_rope_offset_history + self._source_sink_frames) ), height=self._H, width=self._W, device=device, fps=self._fps, ) sink_pe_no_history = sink_pe_history if sink_rope_offset_no_history != sink_rope_offset_history: sink_pe_no_history = _build_rotary_emb_for_absolute_positions( transformer=refiner.transformer, batch_size=B, frame_positions=list( range(sink_rope_offset_no_history, sink_rope_offset_no_history + self._source_sink_frames) ), height=self._H, width=self._W, device=device, fps=self._fps, ) kv_prefix_per_layer: list[dict[str, object]] = [] preconcat_prefix = _env_flag("SANA_WM_REFINER_PRECONCAT_PREFIX") empty_cache_before_prefix = _env_flag("SANA_WM_REFINER_EMPTY_CACHE_BEFORE_PREFIX") for layer_idx in range(self._n_layers): hk = self._history_kv_post[layer_idx] use_history = bool(self._history_layer_mask[layer_idx] and hk is not None and hk[0].shape[1] > 0) sink_pe = sink_pe_history if use_history else sink_pe_no_history prefix: dict[str, object] = { "mode": "rf_shifted_sink", "sink_k_pre": self._sink_kv_pre[layer_idx][0], "sink_v": self._sink_kv_pre[layer_idx][1], "sink_pe": sink_pe, "history_k": (hk[0] if use_history else None), "history_v": (hk[1] if use_history else None), } if preconcat_prefix: prefix_k_parts: list[torch.Tensor] = [] prefix_v_parts: list[torch.Tensor] = [] sink_k_pre, sink_v = self._sink_kv_pre[layer_idx] if sink_k_pre.shape[1] > 0 and sink_v.shape[1] > 0: attn = refiner.transformer.transformer_blocks[layer_idx].attn1 sink_k = _apply_refiner_rotary(attn, sink_k_pre.to(self._dtype), sink_pe) prefix_k_parts.append(sink_k) prefix_v_parts.append(sink_v.to(self._dtype)) if use_history: prefix_k_parts.append(hk[0].to(self._dtype)) prefix_v_parts.append(hk[1].to(self._dtype)) if prefix_k_parts: if empty_cache_before_prefix and device.type == "cuda": torch.cuda.empty_cache() prefix_k = torch.cat(prefix_k_parts, dim=1) prefix_v = torch.cat(prefix_v_parts, dim=1) prefix["prefix_k"] = prefix_k prefix["prefix_v"] = prefix_v kv_prefix_per_layer.append(prefix) # 3) FM endpoint at sigma=sigma0: single epsilon per block. with profiler.section("noise_init"): eps = None if self._precomputed_eps_blocks is not None and int(block_idx) < len(self._precomputed_eps_blocks): candidate_eps = self._precomputed_eps_blocks[int(block_idx)] if tuple(candidate_eps.shape) == tuple(clean_block.shape): eps = candidate_eps if eps is None: eps = torch.randn(clean_block.shape, generator=self._generator, device=device, dtype=self._dtype) x_t = ((1.0 - self._sigma_max) * clean_block.float() + self._sigma_max * eps.float()).to(self._dtype) with profiler.section("active_rope"): active_positions = list(range(int(block_start), int(block_end))) active_video_rotary_emb = _build_rotary_emb_for_absolute_positions( transformer=refiner.transformer, batch_size=B, frame_positions=active_positions, height=self._H, width=self._W, device=device, fps=self._fps, ) fast_kv_capture = _refiner_fast_kv_capture_mode() reuse_final_predict_kv = fast_kv_capture == "last_predict" and not _refiner_fast_kv_needs_clean_block( int(block_idx) ) fill_missing_predict_kv = fast_kv_capture == "fill_missing" captured_kv_post: list[tuple[torch.Tensor, torch.Tensor] | None] | None = None n_sigma_pairs = len(self._sigma_pairs) for step_idx, (sigma_cur, sigma_next) in enumerate(self._sigma_pairs): with profiler.section(f"denoise_step{step_idx}"): capture_predict_kv = bool( (reuse_final_predict_kv or fill_missing_predict_kv) and step_idx == n_sigma_pairs - 1 ) pred_result = refiner._predict_x0_active_block( active=x_t, active_positions=active_positions, sigma_cur=sigma_cur, prompt_embeds=self._prompt_embeds, prompt_attention_mask=self._prompt_attention_mask, fps=self._fps, kv_prefix_per_layer=kv_prefix_per_layer, active_video_rotary_emb=active_video_rotary_emb, capture_post_kv=capture_predict_kv, capture_layer_mask=self._history_layer_mask, ) if isinstance(pred_result, tuple): pred_x0, captured_kv_post = pred_result if fill_missing_predict_kv and captured_kv_post is not None: captured_kv_post = [ ( None if self._exact_capture_layer_mask[layer_idx] else (_store_kv_pair(pair, self._kv_cache_storage_dtype) if pair is not None else None) ) for layer_idx, pair in enumerate(captured_kv_post) ] else: pred_x0 = pred_result if sigma_cur <= 1.0e-6: x_t = pred_x0.to(self._dtype) else: ratio = sigma_next / sigma_cur x_t = (ratio * x_t.float() + (1.0 - ratio) * pred_x0.float()).to(self._dtype) pred_x0 = None if self._max_history_frames <= 0: with profiler.section("history_update"): self._history_frames = 0 for layer_idx in range(self._n_layers): self._history_kv_post[layer_idx] = None profiler.finish() return x_t # 4) Capture POST-RoPE K/V for this refined block under the same prefix. with profiler.section("post_capture"): if reuse_final_predict_kv: if captured_kv_post is None: raise RuntimeError("SANA_WM_REFINER_FAST_KV_CAPTURE=last_predict did not capture post-RoPE K/V.") block_kv_post = captured_kv_post else: if _refiner_empty_cache_before_capture() and device.type == "cuda": torch.cuda.empty_cache() block_kv_post = refiner._capture_block_kv( clean_block=x_t, frame_positions=active_positions, prompt_embeds=self._prompt_embeds, prompt_attention_mask=self._prompt_attention_mask, fps=self._fps, capture_mode="post_rope", kv_prefix_per_layer=kv_prefix_per_layer, capture_layer_mask=self._exact_capture_layer_mask, video_rotary_emb=active_video_rotary_emb, ) if fill_missing_predict_kv: if captured_kv_post is None: raise RuntimeError("SANA_WM_REFINER_FAST_KV_CAPTURE=fill_missing did not capture fallback K/V.") block_kv_post = [ exact_pair if self._exact_capture_layer_mask[layer_idx] else captured_kv_post[layer_idx] for layer_idx, exact_pair in enumerate(block_kv_post) ] with profiler.section("history_update"): for layer_idx in range(self._n_layers): if not self._history_layer_mask[layer_idx]: self._history_kv_post[layer_idx] = None continue raw_pair = block_kv_post[layer_idx] if raw_pair is None: raise RuntimeError(f"Missing post-RoPE K/V capture for history layer {layer_idx}.") raw_k, raw_v = raw_pair new_k = _store_kv_tensor(raw_k, self._kv_cache_storage_dtype) new_v = _store_kv_tensor(raw_v, self._kv_cache_storage_dtype) block_kv_post[layer_idx] = (new_k, new_v) old = self._history_kv_post[layer_idx] if old is None: if self._max_history_frames > 0 and active_len > self._max_history_frames: keep_tokens = self._max_history_frames * self._tokens_per_frame self._history_kv_post[layer_idx] = (new_k[:, -keep_tokens:], new_v[:, -keep_tokens:]) else: self._history_kv_post[layer_idx] = (new_k, new_v) else: if self._max_history_frames > 0: keep_old_frames = max(0, self._max_history_frames - active_len) keep_old_tokens = keep_old_frames * self._tokens_per_frame old = ( old[0][:, -keep_old_tokens:] if keep_old_tokens > 0 else old[0][:, :0], old[1][:, -keep_old_tokens:] if keep_old_tokens > 0 else old[1][:, :0], ) self._history_kv_post[layer_idx] = ( torch.cat([old[0], new_k], dim=1), torch.cat([old[1], new_v], dim=1), ) raw_k = None raw_v = None self._history_frames += active_len if self._max_history_frames > 0 and self._history_frames > self._max_history_frames: keep_tokens = self._max_history_frames * self._tokens_per_frame for layer_idx in range(self._n_layers): hk = self._history_kv_post[layer_idx] if hk is not None: self._history_kv_post[layer_idx] = (hk[0][:, -keep_tokens:], hk[1][:, -keep_tokens:]) self._history_frames = self._max_history_frames profiler.finish() return x_t def _build_rotary_emb_for_absolute_positions( *, transformer: nn.Module, batch_size: int, frame_positions: list[int], height: int, width: int, device: torch.device, fps: float, ) -> tuple[torch.Tensor, torch.Tensor]: """Reimplement ``LTX2VideoRotaryPosEmbed.prepare_video_coords`` with explicit per-frame positions. The default helper assumes contiguous ``torch.arange(num_frames)`` which is fine for bidirectional inference; the sliding-window AR refiner needs to keep each frame's absolute index in the source video so RoPE captures the correct temporal phase across the sink + recent + active window. """ rope = transformer.rope patch_size_t = int(rope.patch_size_t) patch_size = int(rope.patch_size) f_positions = torch.tensor(frame_positions, dtype=torch.float32, device=device) if patch_size_t > 1: # Each patch covers ``patch_size_t`` latent frames; pick the start of each patch. f_positions = f_positions[::patch_size_t] int(f_positions.shape[0]) grid_h = torch.arange(start=0, end=height, step=patch_size, dtype=torch.float32, device=device) grid_w = torch.arange(start=0, end=width, step=patch_size, dtype=torch.float32, device=device) grid = torch.meshgrid(f_positions, grid_h, grid_w, indexing="ij") grid = torch.stack(grid, dim=0) # [3, N_F, N_H, N_W] patch_size_delta = torch.tensor((patch_size_t, patch_size, patch_size), dtype=grid.dtype, device=device) patch_ends = grid + patch_size_delta.view(3, 1, 1, 1) latent_coords = torch.stack([grid, patch_ends], dim=-1) latent_coords = latent_coords.flatten(1, 3).unsqueeze(0).repeat(batch_size, 1, 1, 1) scale_tensor = torch.tensor(rope.scale_factors, device=device) broadcast_shape = [1] * latent_coords.ndim broadcast_shape[1] = -1 pixel_coords = latent_coords * scale_tensor.view(*broadcast_shape) pixel_coords[:, 0, ...] = (pixel_coords[:, 0, ...] + rope.causal_offset - rope.scale_factors[0]).clamp(min=0) pixel_coords[:, 0, ...] = pixel_coords[:, 0, ...] / float(fps) return rope(pixel_coords, device=device) def _forward_video_block( *, block: nn.Module, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor | None, temb: torch.Tensor, video_rotary_emb: tuple[torch.Tensor, torch.Tensor], encoder_attention_mask: torch.Tensor | None, n_context_tokens: int, profiler: _RefinerLayerCudaProfiler | None = None, capture_kv_only: bool = False, ) -> torch.Tensor: batch_size = hidden_states.size(0) if profiler is None: norm_hidden_states = block.norm1(hidden_states) num_ada_params = block.scale_shift_table.shape[0] ada_values = block.scale_shift_table[None, None].to(temb.device) + temb.reshape( batch_size, temb.size(1), num_ada_params, -1 ) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2) norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa if capture_kv_only: _capture_streaming_self_attention_kv( attn=block.attn1, hidden_states=norm_hidden_states, query_rotary_emb=video_rotary_emb, ) return hidden_states attn_hidden_states = _streaming_self_attention( attn=block.attn1, hidden_states=norm_hidden_states, query_rotary_emb=video_rotary_emb, n_context_tokens=n_context_tokens, ) hidden_states = hidden_states + attn_hidden_states * gate_msa norm_hidden_states = block.norm2(hidden_states) cross_kv_cache = getattr(block.attn2, "_sana_cross_attn_kv_cache", None) if cross_kv_cache is None: attn_hidden_states = block.attn2( norm_hidden_states, encoder_hidden_states=encoder_hidden_states, query_rotary_emb=None, attention_mask=encoder_attention_mask, ) else: attn_hidden_states = _cross_attention_with_cached_kv(block.attn2, norm_hidden_states, cross_kv_cache) hidden_states = hidden_states + attn_hidden_states norm_hidden_states = block.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp hidden_states = hidden_states + block.ff(norm_hidden_states) * gate_mlp return hidden_states with _profile_section(profiler, "norm_adaln"): norm_hidden_states = block.norm1(hidden_states) num_ada_params = block.scale_shift_table.shape[0] ada_values = block.scale_shift_table[None, None].to(temb.device) + temb.reshape( batch_size, temb.size(1), num_ada_params, -1 ) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ada_values.unbind(dim=2) norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa with _profile_section(profiler, "self_attn"): if capture_kv_only: _capture_streaming_self_attention_kv( attn=block.attn1, hidden_states=norm_hidden_states, query_rotary_emb=video_rotary_emb, ) return hidden_states else: attn_hidden_states = _streaming_self_attention( attn=block.attn1, hidden_states=norm_hidden_states, query_rotary_emb=video_rotary_emb, n_context_tokens=n_context_tokens, ) hidden_states = hidden_states + attn_hidden_states * gate_msa with _profile_section(profiler, "cross_attn"): norm_hidden_states = block.norm2(hidden_states) cross_kv_cache = getattr(block.attn2, "_sana_cross_attn_kv_cache", None) if cross_kv_cache is None: attn_hidden_states = block.attn2( norm_hidden_states, encoder_hidden_states=encoder_hidden_states, query_rotary_emb=None, attention_mask=encoder_attention_mask, ) else: attn_hidden_states = _cross_attention_with_cached_kv(block.attn2, norm_hidden_states, cross_kv_cache) hidden_states = hidden_states + attn_hidden_states with _profile_section(profiler, "ffn"): norm_hidden_states = block.norm3(hidden_states) * (1 + scale_mlp) + shift_mlp hidden_states = hidden_states + block.ff(norm_hidden_states) * gate_mlp return hidden_states def _streaming_self_attention( *, attn: nn.Module, hidden_states: torch.Tensor, query_rotary_emb: tuple[torch.Tensor, torch.Tensor], n_context_tokens: int, ) -> torch.Tensor: """LTX-2 self-attention with sink/current streaming mask + AR KV-cache hooks. Two modes are layered on top of vanilla diffusers self-attention, selected by ``n_context_tokens`` and per-block hook attributes (set by the AR refiner): * ``n_context_tokens > 0`` (legacy single-shot path) — sink queries attend sink only, current queries attend ``[sink + current]`` via two SDPA calls. * ``n_context_tokens == 0`` (AR mode) — Q comes from the active block only; the per-block ``_tf_kv_prefix`` dict (``rf_shifted_sink``) supplies the pre-RoPE sink K/V (re-RoPE'd here with its sliding offset PE) and the post-RoPE recent-history K/V, concatenated before SDPA. The ``_kv_cache_capture`` and ``_tf_capture_kv`` hooks record K/V into the module for the AR orchestrator to read back. """ from diffusers.models.transformers.transformer_ltx2 import apply_interleaved_rotary_emb, apply_split_rotary_emb gate_logits = attn.to_gate_logits(hidden_states) if attn.to_gate_logits is not None else None fused_qkv = getattr(attn, "_sana_fused_qkv", None) if fused_qkv is not None: query, key, value = fused_qkv(hidden_states) else: query = attn.to_q(hidden_states) key = attn.to_k(hidden_states) value = attn.to_v(hidden_states) query = attn.norm_q(query) key = attn.norm_k(key) # KV-cache capture / inject hooks for ``rf_shifted_sink`` AR refinement. # Mirrors tian's ``diffusion/vendors/ltx/ltx_core/model/transformer/attention.py``: # - ``_kv_cache_capture`` saves PRE-RoPE (post-norm) K/V so a future window # can re-apply RoPE at its shifted sink offset. # - ``_tf_capture_kv`` saves POST-RoPE K/V so the next window can directly # concatenate the recent history. # - ``_tf_kv_prefix`` (a dict with ``mode='rf_shifted_sink'``) prepends a # re-RoPE'd sink + already-post-RoPE recent history before SDPA. if getattr(attn, "_kv_cache_capture", False): attn._cached_kv_pre = (_capture_kv_tensor(key), _capture_kv_tensor(value)) if attn.rope_type == "interleaved": query = apply_interleaved_rotary_emb(query, query_rotary_emb) key = apply_interleaved_rotary_emb(key, query_rotary_emb) elif attn.rope_type == "split": query = apply_split_rotary_emb(query, query_rotary_emb) key = apply_split_rotary_emb(key, query_rotary_emb) else: raise ValueError(f"Unsupported LTX-2 RoPE type: {attn.rope_type}") if getattr(attn, "_tf_capture_kv", False): attn._cached_kv_post = (_capture_kv_tensor(key), _capture_kv_tensor(value)) tf_prefix = getattr(attn, "_tf_kv_prefix", None) if isinstance(tf_prefix, dict) and tf_prefix.get("mode") == "rf_shifted_sink": prefix_k = tf_prefix.get("prefix_k") prefix_v = tf_prefix.get("prefix_v") if prefix_k is not None and prefix_v is not None: key = torch.cat([prefix_k.to(key.dtype), key], dim=1) value = torch.cat([prefix_v.to(value.dtype), value], dim=1) else: prefix_k_parts: list[torch.Tensor] = [] prefix_v_parts: list[torch.Tensor] = [] sink_k_pre = tf_prefix.get("sink_k_pre") sink_v = tf_prefix.get("sink_v") if sink_k_pre is not None and sink_v is not None and sink_k_pre.shape[1] > 0: sink_pe = tf_prefix.get("sink_pe") if sink_pe is None: raise RuntimeError("rf_shifted_sink prefix requires a sink_pe RoPE tuple.") sink_k_pre_dt = sink_k_pre.to(key.dtype) if attn.rope_type == "interleaved": sink_k = apply_interleaved_rotary_emb(sink_k_pre_dt, sink_pe) else: sink_k = apply_split_rotary_emb(sink_k_pre_dt, sink_pe) prefix_k_parts.append(sink_k) prefix_v_parts.append(sink_v.to(value.dtype)) history_k = tf_prefix.get("history_k") history_v = tf_prefix.get("history_v") if history_k is not None and history_v is not None and history_k.shape[1] > 0: prefix_k_parts.append(history_k.to(key.dtype)) prefix_v_parts.append(history_v.to(value.dtype)) if prefix_k_parts: key = torch.cat([*prefix_k_parts, key], dim=1) value = torch.cat([*prefix_v_parts, value], dim=1) query = query.unflatten(2, (attn.heads, -1)) key = key.unflatten(2, (attn.heads, -1)) value = value.unflatten(2, (attn.heads, -1)) processor = attn.processor backend = getattr(processor, "_attention_backend", None) parallel_config = getattr(processor, "_parallel_config", None) # AR mode (n_context_tokens == 0): Q from active block attends to the # injected prefix + current K/V in one SDPA call. Legacy single-shot # mode keeps the sink-self / current-cross split. if n_context_tokens <= 0 or n_context_tokens >= query.shape[1]: hidden_states = _refiner_attention( query, key, value, backend=backend, parallel_config=parallel_config, ) else: context_hidden_states = _refiner_attention( query[:, :n_context_tokens], key[:, :n_context_tokens], value[:, :n_context_tokens], backend=backend, parallel_config=parallel_config, ) current_hidden_states = _refiner_attention( query[:, n_context_tokens:], key, value, backend=backend, parallel_config=parallel_config, ) hidden_states = torch.cat([context_hidden_states, current_hidden_states], dim=1) hidden_states = hidden_states.flatten(2, 3).to(query.dtype) if gate_logits is not None: hidden_states = hidden_states.unflatten(2, (attn.heads, -1)) gates = 2.0 * torch.sigmoid(gate_logits) hidden_states = hidden_states * gates.unsqueeze(-1) hidden_states = hidden_states.flatten(2, 3) hidden_states = attn.to_out[0](hidden_states) hidden_states = attn.to_out[1](hidden_states) return hidden_states def _capture_streaming_self_attention_kv( *, attn: nn.Module, hidden_states: torch.Tensor, query_rotary_emb: tuple[torch.Tensor, torch.Tensor], ) -> None: """Capture the current layer self-attention K/V without computing attention output.""" from diffusers.models.transformers.transformer_ltx2 import apply_interleaved_rotary_emb, apply_split_rotary_emb fused_qkv = getattr(attn, "_sana_fused_qkv", None) if fused_qkv is not None: _, key, value = fused_qkv(hidden_states) else: key = attn.to_k(hidden_states) value = attn.to_v(hidden_states) key = attn.norm_k(key) if getattr(attn, "_kv_cache_capture", False): attn._cached_kv_pre = (_capture_kv_tensor(key), _capture_kv_tensor(value)) if attn.rope_type == "interleaved": key = apply_interleaved_rotary_emb(key, query_rotary_emb) elif attn.rope_type == "split": key = apply_split_rotary_emb(key, query_rotary_emb) else: raise ValueError(f"Unsupported LTX-2 RoPE type: {attn.rope_type}") if getattr(attn, "_tf_capture_kv", False): attn._cached_kv_post = (_capture_kv_tensor(key), _capture_kv_tensor(value)) def _apply_refiner_rotary( attn: nn.Module, tensor: torch.Tensor, rotary_emb: tuple[torch.Tensor, torch.Tensor], ) -> torch.Tensor: from diffusers.models.transformers.transformer_ltx2 import apply_interleaved_rotary_emb, apply_split_rotary_emb if attn.rope_type == "interleaved": return apply_interleaved_rotary_emb(tensor, rotary_emb) if attn.rope_type == "split": return apply_split_rotary_emb(tensor, rotary_emb) raise ValueError(f"Unsupported LTX-2 RoPE type: {attn.rope_type}") def _refiner_attention( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, *, backend: object, parallel_config: object, ) -> torch.Tensor: kernel = _refiner_self_attn_kernel() if kernel in {"flash_attn", "flash-attn", "fa2"}: return _flash_attn_func()(query, key, value, dropout_p=0.0, causal=False) if kernel in {"sdpa", "torch_sdpa", "pytorch_sdpa"}: hidden_states = F.scaled_dot_product_attention( query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2), attn_mask=None, dropout_p=0.0, is_causal=False, ) return hidden_states.transpose(1, 2) if kernel and kernel not in {"default", "dispatch", "diffusers", "0", "off"}: raise ValueError(f"Unsupported SANA_WM_REFINER_SELF_ATTN_KERNEL={kernel!r}.") from diffusers.models.attention_dispatch import dispatch_attention_fn return dispatch_attention_fn( query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, backend=backend, parallel_config=parallel_config, ) def _refiner_self_attn_kernel() -> str: return os.environ.get("SANA_WM_REFINER_SELF_ATTN_KERNEL", "").strip().lower() _FLASH_ATTN_FUNC = None def _flash_attn_func(): global _FLASH_ATTN_FUNC if _FLASH_ATTN_FUNC is None: from flash_attn import flash_attn_func _FLASH_ATTN_FUNC = flash_attn_func return _FLASH_ATTN_FUNC def _set_cross_attention_kv_cache( transformer: nn.Module, prompt_embeds: torch.Tensor, prompt_attention_mask: torch.Tensor | None, ) -> None: blocks = transformer.transformer_blocks if not blocks: return batch_size = int(prompt_embeds.shape[0]) hidden_dim = int(blocks[0].attn2.to_k.in_features) encoder_hidden_states = transformer.caption_projection(prompt_embeds) encoder_hidden_states = encoder_hidden_states.view(batch_size, -1, hidden_dim) encoder_attention_mask = _prepare_encoder_attention_mask(prompt_attention_mask, encoder_hidden_states.dtype) for block in blocks: attn = block.attn2 cross_hidden = encoder_hidden_states if getattr(attn, "norm_cross", False): cross_hidden = attn.norm_encoder_hidden_states(cross_hidden) key = attn.to_k(cross_hidden) value = attn.to_v(cross_hidden) if attn.norm_k is not None: key = attn.norm_k(key) inner_dim = int(key.shape[-1]) head_dim = inner_dim // int(attn.heads) key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) attn._sana_cross_attn_kv_cache = (key.detach(), value.detach(), encoder_attention_mask) def _clear_cross_attention_kv_cache(transformer: nn.Module) -> None: for block in transformer.transformer_blocks: if hasattr(block.attn2, "_sana_cross_attn_kv_cache"): block.attn2._sana_cross_attn_kv_cache = None def _has_cross_attention_kv_cache(transformer: nn.Module) -> bool: blocks = getattr(transformer, "transformer_blocks", None) if not blocks: return False return getattr(blocks[0].attn2, "_sana_cross_attn_kv_cache", None) is not None def _cross_attention_with_cached_kv( attn: nn.Module, hidden_states: torch.Tensor, cache: tuple[torch.Tensor, torch.Tensor, torch.Tensor | None], ) -> torch.Tensor: key, value, attention_mask = cache residual = hidden_states input_ndim = hidden_states.ndim spatial_norm = getattr(attn, "spatial_norm", None) if spatial_norm is not None: hidden_states = spatial_norm(hidden_states, None) if input_ndim == 4: batch_size, channel, height, width = hidden_states.shape hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) else: batch_size = int(hidden_states.shape[0]) channel = height = width = None if attention_mask is not None: source_length = int(key.shape[2]) prepare_attention_mask = getattr(attn, "prepare_attention_mask", None) if prepare_attention_mask is not None: attn_mask = prepare_attention_mask(attention_mask, source_length, batch_size) attn_mask = attn_mask.view(batch_size, attn.heads, -1, attn_mask.shape[-1]) elif attention_mask.ndim == 3: attn_mask = attention_mask[:, None, :, :] elif attention_mask.ndim == 2: attn_mask = attention_mask[:, None, None, :] else: attn_mask = attention_mask else: attn_mask = None group_norm = getattr(attn, "group_norm", None) if group_norm is not None: hidden_states = group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) query = attn.to_q(hidden_states) if attn.norm_q is not None: query = attn.norm_q(query) head_dim = int(key.shape[-1]) query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) hidden_states = F.scaled_dot_product_attention( query, key.to(query.dtype), value.to(query.dtype), attn_mask=attn_mask, dropout_p=0.0, is_causal=False, ) hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim) hidden_states = hidden_states.to(query.dtype) hidden_states = attn.to_out[0](hidden_states) hidden_states = attn.to_out[1](hidden_states) if input_ndim == 4: hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) if getattr(attn, "residual_connection", False): hidden_states = hidden_states + residual return hidden_states / float(getattr(attn, "rescale_output_factor", 1.0)) def _set_kv_prefix_on_blocks( transformer: nn.Module, kv_prefix_per_layer: list[dict[str, object]] | None, ) -> None: """Mirror tian's ``_inject_kv_prefix``: attach a per-layer prefix dict to each ``attn1``.""" blocks = transformer.transformer_blocks if kv_prefix_per_layer is None: _clear_kv_prefix_on_blocks(transformer) return if len(kv_prefix_per_layer) != len(blocks): raise RuntimeError( f"kv_prefix_per_layer has {len(kv_prefix_per_layer)} entries but transformer has {len(blocks)} blocks." ) for block, prefix in zip(blocks, kv_prefix_per_layer): block.attn1._tf_kv_prefix = prefix def _clear_kv_prefix_on_blocks(transformer: nn.Module) -> None: for block in transformer.transformer_blocks: block.attn1._tf_kv_prefix = None def _set_capture_flag_on_blocks( transformer: nn.Module, mode: str, *, enable: bool, layer_mask: list[bool] | None = None, ) -> None: """Toggle ``_kv_cache_capture`` (pre-RoPE) or ``_tf_capture_kv`` (post-RoPE) per block.""" if mode == "pre_rope": attr = "_kv_cache_capture" clear_attr = "_cached_kv_pre" elif mode == "post_rope": attr = "_tf_capture_kv" clear_attr = "_cached_kv_post" else: raise ValueError(f"capture_mode must be 'pre_rope' or 'post_rope', got {mode!r}") blocks = transformer.transformer_blocks if layer_mask is not None and len(layer_mask) != len(blocks): raise RuntimeError(f"layer_mask has {len(layer_mask)} entries but transformer has {len(blocks)} blocks.") for layer_idx, block in enumerate(blocks): enabled = bool(enable and (layer_mask is None or layer_mask[layer_idx])) setattr(block.attn1, attr, enabled) if enabled: # Clear any previous capture so the next forward writes a fresh value. if hasattr(block.attn1, clear_attr): setattr(block.attn1, clear_attr, None) def _collect_captured_kv_from_blocks( transformer: nn.Module, mode: str, layer_mask: list[bool] | None = None, ) -> list[tuple[torch.Tensor, torch.Tensor] | None]: attr = "_cached_kv_pre" if mode == "pre_rope" else "_cached_kv_post" blocks = transformer.transformer_blocks if layer_mask is not None and len(layer_mask) != len(blocks): raise RuntimeError(f"layer_mask has {len(layer_mask)} entries but transformer has {len(blocks)} blocks.") out: list[tuple[torch.Tensor, torch.Tensor] | None] = [] for layer_idx, block in enumerate(blocks): if layer_mask is not None and not layer_mask[layer_idx]: out.append(None) if hasattr(block.attn1, attr): setattr(block.attn1, attr, None) continue cached = getattr(block.attn1, attr, None) if cached is None: raise RuntimeError(f"Expected {attr!r} on attn1 after capture forward, but found None.") out.append(cached) # Release the reference so the orchestrator owns the only handle. setattr(block.attn1, attr, None) return out def _pack_text_embeds( text_hidden_states: torch.Tensor, sequence_lengths: torch.Tensor, device: str | torch.device, padding_side: str = "left", scale_factor: int = 8, eps: float = 1e-6, ) -> torch.Tensor: batch_size, seq_len, hidden_dim, _ = text_hidden_states.shape original_dtype = text_hidden_states.dtype token_indices = torch.arange(seq_len, device=device).unsqueeze(0) if padding_side == "right": mask = token_indices < sequence_lengths[:, None] elif padding_side == "left": start_indices = seq_len - sequence_lengths[:, None] mask = token_indices >= start_indices else: raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}") mask = mask[:, :, None, None] masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0) num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1) masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (num_valid_positions + eps) x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True) x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True) normalized_hidden_states = (text_hidden_states - masked_mean) / (x_max - x_min + eps) normalized_hidden_states = normalized_hidden_states * scale_factor normalized_hidden_states = normalized_hidden_states.flatten(2) mask_flat = mask.squeeze(-1).expand(-1, -1, normalized_hidden_states.shape[-1]) normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0) return normalized_hidden_states.to(dtype=original_dtype) def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: batch_size, _, num_frames, height, width = latents.shape post_patch_num_frames = num_frames // patch_size_t post_patch_height = height // patch_size post_patch_width = width // patch_size latents = latents.reshape( batch_size, -1, post_patch_num_frames, patch_size_t, post_patch_height, patch_size, post_patch_width, patch_size, ) latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) return latents def _unpack_latents( latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1, ) -> torch.Tensor: batch_size = latents.size(0) latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) return latents def _prepare_encoder_attention_mask(mask: torch.Tensor | None, dtype: torch.dtype) -> torch.Tensor | None: if mask is None: return None if mask.ndim != 2: return mask if bool(torch.all(mask)): return None return ((1 - mask.to(dtype)) * -10000.0).unsqueeze(1) def _resolve_kv_cache_storage_dtype() -> torch.dtype | None: raw = os.environ.get("SANA_WM_REFINER_KV_CACHE_DTYPE", "").strip().lower() if not raw or raw in {"bf16", "bfloat16", "none", "off", "0"}: return None if raw in {"fp8", "fp8_e4m3", "fp8_e4m3fn", "float8_e4m3fn", "e4m3"}: return torch.float8_e4m3fn if raw in {"fp8_e5m2", "float8_e5m2", "e5m2"}: return torch.float8_e5m2 raise ValueError(f"Unsupported SANA_WM_REFINER_KV_CACHE_DTYPE={raw!r}.") def _refiner_fast_kv_capture_mode() -> str: raw = os.environ.get("SANA_WM_REFINER_FAST_KV_CAPTURE", "").strip().lower() if not raw or raw in {"clean", "exact", "off", "0"}: return "clean" # Reuses K/V from the final denoise prediction. This avoids the extra # post-refine capture forward, but the cached history is approximate. if raw in {"last_predict", "reuse_last_predict", "final_predict"}: return "last_predict" # Hybrid mode: run exact post-capture only for # SANA_WM_REFINER_EXACT_CAPTURE_LAYERS, and fill the remaining history # layers from the final denoise prediction K/V. This is approximate only # for layers outside the exact-capture mask. if raw in {"fill_missing", "fill-missing", "hybrid_fill", "hybrid"}: return "fill_missing" raise ValueError(f"Unsupported SANA_WM_REFINER_FAST_KV_CAPTURE={raw!r}.") def _refiner_fast_kv_needs_clean_block(block_idx: int) -> bool: raw = os.environ.get("SANA_WM_REFINER_FAST_KV_CLEAN_INTERVAL", "").strip() if not raw: return False interval = int(raw) if interval <= 0: return False # Keep block 0 exact so the sink/first active history starts clean, then # refresh periodically to bound drift in long videos. return block_idx == 0 or ((block_idx + 1) % interval == 0) def _refiner_history_layer_mask(n_layers: int) -> list[bool]: raw_layers = os.environ.get("SANA_WM_REFINER_HISTORY_LAYERS", "").strip() if raw_layers: mask = [False] * int(n_layers) for item in raw_layers.split(","): item = item.strip() if not item: continue if item.lower() == "last": mask[-1] = True continue if "-" in item: start_raw, end_raw = item.split("-", 1) start = int(start_raw) end = int(end_raw) if start < 0: start += n_layers if end < 0: end += n_layers for idx in range(max(0, start), min(n_layers - 1, end) + 1): mask[idx] = True continue idx = int(item) if idx < 0: idx += n_layers if idx < 0 or idx >= n_layers: raise ValueError(f"SANA_WM_REFINER_HISTORY_LAYERS index {item!r} outside 0..{n_layers - 1}.") mask[idx] = True if not any(mask): raise ValueError("SANA_WM_REFINER_HISTORY_LAYERS selected no layers.") return mask stride_raw = os.environ.get("SANA_WM_REFINER_HISTORY_LAYER_STRIDE", "").strip() if not stride_raw: return [True] * int(n_layers) stride = int(stride_raw) if stride <= 1: return [True] * int(n_layers) offset = int(os.environ.get("SANA_WM_REFINER_HISTORY_LAYER_OFFSET", "0")) mask = [((idx - offset) % stride == 0) for idx in range(int(n_layers))] if _env_flag_default_true("SANA_WM_REFINER_HISTORY_KEEP_LAST"): mask[-1] = True if not any(mask): mask[-1] = True return mask def _refiner_exact_capture_layer_mask(n_layers: int, *, default_mask: list[bool]) -> list[bool]: raw_layers = os.environ.get("SANA_WM_REFINER_EXACT_CAPTURE_LAYERS", "").strip() if not raw_layers: return list(default_mask) mask = [False] * int(n_layers) for item in raw_layers.split(","): item = item.strip() if not item: continue if item.lower() == "last": mask[-1] = True continue if "-" in item: start_raw, end_raw = item.split("-", 1) start = int(start_raw) end = int(end_raw) if start < 0: start += n_layers if end < 0: end += n_layers for idx in range(max(0, start), min(n_layers - 1, end) + 1): mask[idx] = True continue idx = int(item) if idx < 0: idx += n_layers if idx < 0 or idx >= n_layers: raise ValueError(f"SANA_WM_REFINER_EXACT_CAPTURE_LAYERS index {item!r} outside 0..{n_layers - 1}.") mask[idx] = True if not any(mask): raise ValueError("SANA_WM_REFINER_EXACT_CAPTURE_LAYERS selected no layers.") return mask def _refiner_empty_cache_before_capture() -> bool: raw = os.environ.get("SANA_WM_REFINER_EMPTY_CACHE_BEFORE_CAPTURE", "1").strip().lower() return raw not in {"0", "false", "no", "off"} def _refiner_profile_enabled() -> bool: return _env_flag("SANA_WM_REFINER_PROFILE") def _refiner_layer_profile_enabled() -> bool: return _env_flag("SANA_WM_REFINER_LAYER_PROFILE") class _RefinerCudaProfiler: """Tiny env-gated CUDA event profiler for one refiner AR block.""" def __init__(self, *, enabled: bool, device: torch.device, block_idx: int) -> None: self.enabled = bool(enabled and device.type == "cuda") self.device = device self.block_idx = int(block_idx) self._events: list[tuple[str, torch.cuda.Event, torch.cuda.Event]] = [] self._block_start: torch.cuda.Event | None = None self._block_end: torch.cuda.Event | None = None if self.enabled: stream = torch.cuda.current_stream(device) self._block_start = torch.cuda.Event(enable_timing=True) self._block_end = torch.cuda.Event(enable_timing=True) self._block_start.record(stream) def section(self, name: str): if not self.enabled: return nullcontext() return _RefinerCudaProfileSection(self, name) def _record_section(self, name: str, start: torch.cuda.Event, end: torch.cuda.Event) -> None: self._events.append((name, start, end)) def finish(self) -> None: if not self.enabled: return stream = torch.cuda.current_stream(self.device) assert self._block_start is not None and self._block_end is not None self._block_end.record(stream) self._block_end.synchronize() totals_ms: dict[str, float] = {} counts: dict[str, int] = {} for name, start, end in self._events: elapsed_ms = float(start.elapsed_time(end)) totals_ms[name] = totals_ms.get(name, 0.0) + elapsed_ms counts[name] = counts.get(name, 0) + 1 block_total_ms = float(self._block_start.elapsed_time(self._block_end)) parts = [f"block_total={block_total_ms / 1000.0:.6f}s"] for name, elapsed_ms in totals_ms.items(): count_suffix = f"x{counts[name]}" if counts[name] != 1 else "" parts.append(f"{name}={elapsed_ms / 1000.0:.6f}s{count_suffix}") print(f"[refiner-profile] block={self.block_idx} " + " ".join(parts), flush=True) class _RefinerCudaProfileSection: def __init__(self, profiler: _RefinerCudaProfiler, name: str) -> None: self._profiler = profiler self._name = str(name) self._start: torch.cuda.Event | None = None self._end: torch.cuda.Event | None = None def __enter__(self): stream = torch.cuda.current_stream(self._profiler.device) self._start = torch.cuda.Event(enable_timing=True) self._end = torch.cuda.Event(enable_timing=True) self._start.record(stream) return self def __exit__(self, exc_type, exc, tb) -> bool: assert self._start is not None and self._end is not None stream = torch.cuda.current_stream(self._profiler.device) self._end.record(stream) self._profiler._record_section(self._name, self._start, self._end) return False def _current_refiner_prefix_tokens(transformer: nn.Module) -> int: blocks = getattr(transformer, "transformer_blocks", None) if not blocks: return 0 prefix = getattr(blocks[0].attn1, "_tf_kv_prefix", None) if not isinstance(prefix, dict): return 0 prefix_k = prefix.get("prefix_k") if isinstance(prefix_k, torch.Tensor): return int(prefix_k.shape[1]) total = 0 sink_k_pre = prefix.get("sink_k_pre") if isinstance(sink_k_pre, torch.Tensor): total += int(sink_k_pre.shape[1]) history_k = prefix.get("history_k") if isinstance(history_k, torch.Tensor): total += int(history_k.shape[1]) return total def _profile_section(profiler: _RefinerLayerCudaProfiler | None, name: str): if profiler is None: return nullcontext() return profiler.section(name) class _RefinerLayerCudaProfiler: """Env-gated CUDA event profiler for one transformer forward.""" def __init__(self, *, enabled: bool, device: torch.device, label: str) -> None: self.enabled = bool(enabled and device.type == "cuda") self.device = device self.label = str(label) self._events: list[tuple[str, torch.cuda.Event, torch.cuda.Event]] = [] self._start: torch.cuda.Event | None = None self._end: torch.cuda.Event | None = None if self.enabled: stream = torch.cuda.current_stream(device) self._start = torch.cuda.Event(enable_timing=True) self._end = torch.cuda.Event(enable_timing=True) self._start.record(stream) def section(self, name: str): if not self.enabled: return nullcontext() return _RefinerLayerCudaProfileSection(self, name) def _record_section(self, name: str, start: torch.cuda.Event, end: torch.cuda.Event) -> None: self._events.append((name, start, end)) def finish(self) -> None: if not self.enabled: return stream = torch.cuda.current_stream(self.device) assert self._start is not None and self._end is not None self._end.record(stream) self._end.synchronize() totals_ms: dict[str, float] = {} counts: dict[str, int] = {} for name, start, end in self._events: elapsed_ms = float(start.elapsed_time(end)) totals_ms[name] = totals_ms.get(name, 0.0) + elapsed_ms counts[name] = counts.get(name, 0) + 1 total_ms = float(self._start.elapsed_time(self._end)) parts = [f"total={total_ms / 1000.0:.6f}s"] for name, elapsed_ms in totals_ms.items(): count_suffix = f"x{counts[name]}" if counts[name] != 1 else "" parts.append(f"{name}={elapsed_ms / 1000.0:.6f}s{count_suffix}") print(f"[refiner-layer-profile] {self.label} " + " ".join(parts), flush=True) class _RefinerLayerCudaProfileSection: def __init__(self, profiler: _RefinerLayerCudaProfiler, name: str) -> None: self._profiler = profiler self._name = str(name) self._start: torch.cuda.Event | None = None self._end: torch.cuda.Event | None = None def __enter__(self): stream = torch.cuda.current_stream(self._profiler.device) self._start = torch.cuda.Event(enable_timing=True) self._end = torch.cuda.Event(enable_timing=True) self._start.record(stream) return self def __exit__(self, exc_type, exc, tb) -> bool: assert self._start is not None and self._end is not None stream = torch.cuda.current_stream(self._profiler.device) self._end.record(stream) self._profiler._record_section(self._name, self._start, self._end) return False def _store_kv_tensor(tensor: torch.Tensor, dtype: torch.dtype | None) -> torch.Tensor: if dtype is None: return tensor return tensor.to(dtype) def _store_kv_pair( pair: tuple[torch.Tensor, torch.Tensor], dtype: torch.dtype | None, ) -> tuple[torch.Tensor, torch.Tensor]: return (_store_kv_tensor(pair[0], dtype), _store_kv_tensor(pair[1], dtype)) def _capture_kv_tensor(tensor: torch.Tensor) -> torch.Tensor: captured = tensor.detach() if _env_flag("SANA_WM_REFINER_NO_CLONE_CAPTURED_KV"): return captured return captured.clone() def _env_tuple(name: str) -> tuple[str, ...]: raw = os.environ.get(name, "") return tuple(item.strip() for item in raw.split(",") if item.strip()) def _env_flag(name: str) -> bool: return os.environ.get(name, "").lower() in {"1", "true", "yes", "on"} def _env_flag_default_true(name: str) -> bool: return os.environ.get(name, "1").strip().lower() not in {"", "0", "false", "no", "off"} def _prepared_module_cache_root() -> Path | None: if os.environ.get("SANA_WM_PREPARED_MODULE_CACHE", "").strip().lower() not in {"1", "true", "yes", "on"}: return None root = os.environ.get("SANA_WM_PREPARED_MODULE_CACHE_DIR", "").strip() return Path(root).expanduser() if root else Path.home() / ".cache" / "sana_wm_prepared_modules" def _prepared_module_cache_hash(payload: dict[str, object]) -> str: blob = json.dumps(payload, sort_keys=True, default=str, separators=(",", ":")).encode("utf-8") return hashlib.sha256(blob).hexdigest()[:20] def _path_fingerprint(path: str | Path) -> dict[str, object]: raw = str(path) try: resolved = Path(raw).expanduser().resolve() except Exception: return {"path": raw} if resolved.is_dir(): markers = [] for rel in ("config.json", "diffusion_pytorch_model.safetensors", "model.safetensors"): item = resolved / rel try: stat = item.stat() except OSError: continue markers.append((rel, int(stat.st_size), int(stat.st_mtime_ns))) return {"path": str(resolved), "markers": markers} try: stat = resolved.stat() except OSError: return {"path": str(resolved)} return {"path": str(resolved), "size": int(stat.st_size), "mtime_ns": int(stat.st_mtime_ns)} def _is_local_callable_for_pickle(value: object) -> bool: if isinstance(value, types.MethodType): value = value.__func__ if not isinstance(value, types.FunctionType): return False qualname = getattr(value, "__qualname__", "") return "" in qualname or getattr(value, "__name__", "") == "" def _strip_local_callables_for_pickle(root: object) -> list[tuple[object, object, object, str]]: """Temporarily remove TE init closures that are not used after construction.""" restore: list[tuple[object, object, object, str]] = [] seen: set[int] = set() leaf_types = (str, bytes, int, float, bool, type(None), Path, torch.device, torch.dtype) def set_value(owner: object, key: object, old_value: object, new_value: object, kind: str) -> None: if kind == "dict": owner[key] = new_value elif kind == "list": owner[key] = new_value else: setattr(owner, str(key), new_value) restore.append((owner, key, old_value, kind)) def scrub_value(value: object) -> tuple[object, bool]: if _is_local_callable_for_pickle(value): return None, True if hasattr(value, "_replace") and hasattr(value, "init_fn"): updates = {} if _is_local_callable_for_pickle(getattr(value, "init_fn", None)): updates["init_fn"] = None if _is_local_callable_for_pickle(getattr(value, "get_rng_state_tracker", None)): updates["get_rng_state_tracker"] = None if updates: return value._replace(**updates), True return value, False def walk(obj: object) -> None: if isinstance(obj, leaf_types) or isinstance(obj, (torch.Tensor, nn.Parameter)): return obj_id = id(obj) if obj_id in seen: return seen.add(obj_id) if isinstance(obj, dict): for key, value in list(obj.items()): new_value, changed = scrub_value(value) if changed: set_value(obj, key, value, new_value, "dict") else: walk(value) return if isinstance(obj, list): for index, value in enumerate(list(obj)): new_value, changed = scrub_value(value) if changed: set_value(obj, index, value, new_value, "list") else: walk(value) return if isinstance(obj, tuple): return try: items = list(vars(obj).items()) except TypeError: return for key, value in items: if key.startswith("__"): continue new_value, changed = scrub_value(value) if changed: set_value(obj, key, value, new_value, "attr") elif key not in {"_parameters", "_buffers"}: walk(value) walk(root) return restore def _restore_stripped_pickle_values(restore: list[tuple[object, object, object, str]]) -> None: for owner, key, value, kind in reversed(restore): if kind == "dict": owner[key] = value elif kind == "list": owner[key] = value else: setattr(owner, str(key), value) def _te_module_name_variants(name: str) -> tuple[str, ...]: if not _env_flag("SANA_WM_TE_NVFP4_NORMALIZE_MODULE_NAMES"): return (name,) variants = {name} stripped = name while stripped.startswith("_orig_mod."): stripped = stripped[len("_orig_mod.") :] variants.add(stripped) variants.add(name.replace("._orig_mod.", ".")) variants.add(name.replace("_orig_mod.", "")) return tuple(dict.fromkeys(item for item in variants if item)) def _te_name_matches(patterns: tuple[str, ...], name: str) -> bool: return any(re.search(pattern, candidate) for pattern in patterns for candidate in _te_module_name_variants(name)) class _FusedQKVLinear(nn.Module): def __init__(self, to_q: nn.Linear, to_k: nn.Linear, to_v: nn.Linear) -> None: super().__init__() if to_q.in_features != to_k.in_features or to_q.in_features != to_v.in_features: raise ValueError("Cannot fuse QKV with mismatched input dimensions.") device = to_q.weight.device dtype = to_q.weight.dtype out_features = to_q.out_features + to_k.out_features + to_v.out_features use_bias = to_q.bias is not None or to_k.bias is not None or to_v.bias is not None self.linear = nn.Linear(to_q.in_features, out_features, bias=use_bias, device=device, dtype=dtype) self._splits = (to_q.out_features, to_k.out_features, to_v.out_features) with torch.no_grad(): self.linear.weight.copy_(torch.cat([to_q.weight, to_k.weight, to_v.weight], dim=0)) if self.linear.bias is not None: bias_parts = [] for src in (to_q, to_k, to_v): if src.bias is None: bias_parts.append(torch.zeros(src.out_features, device=device, dtype=dtype)) else: bias_parts.append(src.bias) self.linear.bias.copy_(torch.cat(bias_parts, dim=0)) def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: return self.linear(hidden_states).split(self._splits, dim=-1) def _fuse_refiner_self_qkv(transformer: nn.Module) -> int: converted = 0 for block in getattr(transformer, "transformer_blocks", ()): attn = getattr(block, "attn1", None) if attn is None or getattr(attn, "_sana_fused_qkv", None) is not None: continue to_q = getattr(attn, "to_q", None) to_k = getattr(attn, "to_k", None) to_v = getattr(attn, "to_v", None) if not all(isinstance(module, nn.Linear) for module in (to_q, to_k, to_v)): continue fused = _FusedQKVLinear(to_q, to_k, to_v) fused.train(bool(to_q.training or to_k.training or to_v.training)) attn._sana_fused_qkv = fused attn.to_q = nn.Identity() attn.to_k = nn.Identity() attn.to_v = nn.Identity() converted += 1 return converted def _replace_linear_with_te_nvfp4( module: nn.Module, *, recipe, params_dtype: torch.dtype, skip_patterns: tuple[str, ...], include_patterns: tuple[str, ...] | None = None, prefix: str = "", ) -> tuple[int, int]: import transformer_engine.pytorch as te converted = 0 skipped = 0 for name, child in list(module.named_children()): child_prefix = f"{prefix}.{name}" if prefix else name if _te_name_matches(skip_patterns, child_prefix): skipped += 1 continue if isinstance(child, nn.Linear): if include_patterns is not None and not _te_name_matches(include_patterns, child_prefix): skipped += 1 continue if child.in_features % 16 != 0 or child.out_features % 16 != 0: skipped += 1 continue use_cpu_staging = _env_flag("SANA_WM_TE_NVFP4_CPU_STAGING") child_training = child.training has_bias = child.bias is not None params_dtype_for_replacement = ( child.weight.dtype if child.weight.dtype in {torch.float16, torch.bfloat16, torch.float32} else params_dtype ) if use_cpu_staging: old_weight = child.weight.detach().to("cpu", copy=True) old_bias = child.bias.detach().to("cpu", copy=True) if child.bias is not None else None setattr(module, name, nn.Identity()) del child gc.collect() _empty_cuda_cache() else: old_weight = child.weight.detach() old_bias = child.bias.detach() if child.bias is not None else None try: ctx = te.fp8_model_init( enabled=True, recipe=recipe, preserve_high_precision_init_val=False, ) except TypeError: ctx = te.fp8_model_init(enabled=True, recipe=recipe) with ctx: replacement = te.Linear( old_weight.shape[1], old_weight.shape[0], bias=has_bias, params_dtype=params_dtype_for_replacement, device=str(torch.device("cuda", torch.cuda.current_device())), ) replacement.train(child_training) with torch.no_grad(): replacement.weight.copy_(old_weight.to(device=replacement.weight.device)) if old_bias is not None: replacement.bias.copy_(old_bias.to(device=replacement.bias.device)) if use_cpu_staging: del old_weight, old_bias _empty_cuda_cache() setattr(module, name, replacement) converted += 1 continue child_converted, child_skipped = _replace_linear_with_te_nvfp4( child, recipe=recipe, params_dtype=params_dtype, skip_patterns=skip_patterns, include_patterns=include_patterns, prefix=child_prefix, ) converted += child_converted skipped += child_skipped return converted, skipped def _offload_video_unused_audio_modules(transformer: nn.Module, device: torch.device | str) -> None: for name in ( "audio_proj_in", "audio_caption_projection", "audio_time_embed", "av_cross_attn_video_scale_shift", "av_cross_attn_audio_scale_shift", "av_cross_attn_video_a2v_gate", "av_cross_attn_audio_v2a_gate", "audio_rope", "cross_attn_rope", "cross_attn_audio_rope", "audio_norm_out", "audio_proj_out", ): child = getattr(transformer, name, None) if isinstance(child, nn.Module): child.to(device) for block in getattr(transformer, "transformer_blocks", ()): for name in ( "audio_norm1", "audio_attn1", "audio_norm2", "audio_attn2", "audio_to_video_norm", "audio_to_video_attn", "video_to_audio_norm", "video_to_audio_attn", "audio_norm3", "audio_ff", ): child = getattr(block, name, None) if isinstance(child, nn.Module): child.to(device) def _move_ltx2_video_modules_to(transformer: nn.Module, device: torch.device | str) -> None: for name in ("proj_in", "caption_projection", "time_embed", "rope", "norm_out", "proj_out"): child = getattr(transformer, name, None) if isinstance(child, nn.Module): child.to(device) _move_tensor_attr(transformer, "scale_shift_table", device) for block in getattr(transformer, "transformer_blocks", ()): _move_tensor_attr(block, "scale_shift_table", device) for name in ("norm1", "attn1", "norm2", "attn2", "norm3", "ff"): child = getattr(block, name, None) if isinstance(child, nn.Module): child.to(device) def _move_tensor_attr(module: nn.Module, name: str, device: torch.device | str) -> None: value = getattr(module, name, None) if isinstance(value, nn.Parameter): if value.device != torch.device(device): setattr(module, name, nn.Parameter(value.to(device), requires_grad=value.requires_grad)) elif isinstance(value, torch.Tensor) and value.device != torch.device(device): setattr(module, name, value.to(device)) def _empty_cuda_cache() -> None: if torch.cuda.is_available(): torch.cuda.empty_cache() gc.collect()