431 lines
18 KiB
Python
431 lines
18 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
# SPDX-FileCopyrightText: Copyright 2026 Lightricks and The HuggingFace Team. All rights reserved.
|
|
#
|
|
# Native conversion rules and decoder behavior copied and modified from
|
|
# Diffusers commit d035dcd7cc7c88e0a154609b62887d50bba9fdc2 (Apache-2.0).
|
|
|
|
"""LTX-2.5 Native-weight adapter for Diffusers' diffusion VAE decoder."""
|
|
|
|
import math
|
|
from collections.abc import Mapping
|
|
from functools import cache
|
|
from typing import Any
|
|
|
|
import torch
|
|
from diffusers.models.autoencoders import (
|
|
LTX2VideoDiffusionDecoderModel as DiffusersLTX2VideoDiffusionDecoderModel,
|
|
)
|
|
from diffusers.models.autoencoders import ltx2_diffusion_decoder as diffusers_decoder
|
|
from diffusers.models.autoencoders.ltx2_diffusion_decoder import (
|
|
LTX2VideoDiffusionDecoder3d as DiffusersLTX2VideoDiffusionDecoder3d,
|
|
)
|
|
from diffusers.models.autoencoders.vae import DecoderOutput
|
|
from safetensors import safe_open
|
|
|
|
from ..ops import is_ltx2_fna_eligible, is_ltx2_fusion_eligible, resolve_ltx2_vae_operators
|
|
|
|
LTX25_NATIVE_DIFFUSION_DECODER_REPO_ID = "Lightricks/LTX-2.5"
|
|
LTX25_NATIVE_ARTIFACT_REVISION = "8a4ff96f581e72bedc1b44367581c49d544a05f1"
|
|
LTX25_NATIVE_DIFFUSION_DECODER_FILENAME = "vae/ltx-2.5-video-vae-bf16.safetensors"
|
|
|
|
_NATIVE_DECODER_PREFIXES = ("vae.decoder.", "decoder.")
|
|
_NATIVE_STATISTICS_KEYS = {
|
|
"per_channel_statistics.mean-of-means": "latents_mean",
|
|
"per_channel_statistics.std-of-means": "latents_std",
|
|
}
|
|
_NATIVE_KEY_REPLACEMENTS = (
|
|
("t_embedder.mlp.0.", "t_embedder.timestep_embedder.linear_1."),
|
|
("t_embedder.mlp.2.", "t_embedder.timestep_embedder.linear_2."),
|
|
(".attn.proj.", ".attn.to_out.0."),
|
|
(".attn.q_norm.", ".attn.norm_q."),
|
|
(".attn.k_norm.", ".attn.norm_k."),
|
|
)
|
|
_GATE_FOLD_TARGETS = {
|
|
".attn.to_out.0.weight": ".gate_msa",
|
|
".attn.to_out.0.bias": ".gate_msa",
|
|
".mlp.w_down.weight": ".gate_mlp",
|
|
".context_proj.weight": ".gate_ctx",
|
|
".context_proj.bias": ".gate_ctx",
|
|
}
|
|
_GATE_SUFFIXES = tuple(_GATE_FOLD_TARGETS.values())
|
|
|
|
|
|
def _strip_native_decoder_prefix(key: str) -> str | None:
|
|
for prefix in _NATIVE_DECODER_PREFIXES:
|
|
if key.startswith(prefix):
|
|
return key.removeprefix(prefix)
|
|
return None
|
|
|
|
|
|
def _native_statistics_target(key: str) -> str | None:
|
|
return _NATIVE_STATISTICS_KEYS.get(key.removeprefix("vae."))
|
|
|
|
|
|
def _rename_native_decoder_key(key: str) -> str:
|
|
for source, target in _NATIVE_KEY_REPLACEMENTS:
|
|
key = key.replace(source, target)
|
|
return key
|
|
|
|
|
|
def _fold_native_gate(key: str, value: torch.Tensor, gates: Mapping[str, torch.Tensor]) -> torch.Tensor:
|
|
target_suffix = next((suffix for suffix in _GATE_FOLD_TARGETS if key.endswith(suffix)), None)
|
|
if target_suffix is None:
|
|
return value
|
|
gate_key = key[: -len(target_suffix)] + _GATE_FOLD_TARGETS[target_suffix]
|
|
gate = gates.get(gate_key)
|
|
if gate is None:
|
|
return value
|
|
gate = gate.to(device=value.device, dtype=torch.float32)
|
|
value_float = value.to(dtype=torch.float32)
|
|
folded = gate.unsqueeze(1) * value_float if value.ndim == 2 else gate * value_float
|
|
return folded.to(dtype=value.dtype)
|
|
|
|
|
|
def convert_ltx25_native_diffusion_decoder_state_dict(
|
|
native_state_dict: Mapping[str, torch.Tensor],
|
|
) -> dict[str, torch.Tensor]:
|
|
"""Convert the canonical Native LTX-2.5 DiffVAE state dict to Diffusers."""
|
|
decoder_state_dict: dict[str, torch.Tensor] = {}
|
|
converted: dict[str, torch.Tensor] = {}
|
|
for key, value in native_state_dict.items():
|
|
statistics_target = _native_statistics_target(key)
|
|
if statistics_target is not None:
|
|
converted[statistics_target] = value
|
|
continue
|
|
decoder_key = _strip_native_decoder_prefix(key)
|
|
if decoder_key is not None:
|
|
decoder_state_dict[decoder_key] = value
|
|
|
|
gates = {key: value for key, value in decoder_state_dict.items() if key.endswith(_GATE_SUFFIXES)}
|
|
for key, value in decoder_state_dict.items():
|
|
# ``type_emb`` is an unused shipping-checkpoint residue. Preview heads
|
|
# and legacy static gates are not parameters of the runtime decoder.
|
|
if key == "type_emb" or key.startswith("coarse_") or ".coarse_" in key or key in gates:
|
|
continue
|
|
|
|
converted_key = _rename_native_decoder_key(key)
|
|
value = _fold_native_gate(converted_key, value, gates)
|
|
if converted_key.endswith((".qkv.weight", ".qkv.bias")):
|
|
leaf = "weight" if converted_key.endswith(".weight") else "bias"
|
|
prefix = converted_key[: -len(f"qkv.{leaf}")]
|
|
if value.shape[0] % 3 != 0:
|
|
raise ValueError(
|
|
f"Fused LTX-2.5 DiffVAE parameter {key!r} has leading dimension "
|
|
f"{value.shape[0]}, which is not divisible by 3."
|
|
)
|
|
chunk = value.shape[0] // 3
|
|
converted[f"decoder.{prefix}to_q.{leaf}"] = value[:chunk].clone()
|
|
converted[f"decoder.{prefix}to_k.{leaf}"] = value[chunk : 2 * chunk].clone()
|
|
converted[f"decoder.{prefix}to_v.{leaf}"] = value[2 * chunk :].clone()
|
|
continue
|
|
|
|
converted[f"decoder.{converted_key}"] = value
|
|
|
|
return converted
|
|
|
|
|
|
def load_ltx25_native_diffusion_decoder_state_dict(path: str) -> dict[str, torch.Tensor]:
|
|
"""Read and convert only DiffVAE tensors from the canonical Native VAE file."""
|
|
with safe_open(path, framework="pt", device="cpu") as handle:
|
|
native_state_dict = {
|
|
key: handle.get_tensor(key)
|
|
for key in handle.keys()
|
|
if _strip_native_decoder_prefix(key) is not None or _native_statistics_target(key) is not None
|
|
}
|
|
return convert_ltx25_native_diffusion_decoder_state_dict(native_state_dict)
|
|
|
|
|
|
@cache
|
|
def _load_natten_token_permutation():
|
|
from kernels import get_kernel
|
|
|
|
natten = get_kernel("shi-labs/natten", version=1)
|
|
return natten.token_permute.token_permute_operation, natten.token_permute.token_unpermute_operation
|
|
|
|
|
|
class LTX2VideoVaeNeighborhoodNattenProcessor(diffusers_decoder.LTX2VideoVaeNeighborhoodNattenProcessor):
|
|
"""NATTEN dispatch with Omni FNA for eligible stage-5 grids."""
|
|
|
|
def __call__(
|
|
self, attn: "LTX2VideoVaeNeighborhoodAttention", hidden_states: torch.Tensor, block_mask=None
|
|
) -> torch.Tensor:
|
|
batch_size, num_frames, height, width, channels = hidden_states.shape
|
|
query, key, value = attn.project_qkv(hidden_states)
|
|
# NATTEN's CUTLASS kernels silently produce wrong output for non-contiguous inputs.
|
|
query, key, value = query.contiguous(), key.contiguous(), value.contiguous()
|
|
operators = resolve_ltx2_vae_operators(hidden_states.device) if is_ltx2_fna_eligible(hidden_states) else None
|
|
# `scale=1.0`: the query is already scaled in `project_qkv`, as in the reference.
|
|
if (
|
|
tuple(attn.kernel_size) == (11, 11, 11)
|
|
and hidden_states.dtype == torch.bfloat16
|
|
and channels == 256
|
|
and attn.head_dim == 64
|
|
and operators is not None
|
|
and operators.fna is not None
|
|
# The native metadata encodes KV tile indices in 15 bits.
|
|
and math.prod((n + tile - 1) // tile for n, tile in zip((num_frames, height, width), (4, 4, 8))) <= 1 << 15
|
|
):
|
|
token_permute, token_unpermute = _load_natten_token_permutation()
|
|
query_permuted, token_shape, _ = token_permute(query, (4, 4, 4), flip_tiled_dims=True)
|
|
key_permuted, _, _ = token_permute(key, (4, 4, 8), flip_tiled_dims=True)
|
|
value_permuted, _, _ = token_permute(value, (4, 4, 8), flip_tiled_dims=True)
|
|
output = operators.fna(
|
|
query_permuted,
|
|
key_permuted,
|
|
value_permuted,
|
|
shape=(num_frames, height, width),
|
|
window=tuple(attn.kernel_size),
|
|
)
|
|
hidden_states = token_unpermute(
|
|
output,
|
|
token_layout_shape=token_shape,
|
|
tile_shape=(4, 4, 4),
|
|
flip_tiled_dims=True,
|
|
)
|
|
else:
|
|
hidden_states = self._na3d(
|
|
query,
|
|
key,
|
|
value,
|
|
kernel_size=attn.kernel_size,
|
|
scale=1.0,
|
|
backend=self.backend,
|
|
)
|
|
hidden_states = hidden_states.reshape(batch_size, num_frames, height, width, channels)
|
|
return attn.to_out[0](hidden_states)
|
|
|
|
|
|
class LTX2VideoVaeNeighborhoodAttention(diffusers_decoder.LTX2VideoVaeNeighborhoodAttention):
|
|
_available_processors = [
|
|
*diffusers_decoder.LTX2VideoVaeNeighborhoodAttention._available_processors,
|
|
LTX2VideoVaeNeighborhoodNattenProcessor,
|
|
]
|
|
|
|
def project_qkv(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
|
batch_size, num_frames, height, width, _ = hidden_states.shape
|
|
shape = (batch_size, num_frames, height, width, self.heads, self.head_dim)
|
|
query = self.to_q(hidden_states).view(shape)
|
|
key = self.to_k(hidden_states).view(shape)
|
|
value = self.to_v(hidden_states).view(shape)
|
|
|
|
operators = resolve_ltx2_vae_operators(query.device) if is_ltx2_fusion_eligible(query) else None
|
|
optimized = (
|
|
None
|
|
if operators is None
|
|
else operators.qk_norm_rope(
|
|
query,
|
|
key,
|
|
self.norm_q.weight,
|
|
self.norm_k.weight,
|
|
self.norm_q.eps,
|
|
self.scale,
|
|
self.rope.rope_dim_split,
|
|
self.rope.base,
|
|
)
|
|
)
|
|
if optimized is not None:
|
|
query, key = optimized
|
|
return query, key, value
|
|
|
|
query = self.norm_q(query)
|
|
key = self.norm_k(key)
|
|
query = query * self.scale
|
|
return self.rope(query), self.rope(key), value
|
|
|
|
|
|
class LTX2VideoVaeSwiGLU(diffusers_decoder.LTX2VideoVaeSwiGLU):
|
|
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
|
operators = resolve_ltx2_vae_operators(hidden_states.device) if is_ltx2_fusion_eligible(hidden_states) else None
|
|
optimized = (
|
|
None
|
|
if operators is None
|
|
else operators.swiglu(hidden_states, self.w_gate.weight, self.w_up.weight, self.w_down.weight)
|
|
)
|
|
return super().forward(hidden_states) if optimized is None else optimized
|
|
|
|
|
|
class LTX2VideoVaeDiffusionNABlock(diffusers_decoder.LTX2VideoVaeDiffusionNABlock):
|
|
def forward(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
latent_context: torch.Tensor,
|
|
modulation: tuple[torch.Tensor, ...],
|
|
block_mask=None,
|
|
) -> torch.Tensor:
|
|
if hidden_states.dtype is not torch.bfloat16 or not is_ltx2_fusion_eligible(hidden_states):
|
|
return super().forward(hidden_states, latent_context, modulation, block_mask)
|
|
operators = resolve_ltx2_vae_operators(hidden_states.device)
|
|
assert operators is not None
|
|
|
|
scale_msa, shift_msa, _, scale_mlp, shift_mlp, _, _ = [
|
|
modulation[i] + self.scale_shift_table[i].view(1, 1, 1, 1, -1) for i in range(self.num_mod_params)
|
|
]
|
|
context_output = self.context_proj(latent_context)
|
|
attention_input = operators.residual_norm(
|
|
hidden_states,
|
|
context_output,
|
|
None,
|
|
self.norm1.weight,
|
|
scale_msa,
|
|
shift_msa,
|
|
self.norm1.eps,
|
|
)
|
|
if attention_input is None:
|
|
hidden_states = hidden_states + context_output
|
|
hidden_states = hidden_states + self.attn(
|
|
self.norm1(hidden_states) * (1 + scale_msa) + shift_msa, block_mask
|
|
)
|
|
hidden_states = hidden_states + self.mlp(self.norm2(hidden_states) * (1 + scale_mlp) + shift_mlp)
|
|
return hidden_states
|
|
|
|
attention_output = self.attn(attention_input, block_mask)
|
|
mlp_input = operators.residual_norm(
|
|
hidden_states,
|
|
context_output,
|
|
attention_output,
|
|
self.norm2.weight,
|
|
scale_mlp,
|
|
shift_mlp,
|
|
self.norm2.eps,
|
|
)
|
|
if mlp_input is None:
|
|
residual = hidden_states + context_output
|
|
residual = residual + attention_output
|
|
mlp_input = self.norm2(residual) * (1 + scale_mlp) + shift_mlp
|
|
mlp_output = self.mlp(mlp_input)
|
|
output = operators.residual_add(
|
|
hidden_states,
|
|
context_output,
|
|
attention_output,
|
|
mlp_output,
|
|
)
|
|
if output is not None:
|
|
return output
|
|
residual = hidden_states + context_output
|
|
residual = residual + attention_output
|
|
return residual + mlp_output
|
|
|
|
|
|
class LTX2VideoDiffusionDecoder3d(DiffusersLTX2VideoDiffusionDecoder3d):
|
|
"""Diffusers decoder core with the short-clip NATTEN context fix."""
|
|
|
|
def forward_stage_4(
|
|
self,
|
|
hidden_states: torch.Tensor,
|
|
drop_leading_frame: bool = True,
|
|
crop_trailing_ghost: bool = True,
|
|
) -> torch.Tensor:
|
|
blocks = self.det_stages[-1]
|
|
block_mask = blocks[0].attn.build_block_mask(hidden_states)
|
|
for block in blocks:
|
|
hidden_states = block(hidden_states, block_mask)
|
|
hidden_states = self.upsamples[-1](hidden_states, drop_leading_frame=drop_leading_frame)
|
|
|
|
num_pad = self.trailing_pad_latent_frames
|
|
if crop_trailing_ghost and num_pad > 0:
|
|
content_frames = max(hidden_states.shape[1] - num_pad * self.temporal_compression_ratio, 1)
|
|
keep_frames = min(hidden_states.shape[1], max(content_frames, self.stage5_kernel[0]))
|
|
hidden_states = hidden_states[:, :keep_frames]
|
|
return hidden_states
|
|
|
|
|
|
class LTX2VideoDiffusionDecoderModel(DiffusersLTX2VideoDiffusionDecoderModel):
|
|
"""Diffusers 0.40 decoder plus canonical Native loading and short-clip support."""
|
|
|
|
def __init__(
|
|
self,
|
|
out_channels: int = 3,
|
|
latent_channels: int = 128,
|
|
patch_size: int = 4,
|
|
scaling_factor: float = 1.0,
|
|
decoder_head_dim: int = 64,
|
|
decoder_stage_channels: tuple[int, ...] = (2048, 1024, 512, 512, 256),
|
|
decoder_stage_depths: tuple[int, ...] = (4, 6, 4, 2, 8),
|
|
decoder_stage_kernels: tuple[tuple[int, int, int], ...] = ((3, 7, 7), (3, 7, 7), (3, 5, 5), (3, 5, 5)),
|
|
decoder_upsample_strides: tuple[tuple[int, int, int], ...] = ((1, 2, 2), (2, 1, 1), (2, 2, 2), (2, 2, 2)),
|
|
decoder_upsample_channel_reductions: tuple[int, ...] = (2, 2, 1, 2),
|
|
decoder_stage5_kernel: tuple[int, int, int] = (11, 11, 11),
|
|
decoder_t_emb_dim: int = 384,
|
|
decoder_timestep_scale_multiplier: float = 1000.0,
|
|
decoder_model_output_type: str = "x0",
|
|
decoder_num_inference_steps: int = 1,
|
|
spatial_compression_ratio: int = 32,
|
|
temporal_compression_ratio: int = 8,
|
|
) -> None:
|
|
super().__init__(
|
|
out_channels=out_channels,
|
|
latent_channels=latent_channels,
|
|
patch_size=patch_size,
|
|
scaling_factor=scaling_factor,
|
|
decoder_head_dim=decoder_head_dim,
|
|
decoder_stage_channels=decoder_stage_channels,
|
|
decoder_stage_depths=decoder_stage_depths,
|
|
decoder_stage_kernels=decoder_stage_kernels,
|
|
decoder_upsample_strides=decoder_upsample_strides,
|
|
decoder_upsample_channel_reductions=decoder_upsample_channel_reductions,
|
|
decoder_stage5_kernel=decoder_stage5_kernel,
|
|
decoder_t_emb_dim=decoder_t_emb_dim,
|
|
decoder_timestep_scale_multiplier=decoder_timestep_scale_multiplier,
|
|
decoder_model_output_type=decoder_model_output_type,
|
|
decoder_num_inference_steps=decoder_num_inference_steps,
|
|
spatial_compression_ratio=spatial_compression_ratio,
|
|
temporal_compression_ratio=temporal_compression_ratio,
|
|
)
|
|
# Diffusers constructs its private decoder core internally. Switching
|
|
# to this behavior-only subclass preserves parameters and state-dict keys.
|
|
self.decoder.__class__ = LTX2VideoDiffusionDecoder3d
|
|
self.decoder.stage5_kernel = tuple(decoder_stage5_kernel)
|
|
overrides = {
|
|
diffusers_decoder.LTX2VideoVaeNeighborhoodAttention: LTX2VideoVaeNeighborhoodAttention,
|
|
diffusers_decoder.LTX2VideoVaeSwiGLU: LTX2VideoVaeSwiGLU,
|
|
diffusers_decoder.LTX2VideoVaeDiffusionNABlock: LTX2VideoVaeDiffusionNABlock,
|
|
}
|
|
# Preserve every module and parameter object; only execution methods change.
|
|
for module in self.decoder.modules():
|
|
if replacement := overrides.get(type(module)):
|
|
module.__class__ = replacement
|
|
|
|
@classmethod
|
|
def from_ltx25_native_checkpoint(
|
|
cls,
|
|
checkpoint_path: str,
|
|
config: Mapping[str, Any],
|
|
dtype: torch.dtype,
|
|
) -> "LTX2VideoDiffusionDecoderModel":
|
|
"""Construct this decoder class directly from the canonical Native checkpoint."""
|
|
with torch.device("meta"):
|
|
model = cls.from_config(dict(config))
|
|
state_dict = load_ltx25_native_diffusion_decoder_state_dict(checkpoint_path)
|
|
try:
|
|
model.load_state_dict(state_dict, strict=True, assign=True)
|
|
except RuntimeError as exc:
|
|
raise ValueError(f"Invalid LTX-2.5 Native DiffVAE checkpoint {checkpoint_path!r}.") from exc
|
|
model.to(device="cpu", dtype=dtype)
|
|
return model
|
|
|
|
def decode(
|
|
self,
|
|
z: torch.Tensor,
|
|
generator: torch.Generator | None = None,
|
|
num_inference_steps: int | None = None,
|
|
return_dict: bool = True,
|
|
) -> DecoderOutput | tuple[torch.Tensor]:
|
|
decoded = super().decode(
|
|
z,
|
|
generator=generator,
|
|
num_inference_steps=num_inference_steps,
|
|
return_dict=False,
|
|
)[0]
|
|
|
|
# Short clips retain replicated trailing context through stage 5. Do
|
|
# not expose those context pixels to the caller.
|
|
target_num_frames = (z.shape[2] - 1) * self.temporal_compression_ratio + 1
|
|
target_height = z.shape[3] * self.spatial_compression_ratio
|
|
target_width = z.shape[4] * self.spatial_compression_ratio
|
|
decoded = decoded[:, :, :target_num_frames, :target_height, :target_width]
|
|
|
|
if not return_dict:
|
|
return (decoded,)
|
|
return DecoderOutput(sample=decoded)
|