From 3f5614fd03e6d7db0f6aed3ab4eba494ad61a1fd Mon Sep 17 00:00:00 2001 From: William Lin Date: Wed, 29 Jul 2026 16:04:10 -0700 Subject: [PATCH] single-file bundle loader: metadata config, prefix routing, per-stream ff bias, encoder-root resolution --- .../fine_tuning/ltx2_5/nvfp4_qat_t2v.yaml | 130 +++++++++++ fastvideo/configs/models/dits/ltx2.py | 7 + fastvideo/models/dits/ltx2.py | 19 ++ fastvideo/models/encoders/gemma.py | 44 +++- fastvideo/models/loader/component_loader.py | 41 +++- fastvideo/models/loader/fsdp_load.py | 11 +- fastvideo/models/loader/ltx_single_file.py | 216 ++++++++++++++++++ fastvideo/models/registry.py | 2 + fastvideo/tests/api/test_ltx_single_file.py | 174 ++++++++++++++ 9 files changed, 634 insertions(+), 10 deletions(-) create mode 100644 examples/train/configs/fine_tuning/ltx2_5/nvfp4_qat_t2v.yaml create mode 100644 fastvideo/models/loader/ltx_single_file.py create mode 100644 fastvideo/tests/api/test_ltx_single_file.py diff --git a/examples/train/configs/fine_tuning/ltx2_5/nvfp4_qat_t2v.yaml b/examples/train/configs/fine_tuning/ltx2_5/nvfp4_qat_t2v.yaml new file mode 100644 index 000000000..8a285ddb8 --- /dev/null +++ b/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 diff --git a/fastvideo/configs/models/dits/ltx2.py b/fastvideo/configs/models/dits/ltx2.py index 135d4ca6e..5c648e320 100644 --- a/fastvideo/configs/models/dits/ltx2.py +++ b/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]) diff --git a/fastvideo/models/dits/ltx2.py b/fastvideo/models/dits/ltx2.py index 2ef569648..81a87a799 100644 --- a/fastvideo/models/dits/ltx2.py +++ b/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, diff --git a/fastvideo/models/encoders/gemma.py b/fastvideo/models/encoders/gemma.py index 944db35ff..d81aea715 100644 --- a/fastvideo/models/encoders/gemma.py +++ b/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: diff --git a/fastvideo/models/loader/component_loader.py b/fastvideo/models/loader/component_loader.py index f11f15255..722ba33dc 100644 --- a/fastvideo/models/loader/component_loader.py +++ b/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, diff --git a/fastvideo/models/loader/fsdp_load.py b/fastvideo/models/loader/fsdp_load.py index 4dfca6efb..00e92778c 100644 --- a/fastvideo/models/loader/fsdp_load.py +++ b/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, diff --git a/fastvideo/models/loader/ltx_single_file.py b/fastvideo/models/loader/ltx_single_file.py new file mode 100644 index 000000000..546f2424b --- /dev/null +++ b/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 diff --git a/fastvideo/models/registry.py b/fastvideo/models/registry.py index 4f9939a62..8bcbbc690 100644 --- a/fastvideo/models/registry.py +++ b/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": ( diff --git a/fastvideo/tests/api/test_ltx_single_file.py b/fastvideo/tests/api/test_ltx_single_file.py new file mode 100644 index 000000000..1080c5eef --- /dev/null +++ b/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)