unsloth/studio/backend/core/inference/video_ltx2.py

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

"""LTX-2.3 pipeline assembly for diffusers 0.39.

diffusers 0.39 ships every LTX-2.3 model class but its single-file loader maps every LTX-2
checkpoint to the 2.0 config, so 2.3 checkpoints fail a shape check at load. The community
transformer-only GGUFs also carry the DiT + connectors but NOT the text projections, VAEs, or
vocoder that 2.3 moved out of the transformer. This assembles the full 2.3 pipeline:

- transformer: from the checkpoint via ``from_single_file`` with the 2.3 config overrides and the
  ``prompt_adaln_single`` keys pre-renamed (the library converter doesn't know them).
- connectors: from the checkpoint's connector keys plus the ``text_embedding_projection`` tensors,
  fetched from the companion file in ``unsloth/LTX-2.3-GGUF`` when not bundled.
- video/audio VAE, vocoder: from the checkpoint when bundled, else the companion files.
- scheduler, text encoder (Gemma3), tokenizer: from the LTX-2.0 base repo, which 2.3 shares.

Every config and rename table mirrors diffusers' ``scripts/convert_ltx2_to_diffusers.py`` (the
authoritative 2.3 mapping the loader hasn't absorbed). Assembled through the constructor, not
``from_pretrained``, because the vocoder class differs from the base pin (``LTX2VocoderWithBWE`` vs
``LTX2Vocoder``) and the type gate would reject it.
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Optional

from loggers import get_logger

logger = get_logger(__name__)

# Companion files (text projections, VAEs incl. vocoder) next to the quants in unsloth's GGUF
# repo: the official Lightricks weights split out of the combined checkpoint. Keyed by variant.
LTX23_EXTRAS_REPO = "unsloth/LTX-2.3-GGUF"
_EXTRAS_TEXT_PROJ = "text_encoders/ltx-2.3-22b-{variant}_embeddings_connectors.safetensors"
_EXTRAS_VIDEO_VAE = "vae/ltx-2.3-22b-{variant}_video_vae.safetensors"
_EXTRAS_AUDIO_VAE = "vae/ltx-2.3-22b-{variant}_audio_vae.safetensors"

# ── configs + rename tables, verbatim from scripts/convert_ltx2_to_diffusers.py ──

# from_single_file config overrides on top of the base 2.0 transformer config.
LTX_2_3_TRANSFORMER_CONFIG_OVERRIDES: dict[str, Any] = {
    "gated_attn": True,
    "cross_attn_mod": True,
    "audio_gated_attn": True,
    "audio_cross_attn_mod": True,
    "use_prompt_embeddings": False,
    "perturbed_attn": True,
}

# Keys the 2.0-era converter doesn't know; renamed before from_single_file. Audio prefix first.
_TRANSFORMER_PRERENAME = (
    ("audio_prompt_adaln_single.", "audio_prompt_adaln."),
    ("prompt_adaln_single.", "prompt_adaln."),
)

_CONNECTOR_KEY_PREFIXES = (
    "video_embeddings_connector",
    "audio_embeddings_connector",
    "transformer_1d_blocks",
    "text_embedding_projection",
    "connectors.",
    "video_connector",
    "audio_connector",
    "text_proj_in",
)

_CONNECTORS_RENAME = {
    "connectors.": "",
    "video_embeddings_connector": "video_connector",
    "audio_embeddings_connector": "audio_connector",
    "transformer_1d_blocks": "transformer_blocks",
    "text_embedding_projection.audio_aggregate_embed": "audio_text_proj_in",
    "text_embedding_projection.video_aggregate_embed": "video_text_proj_in",
    "q_norm": "norm_q",
    "k_norm": "norm_k",
}

_CONNECTORS_CONFIG: dict[str, Any] = {
    "caption_channels": 3840,
    "text_proj_in_factor": 49,
    "video_connector_num_attention_heads": 32,
    "video_connector_attention_head_dim": 128,
    "video_connector_num_layers": 8,
    "video_connector_num_learnable_registers": 128,
    "video_gated_attn": True,
    "audio_connector_num_attention_heads": 32,
    "audio_connector_attention_head_dim": 64,
    "audio_connector_num_layers": 8,
    "audio_connector_num_learnable_registers": 128,
    "audio_gated_attn": True,
    "connector_rope_base_seq_len": 4096,
    "rope_theta": 10000.0,
    "rope_double_precision": True,
    "causal_temporal_positioning": False,
    "rope_type": "split",
    "per_modality_projections": True,
    "video_hidden_dim": 4096,
    "audio_hidden_dim": 2048,
    "proj_bias": True,
}

_VIDEO_VAE_RENAME = {
    # Encoder
    "down_blocks.0": "down_blocks.0",
    "down_blocks.1": "down_blocks.0.downsamplers.0",
    "down_blocks.2": "down_blocks.1",
    "down_blocks.3": "down_blocks.1.downsamplers.0",
    "down_blocks.4": "down_blocks.2",
    "down_blocks.5": "down_blocks.2.downsamplers.0",
    "down_blocks.6": "down_blocks.3",
    "down_blocks.7": "down_blocks.3.downsamplers.0",
    "down_blocks.8": "mid_block",
    # Decoder (2.3 adds up_blocks.7/8: a 4th decoder stage)
    "up_blocks.0": "mid_block",
    "up_blocks.1": "up_blocks.0.upsamplers.0",
    "up_blocks.2": "up_blocks.0",
    "up_blocks.3": "up_blocks.1.upsamplers.0",
    "up_blocks.4": "up_blocks.1",
    "up_blocks.5": "up_blocks.2.upsamplers.0",
    "up_blocks.6": "up_blocks.2",
    "up_blocks.7": "up_blocks.3.upsamplers.0",
    "up_blocks.8": "up_blocks.3",
    "last_time_embedder": "time_embedder",
    "last_scale_shift_table": "scale_shift_table",
    # Common
    "res_blocks": "resnets",
    "per_channel_statistics.mean-of-means": "latents_mean",
    "per_channel_statistics.std-of-means": "latents_std",
}

_VIDEO_VAE_REMOVE_SUFFIXES = (
    "per_channel_statistics.channel",
    "per_channel_statistics.mean-of-stds",
)

_VIDEO_VAE_CONFIG: dict[str, Any] = {
    "in_channels": 3,
    "out_channels": 3,
    "latent_channels": 128,
    "block_out_channels": (256, 512, 1024, 1024),
    "down_block_types": (
        "LTX2VideoDownBlock3D",
        "LTX2VideoDownBlock3D",
        "LTX2VideoDownBlock3D",
        "LTX2VideoDownBlock3D",
    ),
    "decoder_block_out_channels": (256, 512, 512, 1024),
    "layers_per_block": (4, 6, 4, 2, 2),
    "decoder_layers_per_block": (4, 6, 4, 2, 2),
    "spatio_temporal_scaling": (True, True, True, True),
    "decoder_spatio_temporal_scaling": (True, True, True, True),
    "decoder_inject_noise": (False, False, False, False, False),
    "downsample_type": ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
    "upsample_type": ("spatiotemporal", "spatiotemporal", "temporal", "spatial"),
    "upsample_residual": (False, False, False, False),
    "upsample_factor": (2, 2, 1, 2),
    "timestep_conditioning": False,
    "patch_size": 4,
    "patch_size_t": 1,
    "resnet_norm_eps": 1e-6,
    "encoder_causal": True,
    "decoder_causal": False,
    "encoder_spatial_padding_mode": "zeros",
    "decoder_spatial_padding_mode": "zeros",
    "spatial_compression_ratio": 32,
    "temporal_compression_ratio": 8,
}

_AUDIO_VAE_RENAME = {
    "per_channel_statistics.mean-of-means": "latents_mean",
    "per_channel_statistics.std-of-means": "latents_std",
}

# Same config as LTX-2.0 (upstream's comment); the weights are still 2.3-specific.
_AUDIO_VAE_CONFIG: dict[str, Any] = {
    "base_channels": 128,
    "output_channels": 2,
    "ch_mult": (1, 2, 4),
    "num_res_blocks": 2,
    "attn_resolutions": None,
    "in_channels": 2,
    "resolution": 256,
    "latent_channels": 8,
    "norm_type": "pixel",
    "causality_axis": "height",
    "dropout": 0.0,
    "mid_block_add_attention": False,
    "sample_rate": 16000,
    "mel_hop_length": 160,
    "is_causal": True,
    "mel_bins": 64,
    "double_z": True,
}

_VOCODER_RENAME = {
    "resblocks": "resnets",
    "conv_pre": "conv_in",
    "conv_post": "conv_out",
    "act_post": "act_out",
    "downsample.lowpass": "downsample",
}

_VOCODER_CONFIG: dict[str, Any] = {
    "in_channels": 128,
    "hidden_channels": 1536,
    "out_channels": 2,
    "upsample_kernel_sizes": [11, 4, 4, 4, 4, 4],
    "upsample_factors": [5, 2, 2, 2, 2, 2],
    "resnet_kernel_sizes": [3, 7, 11],
    "resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
    "act_fn": "snakebeta",
    "leaky_relu_negative_slope": 0.1,
    "antialias": True,
    "antialias_ratio": 2,
    "antialias_kernel_size": 12,
    "final_act_fn": None,
    "final_bias": False,
    "bwe_in_channels": 128,
    "bwe_hidden_channels": 512,
    "bwe_out_channels": 2,
    "bwe_upsample_kernel_sizes": [12, 11, 4, 4, 4],
    "bwe_upsample_factors": [6, 5, 2, 2, 2],
    "bwe_resnet_kernel_sizes": [3, 7, 11],
    "bwe_resnet_dilations": [[1, 3, 5], [1, 3, 5], [1, 3, 5]],
    "bwe_act_fn": "snakebeta",
    "bwe_leaky_relu_negative_slope": 0.1,
    "bwe_antialias": True,
    "bwe_antialias_ratio": 2,
    "bwe_antialias_kernel_size": 12,
    "bwe_final_act_fn": None,
    "bwe_final_bias": False,
    "filter_length": 512,
    "hop_length": 80,
    "window_length": 512,
    "num_mel_channels": 64,
    "input_sampling_rate": 16000,
    "output_sampling_rate": 48000,
}

_DIT_PREFIX = "model.diffusion_model."


# ── checkpoint inspection ────────────────────────────────────────────────────


def read_checkpoint_header(checkpoint_path: Path | str) -> dict[str, tuple[int, ...]]:
    """Tensor name -> shape from the checkpoint HEADER only (no weight data). GGUF shapes come back
    in GGML (reversed) order, so callers should membership-test, not assume a dimension position."""
    names_shapes: dict[str, tuple[int, ...]] = {}
    path = str(checkpoint_path)
    if path.lower().endswith(".gguf"):
        from gguf import GGUFReader
        for tensor in GGUFReader(path).tensors:
            names_shapes[str(tensor.name)] = tuple(int(x) for x in tensor.shape)
    else:
        from safetensors import safe_open
        with safe_open(path, framework = "pt") as handle:
            for name in handle.keys():
                names_shapes[name] = tuple(handle.get_slice(name).get_shape())
    return names_shapes


def is_ltx23_checkpoint(checkpoint_path: Path | str) -> bool:
    """True when the checkpoint carries the 9-row LTX-2.3 modulation tables (2.0 has 6-row
    per-block scale/shift tables; 2.3 widens them to 9). An unreadable header returns False so the
    caller falls back to the stock 2.0 path."""
    try:
        header = read_checkpoint_header(checkpoint_path)
    except Exception as exc:  # noqa: BLE001
        logger.warning("video.ltx2_header_probe_failed: %s", exc)
        return False
    for name, shape in header.items():
        if name.endswith("transformer_blocks.0.scale_shift_table"):
            return 9 in shape
    return False


# ── state-dict plumbing ──────────────────────────────────────────────────────


def _apply_rename(state: dict[str, Any], rename: dict[str, str]) -> dict[str, Any]:
    out: dict[str, Any] = {}
    for key, value in state.items():
        new_key = key
        for old, new in rename.items():
            new_key = new_key.replace(old, new)
        out[new_key] = value
    return out


def _to_plain_dtype(state: dict[str, Any], torch_dtype: Any) -> dict[str, Any]:
    """Materialise every tensor as a plain torch tensor in torch_dtype. GGUF tensors arrive as
    block-quantized GGUFParameter; the small non-DiT components run dense, so dequantize here."""
    import torch

    try:
        from diffusers.quantizers.gguf.utils import GGUFParameter, dequantize_gguf_tensor
    except Exception:  # noqa: BLE001 -- gguf support not installed; plain tensors only
        GGUFParameter, dequantize_gguf_tensor = (), None

    out: dict[str, Any] = {}
    for key, value in state.items():
        if dequantize_gguf_tensor is not None and isinstance(value, GGUFParameter):
            value = dequantize_gguf_tensor(value)
        out[key] = value.to(torch_dtype) if isinstance(value, torch.Tensor) else value
    return out


def _split_checkpoint(state: dict[str, Any]) -> dict[str, dict[str, Any]]:
    """Partition a combined LTX checkpoint into per-component state dicts. Handles both layouts: the
    official combined single file (``vae.*`` / ``audio_vae.*`` / ``vocoder.*`` / DiT + projections)
    and transformer-only GGUFs (bare DiT + connector keys)."""
    groups: dict[str, dict[str, Any]] = {
        "dit": {},
        "connectors": {},
        "vae": {},
        "audio_vae": {},
        "vocoder": {},
    }
    for key, value in state.items():
        bare = key[len(_DIT_PREFIX) :] if key.startswith(_DIT_PREFIX) else key
        if bare.startswith("vae."):
            groups["vae"][bare[len("vae.") :]] = value
        elif bare.startswith("audio_vae."):
            groups["audio_vae"][bare[len("audio_vae.") :]] = value
        elif bare.startswith("vocoder."):
            groups["vocoder"][bare[len("vocoder.") :]] = value
        elif bare.startswith(_CONNECTOR_KEY_PREFIXES):
            groups["connectors"][bare] = value
        else:
            groups["dit"][bare] = value
    return groups


def _load_extras_file(filename: str, hf_token: Optional[str]) -> dict[str, Any]:
    from safetensors.torch import load_file

    from utils.hf_xet_fallback import hf_hub_download_with_xet_fallback

    path = hf_hub_download_with_xet_fallback(LTX23_EXTRAS_REPO, filename, hf_token)
    return load_file(path)


def checkpoint_variant(checkpoint_path: Path | str) -> str:
    """Which companion-weight set a checkpoint pairs with ("dev"/"distilled"). The distilled-1.1
    refresh only retrained the DiT, so it shares the distilled companions."""
    return "dev" if "dev" in Path(checkpoint_path).name.lower() else "distilled"


def ltx23_extras_files(checkpoint_path: Path | str) -> tuple[str, ...]:
    """The companion files in ``LTX23_EXTRAS_REPO`` a 2.3 checkpoint loads alongside itself.

    Same variant rule as the assembly, so the download plan stages exactly what the load reads
    (they are otherwise fetched inline, outside the panel's progress, cancel and disk preflight)."""
    variant = checkpoint_variant(checkpoint_path)
    return tuple(
        template.format(variant = variant)
        for template in (_EXTRAS_TEXT_PROJ, _EXTRAS_VIDEO_VAE, _EXTRAS_AUDIO_VAE)
    )


# Upstream ltx_core's DISTILLED_SIGMA_VALUES: the fixed 8-step sampling curve the 22B distilled
# DiT was trained against (the scheduler appends the terminal 0 itself). The base scheduler's
# resolution-shifted spacing lands far from it at every mu the pipeline can compute, so the
# distilled default of 8 steps must pass this list verbatim.
LTX23_DISTILLED_SIGMAS: tuple[float, ...] = (
    1.0,
    0.99375,
    0.9875,
    0.98125,
    0.975,
    0.909375,
    0.725,
    0.421875,
)


def ltx2_distilled_ids(*ids: Optional[str]) -> bool:
    """True when any loaded-checkpoint id names the distilled DiT (same substring the
    generation-defaults table keys on, so sigmas and the 8-step default stay in lockstep)."""
    return any("distilled" in str(i or "").lower() for i in ids)


def ltx23_verbatim_sigmas(pipe: Any) -> Any:
    """Context manager neutralising the scheduler transforms that re-shape even explicit
    ``sigmas`` (FlowMatchEulerDiscreteScheduler applies dynamic time-shift and the
    shift_terminal stretch to caller-provided lists): dynamic shifting off, shift 1.0
    (identity), no terminal stretch, restored on exit. Without this the calibrated curve
    above would arrive at the DiT distorted (its 0.421875 tail clamped to 0.1)."""
    import contextlib

    @contextlib.contextmanager
    def _ctx():
        sched = getattr(pipe, "scheduler", None)
        cfg = getattr(sched, "config", None)
        register = getattr(sched, "register_to_config", None)
        if cfg is None or not callable(register):
            yield
            return
        saved = {
            "use_dynamic_shifting": cfg.get("use_dynamic_shifting", False),
            "shift": cfg.get("shift", 1.0),
            "shift_terminal": cfg.get("shift_terminal", None),
        }
        register(use_dynamic_shifting = False, shift = 1.0, shift_terminal = None)
        try:
            yield
        finally:
            register(**saved)

    return _ctx()


# ── component builders ───────────────────────────────────────────────────────


def _build_from_config(
    model_cls: Any,
    config: dict[str, Any],
    state: dict[str, Any],
    rename: dict[str, str],
    torch_dtype: Any,
    remove_suffixes: tuple[str, ...] = (),
) -> Any:
    from accelerate import init_empty_weights

    state = _apply_rename(_to_plain_dtype(state, torch_dtype), rename)
    for key in [k for k in state if k.endswith(remove_suffixes)] if remove_suffixes else []:
        state.pop(key)
    with init_empty_weights():
        model = model_cls.from_config(config)
    model.load_state_dict(state, strict = True, assign = True)
    return model.to(torch_dtype)


def load_ltx23_transformer(
    dit_state: dict[str, Any],
    *,
    base_repo: str,
    torch_dtype: Any,
    is_gguf: bool,
    hf_token: Optional[str],
) -> Any:
    import diffusers
    from diffusers import LTX2VideoTransformer3DModel

    # Pre-rename the 2.3-only keys the converter doesn't know, then from_single_file merges the
    # config overrides into the base 2.0 config and runs the stock key conversion.
    for old, new in _TRANSFORMER_PRERENAME:
        for key in [k for k in dit_state if k.startswith(old)]:
            dit_state[new + key[len(old) :]] = dit_state.pop(key)
    kwargs: dict[str, Any] = {
        "config": base_repo,
        "subfolder": "transformer",
        "torch_dtype": torch_dtype,
        "token": hf_token,
        **LTX_2_3_TRANSFORMER_CONFIG_OVERRIDES,
    }
    if is_gguf:
        kwargs["quantization_config"] = diffusers.GGUFQuantizationConfig(compute_dtype = torch_dtype)
    return LTX2VideoTransformer3DModel.from_single_file(dit_state, **kwargs)


def load_ltx23_connectors(
    connector_state: dict[str, Any], *, variant: str, torch_dtype: Any, hf_token: Optional[str]
) -> Any:
    from diffusers.pipelines.ltx2.connectors import LTX2TextConnectors

    # Transformer-only checkpoints carry the connector stacks but not the per-modality text
    # projections; fetch those from the companion file.
    if not any(k.startswith("text_embedding_projection") for k in connector_state):
        connector_state = dict(connector_state)
        connector_state.update(
            _load_extras_file(_EXTRAS_TEXT_PROJ.format(variant = variant), hf_token)
        )
    return _build_from_config(
        LTX2TextConnectors,
        _CONNECTORS_CONFIG,
        connector_state,
        _CONNECTORS_RENAME,
        torch_dtype,
    )


def load_ltx23_vae(
    vae_state: dict[str, Any], *, variant: str, torch_dtype: Any, hf_token: Optional[str]
) -> Any:
    from diffusers import AutoencoderKLLTX2Video
    if not vae_state:
        vae_state = _load_extras_file(_EXTRAS_VIDEO_VAE.format(variant = variant), hf_token)
    return _build_from_config(
        AutoencoderKLLTX2Video,
        _VIDEO_VAE_CONFIG,
        vae_state,
        _VIDEO_VAE_RENAME,
        torch_dtype,
        remove_suffixes = _VIDEO_VAE_REMOVE_SUFFIXES,
    )


def load_ltx23_audio_vae_and_vocoder(
    audio_vae_state: dict[str, Any],
    vocoder_state: dict[str, Any],
    *,
    variant: str,
    torch_dtype: Any,
    hf_token: Optional[str],
) -> tuple[Any, Any]:
    from diffusers import AutoencoderKLLTX2Audio
    from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE

    if not audio_vae_state or not vocoder_state:
        combined = _load_extras_file(_EXTRAS_AUDIO_VAE.format(variant = variant), hf_token)
        audio_vae_state = {
            k[len("audio_vae.") :]: v for k, v in combined.items() if k.startswith("audio_vae.")
        }
        vocoder_state = {
            k[len("vocoder.") :]: v for k, v in combined.items() if k.startswith("vocoder.")
        }
    audio_vae = _build_from_config(
        AutoencoderKLLTX2Audio,
        _AUDIO_VAE_CONFIG,
        audio_vae_state,
        _AUDIO_VAE_RENAME,
        torch_dtype,
    )
    # The 2.3 vocoder is a composite (base + bandwidth-extension stack + mel STFT buffers); keys line
    # up module-for-module after the renames.
    vocoder_state = _apply_rename(_to_plain_dtype(vocoder_state, torch_dtype), _VOCODER_RENAME)
    for key in [k for k in vocoder_state if ".ups." in k]:
        vocoder_state[key.replace(".ups.", ".upsamplers.")] = vocoder_state.pop(key)
    from accelerate import init_empty_weights

    with init_empty_weights():
        vocoder = LTX2VocoderWithBWE.from_config(_VOCODER_CONFIG)
    vocoder.load_state_dict(vocoder_state, strict = True, assign = True)
    return audio_vae, vocoder.to(torch_dtype)


# ── pipeline assembly ────────────────────────────────────────────────────────


def load_ltx23_pipeline(
    checkpoint_path: Path | str,
    *,
    base_repo: str,
    torch_dtype: Any,
    is_gguf: bool,
    hf_token: Optional[str] = None,
    text_encoder: Optional[Any] = None,
) -> Any:
    """Full LTX-2.3 pipeline from a single-file/GGUF checkpoint. Assembled per-component
    (constructor, not from_pretrained) because the base model_index pins LTX2Vocoder while 2.3
    needs LTX2VocoderWithBWE, which the type gate would reject.

    ``text_encoder`` supplies an already-built encoder (the caller's pre-cast fp8 Gemma3);
    None builds it dense from the base repo. Because the assembly bypasses
    ``from_pretrained``, this is the only way an fp8 request reaches the 2.3 path."""
    import transformers
    from diffusers import LTX2Pipeline
    from diffusers.loaders.single_file_utils import load_single_file_checkpoint

    variant = checkpoint_variant(checkpoint_path)
    logger.info(
        "video.ltx23_assembly: variant=%s gguf=%s extras=%s",
        variant,
        is_gguf,
        LTX23_EXTRAS_REPO,
    )
    state = load_single_file_checkpoint(str(checkpoint_path))
    groups = _split_checkpoint(state)
    del state

    # The Lightricks fp8 single files store SCALED float8 weights (.weight_scale/.input_scale
    # companions). Casting without the scales corrupts every quantized layer, so refuse loudly; use
    # the GGUF quants (Q8_0 for highest fidelity) instead.
    if any(k.endswith((".weight_scale", ".input_scale")) for k in groups["dit"]):
        raise ValueError(
            "This LTX checkpoint stores scaled fp8 weights, which this loader does "
            "not dequantize yet. Use the GGUF quants from unsloth/LTX-2.3-GGUF "
            "instead (Q8_0 for the highest fidelity) or the official bf16 checkpoint."
        )

    transformer = load_ltx23_transformer(
        groups["dit"],
        base_repo = base_repo,
        torch_dtype = torch_dtype,
        is_gguf = is_gguf,
        hf_token = hf_token,
    )
    connectors = load_ltx23_connectors(
        groups["connectors"],
        variant = variant,
        torch_dtype = torch_dtype,
        hf_token = hf_token,
    )
    vae = load_ltx23_vae(groups["vae"], variant = variant, torch_dtype = torch_dtype, hf_token = hf_token)
    audio_vae, vocoder = load_ltx23_audio_vae_and_vocoder(
        groups["audio_vae"],
        groups["vocoder"],
        variant = variant,
        torch_dtype = torch_dtype,
        hf_token = hf_token,
    )

    # Shared 2.0/2.3 components from the base repo, via model_index so upstream class renames break
    # loudly here rather than drift silently.
    index = LTX2Pipeline.load_config(base_repo, token = hf_token)

    def _sub(name: str, **extra: Any) -> Any:
        library, class_name = index[name]
        module = transformers if library == "transformers" else __import__("diffusers")
        return getattr(module, class_name).from_pretrained(
            base_repo, subfolder = name, token = hf_token, **extra
        )

    scheduler = _sub("scheduler")
    tokenizer = _sub("tokenizer")
    if text_encoder is None:
        text_encoder = _sub("text_encoder", torch_dtype = torch_dtype)

    return LTX2Pipeline(
        scheduler = scheduler,
        text_encoder = text_encoder,
        tokenizer = tokenizer,
        connectors = connectors,
        transformer = transformer,
        vae = vae,
        audio_vae = audio_vae,
        vocoder = vocoder,
    )