749 lines
30 KiB
Python
749 lines
30 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
"""Shared component construction helpers for the LTX model family."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import inspect
|
|
import json
|
|
import logging
|
|
import os
|
|
from dataclasses import dataclass, replace
|
|
from pathlib import Path
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import torch
|
|
from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler
|
|
from diffusers.pipelines.ltx2 import LTX2TextConnectors
|
|
from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
|
|
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
|
|
from diffusers.video_processor import VideoProcessor
|
|
from huggingface_hub import hf_hub_download
|
|
from safetensors import safe_open
|
|
from safetensors.torch import load_file
|
|
from transformers import AutoModelForImageTextToText, AutoTokenizer, Gemma3ForConditionalGeneration
|
|
|
|
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
|
|
from vllm_omni.diffusion.attention.layer import Attention as OmniAttention
|
|
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import DistributedAutoencoderKLLTX2Video
|
|
from vllm_omni.diffusion.distributed.utils import get_local_device
|
|
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
|
from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders
|
|
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
|
|
|
|
from .ltx2_request import LTXCheckpointKind, validate_ltx_checkpoint
|
|
from .ltx2_transformer import (
|
|
LTX2VideoTransformer3DModel,
|
|
apply_interleaved_rotary_emb,
|
|
apply_split_rotary_emb,
|
|
to_ltx_padding_mask,
|
|
)
|
|
|
|
try:
|
|
from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE
|
|
except ImportError:
|
|
LTX2VocoderWithBWE = None
|
|
|
|
try:
|
|
from transformers import Gemma4UnifiedForConditionalGeneration as _Gemma4UnifiedForConditionalGeneration
|
|
except ImportError:
|
|
_Gemma4UnifiedForConditionalGeneration = None
|
|
|
|
|
|
_LTX25_TEXT_ENCODER_CLS = AutoModelForImageTextToText if _Gemma4UnifiedForConditionalGeneration is not None else None
|
|
|
|
|
|
_LTX_COMPONENT_SUBFOLDERS = (
|
|
"tokenizer",
|
|
"text_encoder",
|
|
"connectors",
|
|
"vae",
|
|
"audio_vae",
|
|
"vocoder",
|
|
"scheduler",
|
|
"latent_upsampler",
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LTXComponentProfile:
|
|
"""Component construction and discovery contract for an LTX variant."""
|
|
|
|
name: str
|
|
dit_modules: tuple[str, ...]
|
|
encoder_modules: tuple[str, ...]
|
|
vae_modules: tuple[str, ...]
|
|
resident_modules: tuple[str, ...] = ()
|
|
video_vae_cls: type = AutoencoderKLLTX2Video
|
|
vocoder_cls: type = LTX2Vocoder
|
|
text_encoder_cls: type | None = Gemma3ForConditionalGeneration
|
|
vocoder_fallback_cls: type | None = None
|
|
artifact_repo_id: str | None = None
|
|
latent_upsampler_filename: str | None = None
|
|
distilled_lora_filename: str | None = None
|
|
transformer_subfolder: str = "transformer"
|
|
scheduler_use_dynamic_shifting: bool = False
|
|
scheduler_shift_terminal: float | None = None
|
|
preserve_connector_attention_mask: bool = False
|
|
|
|
|
|
LTX2_COMPONENT_PROFILE = LTXComponentProfile(
|
|
name="ltx2",
|
|
dit_modules=("transformer",),
|
|
encoder_modules=("text_encoder", "connectors"),
|
|
vae_modules=("vae", "audio_vae"),
|
|
resident_modules=("vocoder",),
|
|
video_vae_cls=DistributedAutoencoderKLLTX2Video,
|
|
)
|
|
|
|
LTX23_COMPONENT_PROFILE = LTXComponentProfile(
|
|
name="ltx2_3",
|
|
dit_modules=("transformer",),
|
|
encoder_modules=("text_encoder", "connectors"),
|
|
vae_modules=("vae", "audio_vae"),
|
|
resident_modules=("vocoder",),
|
|
video_vae_cls=DistributedAutoencoderKLLTX2Video,
|
|
vocoder_cls=LTX2VocoderWithBWE or LTX2Vocoder,
|
|
vocoder_fallback_cls=LTX2Vocoder,
|
|
)
|
|
|
|
LTX25_FULL_COMPONENT_PROFILE = LTXComponentProfile(
|
|
name="ltx2_5_full",
|
|
dit_modules=("transformer",),
|
|
encoder_modules=("text_encoder", "connectors"),
|
|
vae_modules=("vae", "audio_vae"),
|
|
resident_modules=("vocoder",),
|
|
video_vae_cls=DistributedAutoencoderKLLTX2Video,
|
|
vocoder_cls=LTX2VocoderWithBWE or LTX2Vocoder,
|
|
vocoder_fallback_cls=LTX2Vocoder,
|
|
text_encoder_cls=_LTX25_TEXT_ENCODER_CLS,
|
|
transformer_subfolder="transformer_full",
|
|
scheduler_use_dynamic_shifting=True,
|
|
scheduler_shift_terminal=0.1,
|
|
preserve_connector_attention_mask=True,
|
|
)
|
|
|
|
|
|
LTX2_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile(
|
|
name="ltx2_distilled",
|
|
dit_modules=("transformer",),
|
|
encoder_modules=("text_encoder", "connectors"),
|
|
vae_modules=("vae", "audio_vae"),
|
|
resident_modules=("vocoder", "latent_upsampler"),
|
|
video_vae_cls=DistributedAutoencoderKLLTX2Video,
|
|
)
|
|
|
|
LTX25_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile(
|
|
name="ltx2_5_distilled",
|
|
dit_modules=("transformer",),
|
|
encoder_modules=("text_encoder", "connectors"),
|
|
vae_modules=("vae", "audio_vae"),
|
|
resident_modules=("vocoder", "latent_upsampler"),
|
|
video_vae_cls=DistributedAutoencoderKLLTX2Video,
|
|
vocoder_cls=LTX2VocoderWithBWE or LTX2Vocoder,
|
|
vocoder_fallback_cls=LTX2Vocoder,
|
|
text_encoder_cls=_LTX25_TEXT_ENCODER_CLS,
|
|
preserve_connector_attention_mask=True,
|
|
)
|
|
|
|
LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE = replace(
|
|
LTX2_DISTILLED_COMPONENT_PROFILE,
|
|
name="ltx2_distilled_one_stage",
|
|
resident_modules=LTX2_COMPONENT_PROFILE.resident_modules,
|
|
)
|
|
|
|
LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE = replace(
|
|
LTX25_DISTILLED_COMPONENT_PROFILE,
|
|
name="ltx2_5_distilled_one_stage",
|
|
resident_modules=LTX25_FULL_COMPONENT_PROFILE.resident_modules,
|
|
)
|
|
|
|
|
|
LTX23_DISTILLED_COMPONENT_PROFILE = replace(
|
|
LTX23_COMPONENT_PROFILE,
|
|
name="ltx2_3_distilled",
|
|
resident_modules=(*LTX23_COMPONENT_PROFILE.resident_modules, "latent_upsampler"),
|
|
artifact_repo_id="Lightricks/LTX-2.3",
|
|
latent_upsampler_filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
|
|
)
|
|
|
|
LTX23_DISTILLED_ONE_STAGE_COMPONENT_PROFILE = replace(
|
|
LTX23_DISTILLED_COMPONENT_PROFILE,
|
|
name="ltx2_3_distilled_one_stage",
|
|
resident_modules=LTX23_COMPONENT_PROFILE.resident_modules,
|
|
)
|
|
|
|
LTX2_TWO_STAGE_COMPONENT_PROFILE = replace(
|
|
LTX2_COMPONENT_PROFILE,
|
|
name="ltx2_two_stage",
|
|
resident_modules=(*LTX2_COMPONENT_PROFILE.resident_modules, "latent_upsampler"),
|
|
artifact_repo_id="Lightricks/LTX-2",
|
|
latent_upsampler_filename="ltx-2-spatial-upscaler-x2-1.0.safetensors",
|
|
distilled_lora_filename="ltx-2-19b-distilled-lora-384.safetensors",
|
|
)
|
|
|
|
LTX23_TWO_STAGE_COMPONENT_PROFILE = replace(
|
|
LTX23_COMPONENT_PROFILE,
|
|
name="ltx2_3_two_stage",
|
|
resident_modules=(*LTX23_COMPONENT_PROFILE.resident_modules, "latent_upsampler"),
|
|
artifact_repo_id="Lightricks/LTX-2.3",
|
|
latent_upsampler_filename="ltx-2.3-spatial-upscaler-x2-1.1.safetensors",
|
|
distilled_lora_filename="ltx-2.3-22b-distilled-lora-384-1.1.safetensors",
|
|
)
|
|
|
|
|
|
LTX25_TWO_STAGE_COMPONENT_PROFILE = replace(
|
|
LTX25_FULL_COMPONENT_PROFILE,
|
|
name="ltx2_5_two_stage",
|
|
resident_modules=(*LTX25_FULL_COMPONENT_PROFILE.resident_modules, "latent_upsampler"),
|
|
artifact_repo_id="Lightricks/LTX-2.5",
|
|
latent_upsampler_filename=("latent_upscale_models/ltx-2.5-latent-spatial-upscaler-x2-bf16-1.0.safetensors"),
|
|
distilled_lora_filename="loras/ltx-2.5-22b-distilled-lora-450-bf16.safetensors",
|
|
)
|
|
|
|
_COMPONENT_PROFILES: dict[tuple[str, str], LTXComponentProfile] = {
|
|
("one_stage", "2"): LTX2_COMPONENT_PROFILE,
|
|
("one_stage", "2.3"): LTX23_COMPONENT_PROFILE,
|
|
("one_stage", "2.5"): LTX25_FULL_COMPONENT_PROFILE,
|
|
("two_stage", "2"): LTX2_TWO_STAGE_COMPONENT_PROFILE,
|
|
("two_stage", "2.3"): LTX23_TWO_STAGE_COMPONENT_PROFILE,
|
|
("two_stage", "2.5"): LTX25_TWO_STAGE_COMPONENT_PROFILE,
|
|
("distilled_one_stage", "2"): LTX2_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
|
|
("distilled_one_stage", "2.3"): LTX23_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
|
|
("distilled_one_stage", "2.5"): LTX25_DISTILLED_ONE_STAGE_COMPONENT_PROFILE,
|
|
("distilled_two_stage", "2"): LTX2_DISTILLED_COMPONENT_PROFILE,
|
|
("distilled_two_stage", "2.3"): LTX23_DISTILLED_COMPONENT_PROFILE,
|
|
("distilled_two_stage", "2.5"): LTX25_DISTILLED_COMPONENT_PROFILE,
|
|
("dmd2", "2"): LTX2_COMPONENT_PROFILE,
|
|
("dmd2", "2.3"): LTX23_COMPONENT_PROFILE,
|
|
}
|
|
|
|
|
|
def resolve_ltx_checkpoint_kind(pipeline_kind: str) -> LTXCheckpointKind | None:
|
|
"""Derive checkpoint requirements from the execution contract."""
|
|
if pipeline_kind in {"one_stage", "two_stage"}:
|
|
return "regular"
|
|
if pipeline_kind in {"distilled_one_stage", "distilled_two_stage"}:
|
|
return "distilled"
|
|
if pipeline_kind == "dmd2":
|
|
return None
|
|
raise ValueError(f"Unsupported LTX pipeline kind: {pipeline_kind!r}.")
|
|
|
|
|
|
def resolve_ltx_artifact(
|
|
model: str,
|
|
repo_id: str,
|
|
filename: str,
|
|
) -> str:
|
|
"""Resolve an official LTX sidecar from the model root or its Hub repository."""
|
|
candidate = Path(model) / filename
|
|
if candidate.is_file():
|
|
return str(candidate)
|
|
|
|
try:
|
|
return hf_hub_download(repo_id=repo_id, filename=filename)
|
|
except Exception as exc:
|
|
raise FileNotFoundError(
|
|
f"Unable to resolve LTX artifact {filename!r}. Searched {candidate}; "
|
|
f"place the file in the model root or make {repo_id} available."
|
|
) from exc
|
|
|
|
|
|
def _load_ltx_latent_upsampler_single_file(path: str, dtype: torch.dtype) -> LTX2LatentUpsamplerModel:
|
|
"""Load an official single-file upsampler into the Diffusers module."""
|
|
with safe_open(path, framework="pt", device="cpu") as handle:
|
|
metadata = handle.metadata() or {}
|
|
raw_config = json.loads(metadata.get("config", "{}"))
|
|
config = {
|
|
"in_channels": raw_config.get("in_channels", 128),
|
|
"mid_channels": raw_config.get("mid_channels", 1024),
|
|
"num_blocks_per_stage": raw_config.get("num_blocks_per_stage", 4),
|
|
"dims": raw_config.get("dims", 3),
|
|
"spatial_upsample": raw_config.get("spatial_upsample", True),
|
|
"temporal_upsample": raw_config.get("temporal_upsample", False),
|
|
"rational_spatial_scale": raw_config.get("spatial_scale", 2.0),
|
|
"use_rational_resampler": raw_config.get("rational_resampler", False),
|
|
}
|
|
with torch.device("cpu"):
|
|
upsampler = LTX2LatentUpsamplerModel(**config).to(dtype=dtype)
|
|
|
|
state_dict = load_file(path, device="cpu")
|
|
if "upsampler.0.weight" in state_dict and hasattr(upsampler.upsampler, "conv"):
|
|
state_dict["upsampler.conv.weight"] = state_dict.pop("upsampler.0.weight")
|
|
state_dict["upsampler.conv.bias"] = state_dict.pop("upsampler.0.bias")
|
|
missing, unexpected = upsampler.load_state_dict(state_dict, strict=False)
|
|
unresolved_missing = set(missing) - {"upsampler.blur_down.kernel"}
|
|
if unresolved_missing or unexpected:
|
|
raise ValueError(
|
|
f"Invalid LTX latent upsampler {path}: missing={sorted(unresolved_missing)}, "
|
|
f"unexpected={sorted(unexpected)}."
|
|
)
|
|
return upsampler
|
|
|
|
|
|
def resolve_ltx_component_profile(pipeline_kind: str, model_version: str) -> LTXComponentProfile:
|
|
"""Resolve component construction independently from execution recipes."""
|
|
try:
|
|
return _COMPONENT_PROFILES[(pipeline_kind, model_version)]
|
|
except KeyError as exc:
|
|
raise ValueError(f"Unsupported LTX component kind/version: {pipeline_kind!r}/{model_version!r}.") from exc
|
|
|
|
|
|
def _load_ltx_metadata_json(model: str, filename: str, revision: str | None = None) -> dict[str, Any]:
|
|
"""Load small checkpoint metadata without relying on repository names."""
|
|
if os.path.isdir(model):
|
|
path = os.path.join(model, filename)
|
|
if not os.path.isfile(path):
|
|
return {}
|
|
else:
|
|
try:
|
|
path = hf_hub_download(repo_id=model, filename=filename, revision=revision)
|
|
except Exception:
|
|
return {}
|
|
try:
|
|
with open(path) as config_file:
|
|
value = json.load(config_file)
|
|
except (OSError, TypeError, ValueError):
|
|
return {}
|
|
return value if isinstance(value, dict) else {}
|
|
|
|
|
|
def detect_ltx_model_version(model: str, revision: str | None = None) -> str:
|
|
"""Detect the LTX model version from checkpoint component metadata."""
|
|
model_index = _load_ltx_metadata_json(model, "model_index.json", revision)
|
|
model_version = str(model_index.get("model_version", ""))
|
|
if model_version.startswith("2.5"):
|
|
return "2.5"
|
|
|
|
text_encoder_entry = model_index.get("text_encoder")
|
|
if isinstance(text_encoder_entry, (list, tuple)) and text_encoder_entry:
|
|
text_encoder_class = str(text_encoder_entry[-1])
|
|
elif isinstance(text_encoder_entry, dict):
|
|
text_encoder_class = str(text_encoder_entry.get("_class_name", ""))
|
|
else:
|
|
text_encoder_class = ""
|
|
if text_encoder_class == "Gemma4UnifiedForConditionalGeneration":
|
|
return "2.5"
|
|
|
|
text_encoder_config = _load_ltx_metadata_json(model, "text_encoder/config.json", revision)
|
|
if text_encoder_config.get("model_type") in ("gemma4_unified", "gemma4"):
|
|
return "2.5"
|
|
|
|
if model_version.startswith("2.3"):
|
|
return "2.3"
|
|
|
|
transformer_config = _load_ltx_metadata_json(model, "transformer/config.json", revision)
|
|
if transformer_config.get("ff_bias") is False:
|
|
return "2.5"
|
|
|
|
vocoder_entry = model_index.get("vocoder")
|
|
if isinstance(vocoder_entry, (list, tuple)) and vocoder_entry:
|
|
vocoder_class = str(vocoder_entry[-1])
|
|
elif isinstance(vocoder_entry, dict):
|
|
vocoder_class = str(vocoder_entry.get("_class_name", ""))
|
|
else:
|
|
vocoder_class = ""
|
|
if vocoder_class == "LTX2VocoderWithBWE":
|
|
return "2.3"
|
|
|
|
vocoder_config = _load_ltx_metadata_json(model, "vocoder/config.json", revision)
|
|
if str(vocoder_config.get("model_version", "")).startswith("2.3"):
|
|
return "2.3"
|
|
if vocoder_config.get("_class_name") == "LTX2VocoderWithBWE":
|
|
return "2.3"
|
|
if "bwe" in vocoder_config or "bwe_config" in vocoder_config:
|
|
return "2.3"
|
|
return "2"
|
|
|
|
|
|
def preserves_reference_image_size(*, model: str | None, revision: str | None = None) -> bool:
|
|
"""Preserve source geometry only for the LTX-2.5 CRF-18 I2V path."""
|
|
return model is not None and detect_ltx_model_version(model, revision=revision) == "2.5"
|
|
|
|
|
|
class _LTXConnectorAttnProcessor:
|
|
"""Preserve official connector math around Omni attention dispatch."""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
has_learned_registers: bool = False,
|
|
preserve_learned_register_mask: bool = False,
|
|
) -> None:
|
|
self.has_learned_registers = has_learned_registers
|
|
self.preserve_learned_register_mask = preserve_learned_register_mask
|
|
|
|
def __call__(
|
|
self,
|
|
attn: Any,
|
|
hidden_states: torch.Tensor,
|
|
encoder_hidden_states: torch.Tensor | None = None,
|
|
attention_mask: torch.Tensor | None = None,
|
|
query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
|
|
) -> torch.Tensor:
|
|
encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
|
|
gate_logits = attn.to_gate_logits(hidden_states) if attn.to_gate_logits is not None else None
|
|
|
|
query = attn.to_q(hidden_states)
|
|
key = attn.to_k(encoder_hidden_states)
|
|
value = attn.to_v(encoder_hidden_states)
|
|
# Offload hooks may execute affine-free Q/K norms in FP32. Restore the
|
|
# projection dtype before attention, matching the fully resident path.
|
|
query = attn.norm_q(query).to(dtype=value.dtype)
|
|
key = attn.norm_k(key).to(dtype=value.dtype)
|
|
|
|
if query_rotary_emb is not None:
|
|
# Diffusers builds connector RoPE in FP32, while the official
|
|
# connector materializes it in the hidden-state dtype.
|
|
query_rotary_emb = tuple(component.to(value.dtype) for component in query_rotary_emb)
|
|
key_rotary_emb = key_rotary_emb if key_rotary_emb is not None else query_rotary_emb
|
|
key_rotary_emb = tuple(component.to(value.dtype) for component in key_rotary_emb)
|
|
if attn.rope_type == "interleaved":
|
|
query = apply_interleaved_rotary_emb(query, query_rotary_emb)
|
|
key = apply_interleaved_rotary_emb(key, key_rotary_emb)
|
|
elif attn.rope_type == "split":
|
|
query = apply_split_rotary_emb(query, query_rotary_emb, head_dim=attn.head_dim)
|
|
key = apply_split_rotary_emb(key, key_rotary_emb, head_dim=attn.head_dim)
|
|
else:
|
|
raise ValueError(f"Unsupported LTX connector RoPE type: {attn.rope_type}")
|
|
|
|
# Keep Q/K in the projection dtype expected by the attention backend.
|
|
query = query.to(dtype=value.dtype)
|
|
key = key.to(dtype=value.dtype)
|
|
|
|
batch_size, _, inner_dim = query.shape
|
|
head_dim = inner_dim // attn.heads
|
|
kv_heads = attn.inner_kv_dim // attn.head_dim
|
|
query = query.view(batch_size, -1, attn.heads, head_dim)
|
|
key = key.view(batch_size, -1, kv_heads, head_dim)
|
|
value = value.view(batch_size, -1, kv_heads, head_dim)
|
|
|
|
# LTX-2/2.3 replace padding tokens with learned registers, so their
|
|
# old padding mask becomes a no-op. LTX-2.5 preserves that all-zero
|
|
# additive mask to match the official masked-SDPA dispatch.
|
|
if self.has_learned_registers and not self.preserve_learned_register_mask:
|
|
attention_mask = None
|
|
elif attention_mask is not None and attn.omni_attention.attn_backend.get_name().upper() == "FLASH_ATTN":
|
|
attention_mask = to_ltx_padding_mask(attention_mask)
|
|
attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
|
|
hidden_states = attn.omni_attention(query, key, value, attn_metadata)
|
|
hidden_states = hidden_states.reshape(batch_size, -1, inner_dim)
|
|
|
|
if gate_logits is not None:
|
|
hidden_states = hidden_states.unflatten(2, (attn.heads, -1))
|
|
hidden_states = hidden_states * (2.0 * torch.sigmoid(gate_logits)).unsqueeze(-1)
|
|
hidden_states = hidden_states.flatten(2, 3)
|
|
|
|
hidden_states = attn.to_out[0](hidden_states)
|
|
return attn.to_out[1](hidden_states)
|
|
|
|
|
|
def _install_connector_attention(
|
|
connectors: LTX2TextConnectors,
|
|
*,
|
|
preserve_learned_register_mask: bool = False,
|
|
) -> None:
|
|
for connector_name in ("video_connector", "audio_connector"):
|
|
connector = getattr(connectors, connector_name, None)
|
|
has_learned_registers = getattr(connector, "learnable_registers", None) is not None
|
|
for block_index, block in enumerate(getattr(connector, "transformer_blocks", ())):
|
|
attention = getattr(block, "attn1", None)
|
|
if attention is not None:
|
|
attention.omni_attention = OmniAttention(
|
|
num_heads=attention.heads,
|
|
head_size=attention.head_dim,
|
|
num_kv_heads=attention.inner_kv_dim // attention.head_dim,
|
|
softmax_scale=1.0 / (attention.head_dim**0.5),
|
|
causal=False,
|
|
prefix=f"connectors.{connector_name}.transformer_blocks.{block_index}.attn1",
|
|
role="ltx2.connector",
|
|
role_category="self",
|
|
skip_sequence_parallel=True,
|
|
disable_kv_quant=True,
|
|
)
|
|
attention.set_processor(
|
|
_LTXConnectorAttnProcessor(
|
|
has_learned_registers=has_learned_registers,
|
|
preserve_learned_register_mask=preserve_learned_register_mask,
|
|
)
|
|
)
|
|
|
|
|
|
def _detect_vocoder_output_sample_rate(model: str, revision: str | None = None) -> int | None:
|
|
"""Read the generated waveform sample rate from the vocoder config."""
|
|
vocoder_config_path = os.path.join(model, "vocoder", "config.json")
|
|
if not os.path.exists(vocoder_config_path):
|
|
try:
|
|
vocoder_config_path = hf_hub_download(model, "vocoder/config.json", revision=revision)
|
|
except Exception:
|
|
return None
|
|
try:
|
|
with open(vocoder_config_path) as config_file:
|
|
return json.load(config_file).get("output_sampling_rate")
|
|
except Exception:
|
|
return None
|
|
|
|
|
|
def get_ltx2_post_process_func(od_config: Any):
|
|
"""Build the common LTX engine-output adapter."""
|
|
output_sample_rate = _detect_vocoder_output_sample_rate(
|
|
od_config.model,
|
|
revision=getattr(od_config, "revision", None),
|
|
)
|
|
|
|
def post_process_func(output: tuple[torch.Tensor, torch.Tensor] | torch.Tensor):
|
|
if not (isinstance(output, tuple) and len(output) == 2):
|
|
return output
|
|
video, audio = output
|
|
if isinstance(audio, torch.Tensor):
|
|
audio = audio.detach().cpu()
|
|
result: dict[str, Any] = {"video": video, "audio": audio}
|
|
if output_sample_rate is not None:
|
|
result["audio_sample_rate"] = output_sample_rate
|
|
return result
|
|
|
|
return post_process_func
|
|
|
|
|
|
def _load_component(
|
|
component_cls: type,
|
|
model: str,
|
|
subfolder: str,
|
|
*,
|
|
local_files_only: bool,
|
|
dtype: torch.dtype,
|
|
revision: str | None,
|
|
) -> Any:
|
|
return from_pretrained_with_prefetch(
|
|
component_cls.from_pretrained,
|
|
model,
|
|
subfolder=subfolder,
|
|
prefetch_list=_LTX_COMPONENT_SUBFOLDERS,
|
|
local_files_only=local_files_only,
|
|
revision=revision,
|
|
torch_dtype=dtype,
|
|
)
|
|
|
|
|
|
def _place_aux_components(pipeline: Any) -> None:
|
|
parallel_config = getattr(pipeline.od_config, "parallel_config", None)
|
|
use_managed_placement = bool(
|
|
getattr(pipeline.od_config, "enable_cpu_offload", False)
|
|
or getattr(pipeline.od_config, "enable_layerwise_offload", False)
|
|
or getattr(parallel_config, "use_hsdp", False)
|
|
)
|
|
if use_managed_placement:
|
|
return
|
|
|
|
modules = ModuleDiscovery.discover(pipeline)
|
|
for module in (*modules.encoders, *modules.vaes, *modules.resident_modules):
|
|
module.to(pipeline.device)
|
|
|
|
|
|
def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
|
|
"""Build the common LTX component graph selected by ``component_profile``."""
|
|
profile: LTXComponentProfile = pipeline.component_profile
|
|
pipeline.od_config = od_config
|
|
pipeline.device = get_local_device()
|
|
dtype = getattr(od_config, "dtype", torch.bfloat16)
|
|
model = od_config.model
|
|
revision = getattr(od_config, "revision", None)
|
|
local_files_only = os.path.exists(model)
|
|
|
|
pipeline.weights_sources = [
|
|
DiffusersPipelineLoader.ComponentSource(
|
|
model_or_path=model,
|
|
subfolder=profile.transformer_subfolder,
|
|
revision=revision,
|
|
prefix="transformer.",
|
|
fall_back_to_pt=True,
|
|
),
|
|
]
|
|
prefetch_subfolders(model, _LTX_COMPONENT_SUBFOLDERS, local_files_only=local_files_only, revision=revision)
|
|
|
|
pipeline.tokenizer = AutoTokenizer.from_pretrained(
|
|
model,
|
|
subfolder="tokenizer",
|
|
local_files_only=local_files_only,
|
|
revision=revision,
|
|
)
|
|
if profile.text_encoder_cls is None:
|
|
raise ImportError("LTX-2.5 requires Gemma4UnifiedForConditionalGeneration; install transformers>=5.10.1,<5.15.")
|
|
with torch.device("cpu"):
|
|
pipeline.text_encoder = _load_component(
|
|
profile.text_encoder_cls,
|
|
model,
|
|
"text_encoder",
|
|
local_files_only=local_files_only,
|
|
dtype=dtype,
|
|
revision=revision,
|
|
)
|
|
pipeline.connectors = _load_component(
|
|
LTX2TextConnectors,
|
|
model,
|
|
"connectors",
|
|
local_files_only=local_files_only,
|
|
dtype=dtype,
|
|
revision=revision,
|
|
)
|
|
_install_connector_attention(
|
|
pipeline.connectors,
|
|
preserve_learned_register_mask=profile.preserve_connector_attention_mask,
|
|
)
|
|
pipeline.vae = _load_component(
|
|
profile.video_vae_cls,
|
|
model,
|
|
"vae",
|
|
local_files_only=local_files_only,
|
|
dtype=dtype,
|
|
revision=revision,
|
|
)
|
|
pipeline.audio_vae = _load_component(
|
|
AutoencoderKLLTX2Audio,
|
|
model,
|
|
"audio_vae",
|
|
local_files_only=local_files_only,
|
|
dtype=dtype,
|
|
revision=revision,
|
|
)
|
|
try:
|
|
pipeline.vocoder = _load_component(
|
|
profile.vocoder_cls,
|
|
model,
|
|
"vocoder",
|
|
local_files_only=local_files_only,
|
|
dtype=dtype,
|
|
revision=revision,
|
|
)
|
|
except (TypeError, OSError, ValueError):
|
|
if profile.vocoder_fallback_cls is None or profile.vocoder_fallback_cls is profile.vocoder_cls:
|
|
raise
|
|
pipeline.vocoder = _load_component(
|
|
profile.vocoder_fallback_cls,
|
|
model,
|
|
"vocoder",
|
|
local_files_only=local_files_only,
|
|
dtype=dtype,
|
|
revision=revision,
|
|
)
|
|
|
|
if "latent_upsampler" in profile.resident_modules:
|
|
upsampler_config = os.path.join(model, "latent_upsampler", "config.json")
|
|
if os.path.isfile(upsampler_config) or not local_files_only:
|
|
try:
|
|
pipeline.latent_upsampler = _load_component(
|
|
LTX2LatentUpsamplerModel,
|
|
model,
|
|
"latent_upsampler",
|
|
local_files_only=local_files_only,
|
|
dtype=dtype,
|
|
revision=revision,
|
|
)
|
|
except (OSError, ValueError):
|
|
if profile.latent_upsampler_filename is None or profile.artifact_repo_id is None:
|
|
raise
|
|
upsampler_path = resolve_ltx_artifact(
|
|
model,
|
|
profile.artifact_repo_id,
|
|
profile.latent_upsampler_filename,
|
|
)
|
|
pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype)
|
|
else:
|
|
if profile.latent_upsampler_filename is None or profile.artifact_repo_id is None:
|
|
raise FileNotFoundError(f"LTX latent upsampler component not found under {model}.")
|
|
upsampler_path = resolve_ltx_artifact(
|
|
model,
|
|
profile.artifact_repo_id,
|
|
profile.latent_upsampler_filename,
|
|
)
|
|
pipeline.latent_upsampler = _load_ltx_latent_upsampler_single_file(upsampler_path, dtype)
|
|
|
|
transformer_config = load_transformer_config(
|
|
model, profile.transformer_subfolder, local_files_only, revision=revision
|
|
)
|
|
quant_config = getattr(od_config, "quantization_config", None)
|
|
pipeline.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config)
|
|
_place_aux_components(pipeline)
|
|
pipeline.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
|
|
model,
|
|
subfolder="scheduler",
|
|
local_files_only=local_files_only,
|
|
revision=revision,
|
|
)
|
|
if profile.scheduler_use_dynamic_shifting:
|
|
pipeline.scheduler = FlowMatchEulerDiscreteScheduler.from_config(
|
|
pipeline.scheduler.config,
|
|
use_dynamic_shifting=True,
|
|
shift_terminal=profile.scheduler_shift_terminal,
|
|
)
|
|
validate_ltx_checkpoint(
|
|
pipeline.scheduler.config,
|
|
expected_kind=resolve_ltx_checkpoint_kind(pipeline.pipeline_kind),
|
|
pipeline_name=type(pipeline).__name__,
|
|
)
|
|
|
|
pipeline.vae_spatial_compression_ratio = pipeline.vae.spatial_compression_ratio
|
|
pipeline.vae_temporal_compression_ratio = pipeline.vae.temporal_compression_ratio
|
|
pipeline.audio_vae_mel_compression_ratio = pipeline.audio_vae.mel_compression_ratio
|
|
pipeline.audio_vae_temporal_compression_ratio = pipeline.audio_vae.temporal_compression_ratio
|
|
pipeline.transformer_spatial_patch_size = pipeline.transformer.config.patch_size
|
|
pipeline.transformer_temporal_patch_size = pipeline.transformer.config.patch_size_t
|
|
pipeline.audio_sampling_rate = pipeline.audio_vae.config.sample_rate
|
|
pipeline.audio_hop_length = pipeline.audio_vae.config.mel_hop_length
|
|
pipeline.video_processor = VideoProcessor(vae_scale_factor=pipeline.vae_spatial_compression_ratio)
|
|
|
|
tokenizer_max_length = pipeline.tokenizer.model_max_length
|
|
if tokenizer_max_length is None or tokenizer_max_length > 100000:
|
|
encoder_config = getattr(pipeline.text_encoder, "config", None)
|
|
tokenizer_max_length = getattr(encoder_config, "max_position_embeddings", None)
|
|
if tokenizer_max_length is None:
|
|
tokenizer_max_length = getattr(encoder_config, "max_seq_len", None)
|
|
pipeline.tokenizer_max_length = int(tokenizer_max_length or 1024)
|
|
|
|
pipeline._interrupt = False
|
|
|
|
|
|
def load_transformer_config(
|
|
model_path: str,
|
|
subfolder: str = "transformer",
|
|
local_files_only: bool = True,
|
|
*,
|
|
revision: str | None = None,
|
|
) -> dict:
|
|
"""Load an LTX transformer config from a local model or the HF Hub."""
|
|
if local_files_only:
|
|
config_path = os.path.join(model_path, subfolder, "config.json")
|
|
if not os.path.exists(config_path):
|
|
raise FileNotFoundError(f"LTX transformer config not found: {config_path}")
|
|
else:
|
|
config_path = hf_hub_download(
|
|
repo_id=model_path,
|
|
filename=f"{subfolder}/config.json",
|
|
revision=revision,
|
|
)
|
|
with open(config_path) as config_file:
|
|
return json.load(config_file)
|
|
|
|
|
|
def create_transformer_from_config(
|
|
config: dict,
|
|
quant_config: QuantizationConfig | None = None,
|
|
) -> LTX2VideoTransformer3DModel:
|
|
"""Construct the shared LTX transformer from a Diffusers config."""
|
|
if not config and quant_config is None:
|
|
return LTX2VideoTransformer3DModel()
|
|
|
|
signature = inspect.signature(LTX2VideoTransformer3DModel.__init__)
|
|
allowed_keys = set(signature.parameters)
|
|
kwargs = {key: value for key, value in config.items() if key in allowed_keys}
|
|
if quant_config is not None:
|
|
kwargs["quant_config"] = quant_config
|
|
|
|
return LTX2VideoTransformer3DModel(**kwargs)
|