resolve pipeline config for single-file bundles from declared transformer class

A bundle path is a file, so config resolution took the directory branch and
was rejected for having no model_index.json. Resolve it from the transformer
_class_name the checkpoint declares about itself, via an explicit alias table.
Keyed on the class name, not the registry model_id (a positional index that
rebinds on registration-order changes). No wildcard fallback and no file-name
inference; an unmapped class raises naming the table and the override.

verify_model_config_and_directory is unchanged.
This commit is contained in:
William Lin
2026-08-10 12:51:08 -07:00
committed by SolitaryThinker
parent 3f5614fd03
commit 6343eadbec
2 changed files with 111 additions and 11 deletions
+63 -1
View File
@@ -142,6 +142,18 @@ _MODEL_HF_PATH_TO_NAME: dict[str, str] = {}
# Detectors to identify model families from paths or class names
_MODEL_NAME_DETECTORS: list[tuple[str, Callable[[str], bool]]] = []
# Single-file checkpoint bundles carry no `model_index.json`, so their pipeline
# config is resolved from the transformer class the file declares about itself.
# Keyed on that `_class_name` -- a stable, checkpoint-declared identifier --
# and NOT on `model_id`, which is a positional index (`str(len(_CONFIG_REGISTRY))`)
# and would silently rebind whenever registration order changes. The value pins
# the ConfigInfo by its stable `(model_family, default_preset)` pair. A bundle's
# file NAME is never consulted.
_BUNDLE_TRANSFORMER_TO_CONFIG: dict[str, tuple[str, str]] = {
# transformer _class_name -> (model_family, default_preset)
"AVTransformer3DModel": ("ltx2", "ltx2_distilled"),
}
def register_configs(
sampling_param_cls: type[SamplingParam] | None,
@@ -187,6 +199,50 @@ def get_model_short_name(model_id: str) -> str:
return model_id
def _bundle_config_info(model_path: str) -> ConfigInfo:
"""Resolve the pipeline config for a single-file checkpoint bundle.
The bundle declares its own component classes in ``__metadata__``; reading
that is the only inference done here. The transformer's ``_class_name`` is
mapped through :data:`_BUNDLE_TRANSFORMER_TO_CONFIG`, an explicit table, to
the ``(model_family, default_preset)`` of a registered config. Nothing is
matched by wildcard, and nothing is read off the file name -- a file name is
incidental, and making resolution depend on it makes correct loading depend
on what a checkpoint happens to be called.
An unrecognized transformer class raises, naming both the table to extend
and the per-run override, rather than falling through to a guess.
ponytail: LTX is the only family shipping a bundle, so this reads LTX
metadata directly. Dispatch on a per-format reader if a second bundle
format ever appears.
"""
from fastvideo.models.loader.ltx_single_file import read_ltx_metadata
try:
cls_name = read_ltx_metadata(model_path).config.get("transformer", {}).get("_class_name")
except KeyError: # no `config` section: not a bundle we know how to read
cls_name = None
entry = _BUNDLE_TRANSFORMER_TO_CONFIG.get(cls_name) if cls_name else None
if entry is None:
raise ValueError(
f"Single-file checkpoint {model_path} declares transformer class {cls_name!r}, which is not in "
f"_BUNDLE_TRANSFORMER_TO_CONFIG (known: {sorted(_BUNDLE_TRANSFORMER_TO_CONFIG)}). Add the class to that "
"table, or select the pipeline explicitly for this run with `override_pipeline_cls_name`. The pipeline is "
"never inferred from the checkpoint's file name.")
model_family, default_preset = entry
for config_info in _CONFIG_REGISTRY.values():
if config_info.model_family == model_family and config_info.default_preset == default_preset:
return config_info
raise ValueError(
f"_BUNDLE_TRANSFORMER_TO_CONFIG maps {cls_name!r} to model_family={model_family!r} "
f"default_preset={default_preset!r}, but no registered config declares that pair. "
"Fix the table or the registration.")
def _get_config_info(
model_path: str,
*,
@@ -209,7 +265,13 @@ def _get_config_info(
model_id = _MODEL_HF_PATH_TO_NAME[registered_model_hf_id]
return _CONFIG_REGISTRY.get(model_id)
# 3. Use detectors (path or model_index pipeline name).
# 3. A single-file bundle is a FILE, so it would take the directory branch
# below and be rejected for having no model_index.json. Its config lives
# in its own `__metadata__` instead.
if model_path.endswith(".safetensors"):
return _bundle_config_info(model_path)
# 4. Use detectors (path or model_index pipeline name).
if os.path.exists(model_path):
config = verify_model_config_and_directory(model_path, required_component_dirs=[])
else:
+48 -10
View File
@@ -14,6 +14,7 @@ import torch
from safetensors.torch import save_file
from fastvideo.models.loader.ltx_single_file import (
LTXCheckpointMetadata,
resolve_text_encoder_root,
build_dit_config,
component_weights,
@@ -40,7 +41,10 @@ TRANSFORMER_SECTION = {
}
def _write_bundle(path: Path, *, with_gemma_source: bool) -> None:
def _write_bundle(path: Path,
*,
with_gemma_source: bool,
transformer_cls: str = "AVTransformer3DModel") -> 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),
@@ -53,8 +57,9 @@ def _write_bundle(path: Path, *, with_gemma_source: bool) -> None:
"text_embedding_projection.video_aggregate_embed.weight": torch.zeros(2),
"duration_head.linear.weight": torch.zeros(2),
}
section = dict(TRANSFORMER_SECTION, _class_name=transformer_cls)
metadata = {
"config": json.dumps({"transformer": TRANSFORMER_SECTION, "vae": {}}),
"config": json.dumps({"transformer": section, "vae": {}}),
"model_version": "9.9.9",
"license": "x" * 128,
}
@@ -131,14 +136,6 @@ def test_defaults_unchanged_without_metadata() -> None:
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
@@ -172,3 +169,44 @@ def test_declared_root_survives_a_version_mismatch() -> None:
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)
def test_bundle_resolves_its_pipeline_config_from_declared_transformer_class() -> None:
"""A bundle path resolves through the alias table, not through model_index.json."""
from fastvideo.pipelines.basic.ltx2.pipeline_configs import LTX2T2VConfig
from fastvideo.registry import get_pipeline_config_cls_from_name
with tempfile.TemporaryDirectory() as tmp:
# Deliberately uninformative name: resolution must not read the file name.
path = Path(tmp) / "anonymous.safetensors"
_write_bundle(path, with_gemma_source=True)
assert get_pipeline_config_cls_from_name(str(path)) is LTX2T2VConfig
def test_unknown_bundle_transformer_class_raises_naming_table_and_override() -> None:
"""No wildcard fallback: an unmapped class fails loud and says how to fix it."""
import pytest
from fastvideo.registry import get_pipeline_config_cls_from_name
with tempfile.TemporaryDirectory() as tmp:
# Name it after a model the detectors DO know, to prove the file name
# is not consulted -- only the class the checkpoint declares.
path = Path(tmp) / "ltx2-distilled.safetensors"
_write_bundle(path, with_gemma_source=True, transformer_cls="NotAModel")
with pytest.raises(ValueError) as excinfo:
get_pipeline_config_cls_from_name(str(path))
message = str(excinfo.value)
assert "_BUNDLE_TRANSFORMER_TO_CONFIG" in message
assert "override_pipeline_cls_name" in message
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()
test_encoder_root_override_beats_config()
test_bundle_resolves_its_pipeline_config_from_declared_transformer_class()
test_unknown_bundle_transformer_class_raises_naming_table_and_override()
print("ok")