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:
committed by
SolitaryThinker
parent
3f5614fd03
commit
6343eadbec
+63
-1
@@ -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:
|
||||
|
||||
@@ -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")
|
||||
|
||||
Reference in New Issue
Block a user