single-file bundle loader: metadata config, prefix routing, per-stream ff bias, encoder-root resolution

examples/train/configs/fine_tuning/ltx2_5/nvfp4_qat_t2v.yaml

@@ -0,0 +1,130 @@
+# Dual-stream T2V NVFP4 quantization-aware fine-tune (QAT).
+#
+# Adapted from fine_tuning/ltx2_3/nvfp4_qat_t2v.yaml. The quantization wiring
+# is UNCHANGED and deliberately so: the curated linear prefix set in
+# fastvideo/layers/quantization/nvfp4_config.py targets
+# ltx2.blocks.{0..47}.* and attaches to this architecture unmodified --
+# verified against the constructed module tree, 576 linears / 12.9B params,
+# the same 12 suffix classes x 48 blocks as before. Do NOT narrow the set for
+# this architecture: the cross-modal projections it targets
+# (audio_to_video_attn, video_to_audio_attn) are deliberately enumerated in the
+# DEPLOYMENT set, so dropping them in training would break train/deploy
+# symmetry.
+#
+# What is NOT targeted, and why (all deliberate, none accidental):
+#   * head-dim-64 audio self/cross attention -- same constraint that kept it
+#     dense before;
+#   * per-head gate projections (to_gate_logits, [32, width]) -- tiny and
+#     precision-sensitive;
+#   * the embeddings connectors -- they live in the TEXT ENCODER, not the DiT,
+#     so deployment targeting must not reach across the component boundary.
+# The coverage audit prints these as skipped-by-rule so a reader can see the
+# exclusions were chosen, and FAILS if any declared class matches zero modules.
+#
+# Video attention uses ATTN_QAT_TRAIN (quantized forward, STE backward).
+#
+# Data paths, step counts and learning rate below are placeholders to adapt.
+
+models:
+  student:
+    _target_: fastvideo.train.models.ltx2.LTX2Model
+    # REQUIRED: single-file bundle path. The config parser does NOT expand
+    # environment interpolations, so override this on the command line:
+    #   --models.student.init_from /path/to/bundle.safetensors
+    init_from: SET_ME_ON_THE_COMMAND_LINE
+    trainable: true
+    enable_gradient_checkpointing_type: full
+    attention_backend: ATTN_QAT_TRAIN
+
+method:
+  _target_: fastvideo.train.methods.fine_tuning.finetune.FineTuneMethod
+
+training:
+  distributed:
+    # ~19B params: bf16 weights 35.4 GiB + bf16 grads 35.4 GiB + Adam fp32
+    # m/v 141.5 GiB = ~212 GiB, against ~173 GiB usable per GPU. Single-GPU is
+    # arithmetically impossible; 4-way shard puts it at ~53 GiB/GPU.
+    num_gpus: 4
+    sp_size: 1
+    tp_size: 1
+    hsdp_replicate_dim: 1
+    hsdp_shard_dim: 4
+
+  data:
+    # REQUIRED: path to your preprocessed dataset.
+    data_path: data/your_dataset_preprocessed
+    dataloader_num_workers: 4
+    train_batch_size: 1
+    # LTX2Model requires 0.0: CFG dropout would zero post-connector
+    # embeddings, which is not the model's unconditional input.
+    training_cfg_rate: 0.0
+    seed: 42
+    # Must match your preprocessed resolution/length.
+    # num_latent_t = (num_frames - 1) / 8 + 1.
+    num_latent_t: 11
+    num_height: 480
+    num_width: 832
+    num_frames: 81
+
+  optimizer:
+    # Carried from the validated LTX-2.0 QAT overfit run as a starting
+    # point; tune for your dataset size and batch configuration.
+    learning_rate: 5.0e-5
+    betas: [0.9, 0.999]
+    weight_decay: 0.0
+    lr_scheduler: constant
+    lr_warmup_steps: 0
+
+  loop:
+    # Workload-dependent: set from your dataset size and target epochs.
+    max_train_steps: 2000
+    gradient_accumulation_steps: 1
+
+  checkpoint:
+    output_dir: outputs/dual_stream_t2v_nvfp4_qat_finetune
+    # A full training-state checkpoint is ~150GB for the 13B trainable
+    # video branch — size the interval and total limit to your storage.
+    training_state_checkpointing_steps: 500
+    checkpoints_total_limit: 2
+    resume_from_checkpoint: latest
+
+  tracker:
+    # No experiment tracking.
+    # It must be ["none"], NOT []: build_tracker appends wandb whenever the
+    # list is empty and project_name is non-empty, so [] silently means wandb.
+    trackers: ["none"]
+
+  model:
+    # LTX2Model.predict_noise converts the DiT's x0 output to velocity,
+    # so the default noise-minus-clean target reproduces the official
+    # unweighted masked-MSE (mask is all-ones for plain T2V).
+    precondition_outputs: false
+    enable_gradient_checkpointing_type: full
+
+callbacks:
+  grad_clip:
+    _target_: fastvideo.train.callbacks.grad_clip.GradNormClipCallback
+    max_grad_norm: 1.0
+  validation:
+    _target_: fastvideo.train.callbacks.validation.ValidationCallback
+    pipeline_target: fastvideo.pipelines.basic.ltx2.ltx2_pipeline.LTX2Pipeline
+    # REQUIRED: prompts to sample during training.
+    dataset_file: data/your_dataset_preprocessed/validation_prompts.json
+    every_steps: 250
+    # 8-step single-pass sampling matches the distilled checkpoint
+    # (validated in the LTX-2.3 overfit runs).
+    sampling_steps: [8]
+    guidance_scale: 1.0
+    num_frames: 81
+    # Validation reuses the live transformer and temporarily replaces its
+    # ATTN_QAT_TRAIN implementations with the sm120 inference kernel.
+    # Set false on GB200 (see header).
+    attn_qat_infer: true
+
+# quant_config selects NVFP4 QAT for the same curated linear prefixes as
+# LTX-2 NVFP4 deployment (see the LTX-2.3 applicability note in the
+# header). They fake-quantize through real FP4 GEMMs with an STE
+# backward. The string resolves to NVFP4QATTrainConfig at parse.
+pipeline:
+  dit_config:
+    quant_config: nvfp4_qat_train

fastvideo/configs/models/dits/ltx2.py

@@ -55,6 +55,13 @@ class LTX2VideoArchConfig(DiTArchConfig):
     # LTX-2.3 gated extensions. All default OFF == LTX-2.0 behavior.
     cross_attention_adaln: bool = False
     caption_proj_before_connector: bool = False
+    # FFN bias, per stream. Some checkpoints ship the video FFN without bias
+    # (``ff_bias: false`` in their metadata), which drops the 96
+    # ``transformer_blocks.*.ff.net.{0.proj,2}.bias`` tensors; the audio FFN is
+    # configured independently and commonly keeps its biases. Both default True,
+    # which is the existing behavior.
+    ff_bias: bool = True
+    audio_ff_bias: bool = True
 
     positional_embedding_theta: float = 10000.0
     positional_embedding_max_pos: list[int] = field(default_factory=lambda: [20, 2048, 2048])

fastvideo/models/dits/ltx2.py

@@ -333,6 +333,7 @@ class GELUApprox(nn.Module):
         self,
         in_features: int,
         out_features: int,
+        bias: bool = True,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ):
@@ -340,6 +341,7 @@ class GELUApprox(nn.Module):
         self.proj = ReplicatedLinear(
             in_features,
             out_features,
+            bias=bias,
             quant_config=quant_config,
             prefix=f"{prefix}.fc_in",
         )
@@ -357,6 +359,7 @@ class FeedForward(nn.Module):
         dim: int,
         dim_out: int,
         mult: int = 4,
+        bias: bool = True,
         quant_config: QuantizationConfig | None = None,
         prefix: str = "",
     ) -> None:
@@ -365,12 +368,14 @@ class FeedForward(nn.Module):
         project_in = GELUApprox(
             dim,
             inner_dim,
+            bias=bias,
             quant_config=quant_config,
             prefix=f"{prefix}.ffn",
         )
         project_out = ReplicatedLinear(
             inner_dim,
             dim_out,
+            bias=bias,
             quant_config=quant_config,
             prefix=f"{prefix}.ffn.fc_out",
         )
@@ -1206,6 +1211,9 @@ class TransformerConfig:
     # LTX-2.3 gated extensions (default OFF == LTX-2.0 behavior).
     apply_gated_attention: bool = False
     cross_attention_adaln: bool = False
+    # FFN bias is per stream: a checkpoint may drop it on video while keeping
+    # it on audio. Default True preserves existing behavior.
+    ff_bias: bool = True
 
 
 class LTXDistributedAttention(DistributedAttention):
@@ -1846,6 +1854,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
             self.ff = FeedForward(
                 video.dim,
                 dim_out=video.dim,
+                bias=video.ff_bias,
                 quant_config=quant_config,
                 prefix=f"{prefix}.blocks.{idx}",
             )
@@ -1884,6 +1893,7 @@ class BasicAVTransformerBlock(torch.nn.Module):
             self.audio_ff = FeedForward(
                 audio.dim,
                 dim_out=audio.dim,
+                bias=audio.ff_bias,
                 quant_config=quant_config,
                 prefix=f"{prefix}.blocks.{idx}.audio",
             )
@@ -2353,6 +2363,8 @@ class LTXModel(torch.nn.Module):
         cross_attention_adaln: bool = False,
         caption_proj_before_connector: bool = False,
         apply_gated_attention: bool = False,
+        ff_bias: bool = True,
+        audio_ff_bias: bool = True,
         stg_block_idx: int = 29,
         use_distributed_attention: bool = False,
         quant_config: QuantizationConfig | None = None,
@@ -2364,6 +2376,9 @@ class LTXModel(torch.nn.Module):
         self.cross_attention_adaln = cross_attention_adaln
         self.caption_proj_before_connector = caption_proj_before_connector
         self.apply_gated_attention = apply_gated_attention
+        # Per-stream FFN bias; default True preserves existing behavior.
+        self.ff_bias = ff_bias
+        self.audio_ff_bias = audio_ff_bias
         self.stg_block_idx = stg_block_idx
         self.use_middle_indices_grid = use_middle_indices_grid
         self.rope_type = rope_type
@@ -2584,6 +2599,7 @@ class LTXModel(torch.nn.Module):
             context_dim=cross_attention_dim,
             apply_gated_attention=self.apply_gated_attention,
             cross_attention_adaln=self.cross_attention_adaln,
+            ff_bias=self.ff_bias,
         ) if self.model_type.is_video_enabled() else None)
         audio_config = (TransformerConfig(
             dim=self.audio_inner_dim,
@@ -2592,6 +2608,7 @@ class LTXModel(torch.nn.Module):
             context_dim=audio_cross_attention_dim,
             apply_gated_attention=self.apply_gated_attention,
             cross_attention_adaln=self.cross_attention_adaln,
+            ff_bias=self.audio_ff_bias,
         ) if self.model_type.is_audio_enabled() else None)
         self.use_distributed_attention = use_distributed_attention
         self.transformer_blocks = torch.nn.ModuleList([
@@ -2778,6 +2795,8 @@ class LTX2Transformer3DModel(BaseDiT):
             cross_attention_adaln=arch.cross_attention_adaln,
             caption_proj_before_connector=arch.caption_proj_before_connector,
             apply_gated_attention=arch.apply_gated_attention,
+            ff_bias=arch.ff_bias,
+            audio_ff_bias=arch.audio_ff_bias,
             stg_block_idx=arch.stg_block_idx,
             use_distributed_attention=use_distributed_attention,
             quant_config=config.quant_config,

fastvideo/models/encoders/gemma.py

@@ -7,7 +7,7 @@ from typing import Any, Iterable
 
 import torch
 from torch import nn
-from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
+from transformers import AutoModelForImageTextToText, AutoTokenizer
 
 from fastvideo.configs.models.encoders import BaseEncoderOutput, TextEncoderConfig
 from fastvideo.models.encoders.base import TextEncoder
@@ -398,7 +398,7 @@ class LTX2GemmaTextEncoderModel(TextEncoder):
         self.gemma_model_path = arch.gemma_model_path
         self.gemma_dtype = arch.gemma_dtype
         self.padding_side = arch.padding_side
-        self._gemma_model: Gemma3ForConditionalGeneration | None = None
+        self._gemma_model: nn.Module | None = None
 
     def named_parameters(self, prefix: str = "", recurse: bool = True):
         for name, param in super().named_parameters(prefix=prefix, recurse=recurse):
@@ -436,17 +436,22 @@ class LTX2GemmaTextEncoderModel(TextEncoder):
             )
 
     @property
-    def gemma_model(self) -> Gemma3ForConditionalGeneration:
+    def gemma_model(self) -> nn.Module:
         if self._gemma_model is None:
             gemma_path = self.gemma_model_path
             if not gemma_path:
                 raise ValueError("gemma_model_path must be set (expected text_encoder/gemma).")
             dtype = getattr(torch, self.gemma_dtype, torch.bfloat16)
-            self._gemma_model = Gemma3ForConditionalGeneration.from_pretrained(
+            # Resolve the class from the root config's own ``model_type``
+            # instead of hardcoding one: a checkpoint may declare its own
+            # encoder pairing. The auto class still resolves to the previously
+            # hardcoded class for the roots shipped so far.
+            self._gemma_model = AutoModelForImageTextToText.from_pretrained(
                 gemma_path,
                 local_files_only=True,
                 torch_dtype=dtype,
             )
+            self._sync_arch_from_gemma_config(self._gemma_model.config)
             # Configure model-level attention implementation when using TORCH_SDPA.
             # Note: torch.backends.cuda.enable_*_sdp() settings should be configured
             # at application/pipeline initialization level, not here, to avoid
@@ -464,6 +469,33 @@ class LTX2GemmaTextEncoderModel(TextEncoder):
             self._gemma_model.eval()
         return self._gemma_model
 
+    def _sync_arch_from_gemma_config(self, gemma_config: Any) -> None:
+        """Take language-model geometry from the loaded Gemma config.
+
+        A multimodal root nests it under ``text_config``; a text-only config
+        carries it at the root. The dataclass defaults hold for every encoder
+        root shipped so far, so this is not a behavior change -- it just stops
+        the defaults from being authoritative.
+        The feature-extractor linears are sized in ``__init__`` from
+        ``feature_extractor_in_features``, before the Gemma weights load, so a
+        geometry mismatch is raised here rather than surfacing later as an
+        opaque matmul shape error.
+        """
+        text_config = getattr(gemma_config, "text_config", gemma_config)
+        arch = self.config.arch_config
+        arch.hidden_size = text_config.hidden_size
+        arch.num_hidden_layers = text_config.num_hidden_layers
+        # +1: the stacked hidden states include the embedding output.
+        expected_in_features = arch.hidden_size * (arch.num_hidden_layers + 1)
+        if expected_in_features != arch.feature_extractor_in_features:
+            raise ValueError(
+                "Gemma geometry does not match the configured feature "
+                f"extractor: loaded config gives hidden_size={arch.hidden_size} "
+                f"x (num_hidden_layers={arch.num_hidden_layers} + 1) = "
+                f"{expected_in_features}, but feature_extractor_in_features is "
+                f"{arch.feature_extractor_in_features}."
+            )
+
     def _run_feature_extractor(
         self,
         hidden_states: tuple[torch.Tensor, ...],
@@ -670,6 +702,10 @@ class LTX2GemmaTextEncoderModel(TextEncoder):
                 name = "feature_extractor_linear.aggregate_embed.weight"
             elif name.startswith("video_connector."):
                 name = name.replace("video_connector.", "embeddings_connector.", 1)
+            elif name.startswith("video_embeddings_connector."):
+                # Some checkpoints name the video connector sub-tree after its
+                # modality; the module is just ``embeddings_connector``.
+                name = name.replace("video_embeddings_connector.", "embeddings_connector.", 1)
             elif name.startswith("audio_connector."):
                 name = name.replace("audio_connector.", "audio_embeddings_connector.", 1)
             if name not in params_dict:

fastvideo/models/loader/component_loader.py

@@ -33,6 +33,10 @@ from fastvideo.logger import init_logger
 from fastvideo.models.encoders.base import TextEncoder
 from fastvideo.models.hf_transformer_utils import get_diffusers_config
 from fastvideo.models.loader.fsdp_load import maybe_load_fsdp_model, shard_model
+from fastvideo.models.loader.ltx_single_file import (
+    component_weights,
+    read_ltx_metadata,
+)
 from fastvideo.models.loader.utils import set_default_torch_dtype
 from fastvideo.models.loader.weight_utils import (
     filter_duplicate_safetensors_files,
@@ -290,6 +294,18 @@ class TextEncoderLoader(ComponentLoader):
                     rope_type = transformer_config.get("rope_type")
                     if rope_type is not None:
                         model_config["connector_rope_type"] = rope_type
+                # Each per-modality projection feeds its connector, so its
+                # output width is that stream's cross-attention width, which
+                # only the transformer config declares. Absent from both
+                # configs -> the arch config default stands.
+                for src, dst in (
+                    ("cross_attention_dim",
+                     "video_feature_extractor_out_features"),
+                    ("audio_cross_attention_dim",
+                     "audio_feature_extractor_out_features"),
+                ):
+                    if dst not in model_config and src in transformer_config:
+                        model_config[dst] = transformer_config[src]
             except json.JSONDecodeError:
                 pass
         logger.info("HF Model config: %s", model_config)
@@ -974,7 +990,16 @@ class TransformerLoader(ComponentLoader):
 
     def load(self, model_path: str, fastvideo_args: FastVideoArgs):
         """Load the transformer based on the model path, and inference args."""
-        config = get_diffusers_config(model=model_path)
+        # LTX ships one bundle holding every component instead of a
+        # per-component directory: the config lives in the file's
+        # ``__metadata__`` and the transformer's tensors are selected by their
+        # top-level key prefix.
+        single_file = str(model_path).endswith(".safetensors")
+        if single_file:
+            config = deepcopy(
+                read_ltx_metadata(model_path).config["transformer"])
+        else:
+            config = get_diffusers_config(model=model_path)
         hf_config = deepcopy(config)
         cls_name = config.pop("_class_name")
         config.pop("_name_or_path", None)
@@ -1008,9 +1033,12 @@ class TransformerLoader(ComponentLoader):
         model_cls, _ = ModelRegistry.resolve_model_cls(cls_name)
 
         # Find all safetensors files
-        safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
-        if not safetensors_list:
-            raise ValueError(f"No safetensors files found in {model_path}")
+        if single_file:
+            safetensors_list = [str(model_path)]
+        else:
+            safetensors_list = glob.glob(os.path.join(str(model_path), "*.safetensors"))
+            if not safetensors_list:
+                raise ValueError(f"No safetensors files found in {model_path}")
 
         # arch_config can infer architecture from weight keys (e.g. Flux2 layer counts)
         update_fn = getattr(dit_config.arch_config, "update_from_weight_keys", None)
@@ -1065,6 +1093,11 @@ class TransformerLoader(ComponentLoader):
                     "hf_config": hf_config
                 },
                 weight_dir_list=safetensors_list,
+                # Route the bundle's tensors to the transformer, prefix
+                # stripped. Skipped when custom init weights replaced the
+                # file list above.
+                weight_iterator=(component_weights(model_path, "transformer")
+                                 if single_file and not use_custom_weights else None),
                 device=get_local_torch_device(),
                 hsdp_replicate_dim=fastvideo_args.hsdp_replicate_dim,
                 hsdp_shard_dim=fastvideo_args.hsdp_shard_dim,

fastvideo/models/loader/fsdp_load.py

@@ -7,7 +7,7 @@
 from __future__ import annotations
 import os
 import contextlib
-from collections.abc import Callable, Generator
+from collections.abc import Callable, Generator, Iterable
 from itertools import chain
 from typing import Any
 
@@ -136,9 +136,15 @@ def maybe_load_fsdp_model(
     pin_cpu_memory: bool = True,
     enable_torch_compile: bool = False,
     torch_compile_kwargs: dict[str, Any] | None = None,
+    weight_iterator: Iterable[tuple[str, torch.Tensor]] | None = None,
 ) -> torch.nn.Module:
     """
     Load the model with FSDP if is training, else load the model without FSDP.
+
+    ``weight_iterator`` overrides reading ``weight_dir_list``, for checkpoints
+    whose keys need routing before they reach the model (see
+    ``ltx_single_file.component_weights``). It must yield the same
+    ``(name, cpu_tensor)`` pairs that ``safetensors_weights_iterator`` does.
     """
     # NOTE(will): cast_forward_inputs=True shouldn't be needed as we are
     # manually casting the inputs to the model
@@ -201,7 +207,8 @@ def maybe_load_fsdp_model(
                     fsdp_shard_conditions=model._fsdp_shard_conditions,
                     pin_cpu_memory=pin_cpu_memory)
 
-    weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=True)
+    if weight_iterator is None:
+        weight_iterator = safetensors_weights_iterator(weight_dir_list, to_cpu=True)
     param_names_mapping_fn = get_param_names_mapping(model.param_names_mapping)
     load_model_from_full_model_state_dict(
         model,

fastvideo/models/loader/ltx_single_file.py

@@ -0,0 +1,216 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Single-file LTX checkpoint support.
+
+LTX ships one ``.safetensors`` bundle holding every component, so the two
+assumptions the directory-based loaders make do not hold:
+
+* there is no per-component ``config.json`` -- every component's config lives
+  in the file's ``__metadata__`` (safetensors stores an 8-byte little-endian
+  header length, then that many bytes of JSON whose ``__metadata__`` object is
+  a flat ``str -> str`` map; ``safe_open(...).metadata()`` returns it);
+* there is no per-component directory -- components are told apart by the
+  top-level prefix on each tensor key.
+
+This module reads that metadata and routes tensors by prefix. It does not
+convert anything to a diffusers layout.
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from collections.abc import Generator
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from safetensors import safe_open
+
+from fastvideo.configs.models.dits.ltx2 import LTX2VideoConfig
+from fastvideo.logger import init_logger
+
+logger = init_logger(__name__)
+
+# Top-level tensor key prefix per component. Verified against the LTX bundle
+# header; every tensor in the file starts with one of these except
+# ``duration_head.``, which has no FastVideo component today.
+COMPONENT_PREFIXES: dict[str, str] = {
+    "transformer": "model.diffusion_model.",
+    "vae": "vae.",
+    "audio_vae": "audio_vae.",
+    "vocoder": "vocoder.",
+    "text_encoder": "text_embedding_projection.",
+}
+
+# Both embeddings connectors are stored *under* the transformer prefix, but
+# FastVideo builds and runs them in the text encoder
+# (``LTX2GemmaTextEncoderModel``'s ``Embeddings1DConnector``), not in the DiT,
+# so they are routed to the text encoder instead. Only the transformer prefix
+# is stripped: the sub-tree name survives so the text encoder's own rename
+# table can map it onto the module name, the same way the already-converted
+# repo layout is remapped. Matches how ``convert_ltx2_weights.py`` splits them.
+TEXT_STACK_SUBPREFIXES: tuple[str, ...] = (
+    "video_embeddings_connector.",
+    "audio_embeddings_connector.",
+)
+
+
+@dataclass(frozen=True)
+class LTXCheckpointMetadata:
+    """Parsed ``__metadata__`` of a single-file LTX checkpoint.
+
+    ``config`` holds the per-component config sections keyed by component
+    (``transformer``, ``vae``, ``scheduler``, ``audio_vae``, ``vocoder``).
+    ``model_version`` and ``gemma_source_checkpoint`` are carried through so
+    callers can validate the text encoder against the checkpoint that trained
+    it; ``gemma_source_checkpoint`` is absent on some variants.
+    """
+
+    config: dict[str, Any]
+    model_version: str | None
+    gemma_source_checkpoint: dict[str, Any] | None
+
+
+def read_ltx_metadata(path: str) -> LTXCheckpointMetadata:
+    """Read the ``__metadata__`` config out of a single-file LTX checkpoint.
+
+    Only the header is parsed; no tensor data is touched. The remaining
+    metadata entries are deliberately not returned or logged -- the bundle also
+    carries a full license text and an ``encrypted_wandb_properties`` blob.
+    """
+    with safe_open(path, framework="pt") as f:
+        metadata = f.metadata() or {}
+
+    config = json.loads(metadata["config"])
+
+    transformer = config.get("transformer")
+    if transformer is not None:
+        # ``frequencies_precision`` is the checkpoint's name for what the arch
+        # config calls ``double_precision_rope``; without this the field would
+        # silently fall back to its dataclass default.
+        precision = transformer.get("frequencies_precision")
+        if precision is not None:
+            transformer["double_precision_rope"] = precision == "float64"
+
+    gemma_source = metadata.get("gemma_source_checkpoint")
+    return LTXCheckpointMetadata(
+        config=config,
+        model_version=metadata.get("model_version"),
+        gemma_source_checkpoint=(json.loads(gemma_source)
+                                 if gemma_source is not None else None),
+    )
+
+
+def build_dit_config(metadata: LTXCheckpointMetadata) -> LTX2VideoConfig:
+    """Build the LTX-2 DiT config from checkpoint metadata.
+
+    Reuses ``update_model_arch``, so the metadata keys that name an arch field
+    win and the rest of the section is ignored -- no hand-written constants.
+    """
+    config = LTX2VideoConfig()
+    config.update_model_arch(metadata.config["transformer"])
+    return config
+
+
+def component_weights(
+    path: str,
+    component: str,
+) -> Generator[tuple[str, torch.Tensor], None, None]:
+    """Yield ``(key_with_prefix_stripped, tensor)`` for one component.
+
+    The transformer prefix covers two owners -- see
+    ``TEXT_STACK_SUBPREFIXES`` -- so keys under it are split before the
+    component filter applies.
+
+    ``safe_open`` mmaps the file and ``get_tensor`` materializes one tensor at
+    a time, so iterating never holds more than a single tensor in RAM. Callers
+    that need a state dict for a small component can do
+    ``dict(component_weights(path, "vocoder"))``; do not do that for the
+    transformer.
+
+    ponytail: every rank reads the file itself, unlike
+    ``safetensors_weights_iterator``, which has local rank 0 read and broadcast.
+    Correct either way, but the ceiling is filesystem bandwidth -- N ranks read
+    N copies. Add the broadcast if loading a bundle over a shared filesystem
+    turns out to be the bottleneck.
+    """
+    prefix = COMPONENT_PREFIXES[component]
+    transformer_prefix = COMPONENT_PREFIXES["transformer"]
+    with safe_open(path, framework="pt") as f:
+        for key in f.keys():
+            if key.startswith(transformer_prefix):
+                name = key[len(transformer_prefix):]
+                # str.startswith takes a tuple.
+                owner = ("text_encoder"
+                         if name.startswith(TEXT_STACK_SUBPREFIXES) else
+                         "transformer")
+                if owner != component:
+                    continue
+            elif key.startswith(prefix):
+                name = key[len(prefix):]
+            else:
+                continue
+            yield name, f.get_tensor(key)
+
+
+def resolve_text_encoder_root(
+    configured_path: str | None,
+    override_path: str | None = None,
+    metadata: LTXCheckpointMetadata | None = None,
+) -> str:
+    """Locate the text-encoder root that goes with a single-file bundle.
+
+    A bundle carries every component's *weights* but no pointer to the text
+    encoder, which lives outside it. The root therefore has to be declared,
+    and there are exactly two ways to declare it:
+
+    * ``override_path`` -- an explicit per-run argument, so a one-off run needs
+      no config edit;
+    * ``configured_path`` -- the pipeline config's encoder path, which is where
+      model composition already lives and which survives files being moved.
+
+    Deliberately absent: any search of the bundle's directory. Picking up an
+    encoder nobody asked for changes what gets loaded without being requested,
+    and silently picks the wrong one when two sit side by side. When neither
+    source is set this raises and names both, rather than guessing.
+
+    ``metadata`` is used only to *validate* a declared root, never to find one:
+    a bundle may not declare an encoder pairing at all, so discovery cannot
+    depend on it.
+    """
+    root = override_path or configured_path
+    if not root:
+        raise ValueError(
+            "A single-file checkpoint does not carry its text encoder, so the "
+            "encoder root must be declared. Set `gemma_model_path` in the "
+            "pipeline config, or pass the text-encoder path explicitly for "
+            "this run. It is not inferred from the checkpoint's directory: "
+            "loading whichever encoder happens to sit beside the file would "
+            "silently pick the wrong one when several are present.")
+
+    expected = (metadata.gemma_source_checkpoint or {}).get(
+        "gemma_version") if metadata is not None else None
+    if expected:
+        # ponytail: warn, don't raise -- the declared root is the user's
+        # explicit instruction and the pairing is advisory. Promote to a hard
+        # error if a mismatch ever turns out to produce silent garbage rather
+        # than an obvious shape failure.
+        actual = _read_encoder_version(root)
+        if actual is not None and actual != expected:
+            logger.warning(
+                "Checkpoint expects text-encoder version %r but the encoder at "
+                "%s reports %r; continuing with the declared root.", expected,
+                root, actual)
+    return root
+
+
+def _read_encoder_version(root: str) -> str | None:
+    """The encoder's declared version, or None if it declares none."""
+    config_path = os.path.join(root, "config.json")
+    if not os.path.isfile(config_path):
+        return None
+    try:
+        with open(config_path, encoding="utf-8") as f:
+            return json.load(f).get("gemma_version")
+    except (OSError, json.JSONDecodeError):
+        return None

fastvideo/models/registry.py

@@ -38,6 +38,8 @@ _TEXT_TO_VIDEO_DIT_MODELS = {
     ("dits", "longcat_video_dit", "LongCatVideoTransformer3DModel"),  # Wrapper (Phase 1)
     "LongCatTransformer3DModel": ("dits", "longcat", "LongCatTransformer3DModel"),  # Native (Phase 2)
     "LTX2Transformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"),
+    # ``_class_name`` declared by single-file LTX checkpoint metadata.
+    "AVTransformer3DModel": ("dits", "ltx2", "LTX2Transformer3DModel"),
     "SD3Transformer2DModel": ("dits", "sd3", "SD3Transformer2DModel"),
     "LingBotWorldTransformer3DModel": ("dits", "lingbotworld", "LingBotWorldTransformer3DModel"),
     "LingBotWorld2CausalFastTransformer3DModel": (

fastvideo/tests/api/test_ltx_single_file.py

@@ -0,0 +1,174 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Checks for single-file LTX metadata parsing and prefix routing.
+
+Builds a tiny synthetic bundle shaped like the real one -- same metadata keys,
+same top-level tensor prefixes -- so no checkpoint is needed. Run directly
+(``python fastvideo/tests/api/test_ltx_single_file.py``) or under pytest.
+"""
+
+import json
+import tempfile
+from pathlib import Path
+
+import torch
+from safetensors.torch import save_file
+
+from fastvideo.models.loader.ltx_single_file import (
+    resolve_text_encoder_root,
+    build_dit_config,
+    component_weights,
+    read_ltx_metadata,
+)
+
+# Subset of the transformer section as the checkpoint declares it.
+TRANSFORMER_SECTION = {
+    "_class_name": "AVTransformer3DModel",
+    "num_layers": 48,
+    "num_attention_heads": 32,
+    "attention_head_dim": 128,
+    "audio_num_attention_heads": 32,
+    "audio_attention_head_dim": 64,
+    "cross_attention_dim": 4096,
+    "audio_cross_attention_dim": 2048,
+    "caption_channels": 3840,
+    "ff_bias": False,
+    "cross_attention_adaln": True,
+    "caption_proj_before_connector": True,
+    "apply_gated_attention": True,
+    "rope_type": "split",
+    "frequencies_precision": "float64",
+}
+
+
+def _write_bundle(path: Path, *, with_gemma_source: bool) -> None:
+    tensors = {
+        "model.diffusion_model.patchify_proj.weight": torch.zeros(2, 2),
+        "model.diffusion_model.transformer_blocks.0.ff.net.2.weight": torch.zeros(2, 2),
+        # Stored under the transformer prefix but owned by the text stack.
+        "model.diffusion_model.video_embeddings_connector.x.weight": torch.zeros(2),
+        "model.diffusion_model.audio_embeddings_connector.x.weight": torch.zeros(2),
+        "vae.encoder.conv.weight": torch.zeros(2),
+        "audio_vae.decoder.conv.weight": torch.zeros(2),
+        "vocoder.conv.weight": torch.zeros(2),
+        "text_embedding_projection.video_aggregate_embed.weight": torch.zeros(2),
+        "duration_head.linear.weight": torch.zeros(2),
+    }
+    metadata = {
+        "config": json.dumps({"transformer": TRANSFORMER_SECTION, "vae": {}}),
+        "model_version": "9.9.9",
+        "license": "x" * 128,
+    }
+    if with_gemma_source:
+        metadata["gemma_source_checkpoint"] = json.dumps(
+            {"ltx_version": "9.9.9", "gemma_version": "fake-encoder-v0"})
+    save_file(tensors, str(path), metadata=metadata)
+
+
+def test_read_metadata_and_routing() -> None:
+    with tempfile.TemporaryDirectory() as tmp:
+        path = Path(tmp) / "bundle.safetensors"
+        _write_bundle(path, with_gemma_source=True)
+
+        meta = read_ltx_metadata(str(path))
+        assert meta.model_version == "9.9.9"
+        assert meta.gemma_source_checkpoint == {
+            "ltx_version": "9.9.9",
+            "gemma_version": "fake-encoder-v0",
+        }
+        assert set(meta.config) == {"transformer", "vae"}
+        # frequencies_precision is the checkpoint's name for double_precision_rope.
+        assert meta.config["transformer"]["double_precision_rope"] is True
+
+        # Prefix stripped, connectors excluded, nothing from other components.
+        transformer_keys = {n for n, _ in component_weights(str(path), "transformer")}
+        assert transformer_keys == {
+            "patchify_proj.weight",
+            "transformer_blocks.0.ff.net.2.weight",
+        }
+        for component, expected in (
+            ("vae", {"encoder.conv.weight"}),
+            ("audio_vae", {"decoder.conv.weight"}),
+            ("vocoder", {"conv.weight"}),
+            # The connectors come off the transformer prefix with their
+            # sub-tree name intact; the text encoder renames from there.
+            ("text_encoder", {
+                "video_aggregate_embed.weight",
+                "video_embeddings_connector.x.weight",
+                "audio_embeddings_connector.x.weight",
+            }),
+        ):
+            assert {n for n, _ in component_weights(str(path), component)} == expected
+
+
+def test_missing_gemma_source_is_none() -> None:
+    with tempfile.TemporaryDirectory() as tmp:
+        path = Path(tmp) / "bundle.safetensors"
+        _write_bundle(path, with_gemma_source=False)
+        assert read_ltx_metadata(str(path)).gemma_source_checkpoint is None
+
+
+def test_dit_config_takes_ff_bias_from_metadata() -> None:
+    with tempfile.TemporaryDirectory() as tmp:
+        path = Path(tmp) / "bundle.safetensors"
+        _write_bundle(path, with_gemma_source=True)
+        arch = build_dit_config(read_ltx_metadata(str(path))).arch_config
+
+        # ff_bias comes from metadata; audio_ff_bias has no metadata key and
+        # must keep the default, matching the audio_ff.*.bias tensors that the
+        # checkpoint does carry.
+        assert arch.ff_bias is False
+        assert arch.audio_ff_bias is True
+        assert arch.num_layers == 48
+        assert arch.cross_attention_adaln is True
+        assert arch.caption_channels == 3840
+
+
+def test_defaults_unchanged_without_metadata() -> None:
+    from fastvideo.configs.models.dits.ltx2 import LTX2VideoArchConfig
+
+    arch = LTX2VideoArchConfig()
+    assert arch.ff_bias is True
+    assert arch.audio_ff_bias is True
+
+
+if __name__ == "__main__":
+    test_read_metadata_and_routing()
+    test_missing_gemma_source_is_none()
+    test_dit_config_takes_ff_bias_from_metadata()
+    test_defaults_unchanged_without_metadata()
+    print("ok")
+
+
+def test_encoder_root_requires_an_explicit_declaration() -> None:
+    """Neither source set -> raise, naming both. Never guess from the filesystem."""
+    import pytest
+    with pytest.raises(ValueError, match="must be declared"):
+        resolve_text_encoder_root(configured_path=None, override_path=None)
+
+
+def test_encoder_root_override_beats_config() -> None:
+    assert resolve_text_encoder_root("/from/config", "/from/override") == "/from/override"
+    assert resolve_text_encoder_root("/from/config", None) == "/from/config"
+
+
+def test_encoder_root_is_not_discovered_from_a_sibling_directory() -> None:
+    """A plausible encoder sitting next to the bundle must NOT be picked up."""
+    with tempfile.TemporaryDirectory() as tmp:
+        sibling = Path(tmp) / "gemma"
+        sibling.mkdir()
+        (sibling / "config.json").write_text("{}")
+        import pytest
+        with pytest.raises(ValueError):
+            resolve_text_encoder_root(configured_path=None, override_path=None)
+
+
+def test_declared_root_survives_a_version_mismatch() -> None:
+    """Pairing metadata validates a declared root; it never overrides it."""
+    with tempfile.TemporaryDirectory() as tmp:
+        root = Path(tmp) / "encoder"
+        root.mkdir()
+        (root / "config.json").write_text(json.dumps({"gemma_version": "other-v0"}))
+        meta = LTXCheckpointMetadata(
+            config={}, model_version="9.9.9",
+            gemma_source_checkpoint={"gemma_version": "fake-encoder-v0"})
+        assert resolve_text_encoder_root(str(root), None, meta) == str(root)