wehub-resource-sync
2631 lines
111 KiB
Python
2631 lines
111 KiB
Python
# 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 "<locals>" in qualname or getattr(value, "__name__", "") == "<lambda>"
|
||
|
||
|
||
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()
|