[refactor] v2: shared model registry (HF-id primary + arch fallback) for all entrypoints
Per review: dispatch should be a directly-mapped HF-string -> card registry (like fastvideo), shared by every entrypoint (VideoGenerator + a future CLI / server), not buried in the generator. - New v2/registry.py mirrors fastvideo's fastvideo/registry.py hybrid resolution: (1) exact HF repo id in an explicit ModelEntry registry [PRIMARY — correct per-model card/capabilities, and the only way to split same-architecture capability variants like Wan2.1 T2V vs the i2v 'InP' 1.3B], (2) short repo-name match, (3) architecture inference [FALLBACK — local paths / unregistered repos]. resolve(model_path[, root]) is the single shared entry point. - video_generator.py: moved _read_arch_signature/_select_builders into the registry; from_config now calls resolve() — registered ids resolve with no config read, else arch inference on a cheap *.json snapshot. Reconciles the earlier 'no brittle table' refactor with the 'map hf string -> card' ask: one clean registry + a fallback, not three coupled structures. - Verified: registry resolves exact-id / short-name / arch-fallback / unregistered correctly; CPU suite green (202+2); wan21 GPU smoke generates via the new resolve path.
This commit is contained in:
SolitaryThinker
committed by
Peiyuan Zhang
Peiyuan Zhang
parent
fa9c58b419
commit
7ac2ff0d1c
@@ -65,8 +65,10 @@ Each would need a new card + adapters (DiT/VAE/encoder forwards) + maybe a sampl
|
||||
|
||||
## How to add a model to the v2 VideoGenerator
|
||||
1. Ensure a v2 card exists (`v2/models/<family>/`) whose `load_id`s resolve to real `fastvideo.models.*`.
|
||||
2. Dispatch is automatic: add a branch in `v2/video_generator.py:_select_builders` keyed on the new
|
||||
architecture (transformer/pipeline/VAE class names) — no HF-id table to edit.
|
||||
2. Register it in the shared `v2/registry.py` (used by every entrypoint — VideoGenerator, CLI, server):
|
||||
add a `ModelEntry` mapping its HF id(s) → builders (the explicit PRIMARY path — the only way to split
|
||||
same-architecture capability variants, e.g. T2V vs the i2v "InP"), and/or a branch in
|
||||
`select_by_architecture` for the arch-inference FALLBACK (local paths / unregistered repos).
|
||||
3. If a DiT/VAE/encoder/upsampler forward differs, add an adapter + class-detection in `torch_adapters.py`
|
||||
and register the component kind in `torch_cuda.py`.
|
||||
4. Add `examples/inference/basic/v2_<name>.py` mirroring the official example; generate to confirm.
|
||||
|
||||
+154
@@ -0,0 +1,154 @@
|
||||
"""v2 model registry — the SHARED resolution of a model (HF repo id or local checkpoint dir) to its
|
||||
``(build_card, build_program)`` builders, so every v2 entrypoint (the typed ``VideoGenerator``, a CLI,
|
||||
the server) uses one source of truth instead of each re-deriving the mapping.
|
||||
|
||||
Mirrors fastvideo's ``fastvideo/registry.py`` hybrid resolution:
|
||||
|
||||
1. exact HF repo id in the explicit registry — PRIMARY: correct per-model card + capabilities, and
|
||||
the only way to split same-architecture *capability*
|
||||
variants (e.g. a Wan2.1 T2V base vs the i2v "InP"
|
||||
1.3B) that pure architecture inference cannot tell apart.
|
||||
2. short repo-name match — a renamed / forked copy of a known repo.
|
||||
3. architecture inference from the checkpoint — FALLBACK: local paths / unregistered repos, keyed on
|
||||
the pipeline / transformer / VAE class names.
|
||||
|
||||
Adding a model = one ``ModelEntry`` below (its HF ids -> builders). New architectures (for the fallback)
|
||||
are added in ``select_by_architecture``. All builders are CPU-clean v2 submodules imported lazily, so
|
||||
``import v2.registry`` stays cheap and torch-free.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Callable
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ModelEntry:
|
||||
"""A registered model family: the exact HF repo ids it serves + its card/program builders."""
|
||||
hf_ids: tuple[str, ...]
|
||||
build_card: Callable[[], Any]
|
||||
build_program: Callable[[], Any]
|
||||
|
||||
|
||||
def _entries() -> list[ModelEntry]:
|
||||
"""The explicit registry (PRIMARY). Builders imported lazily to avoid import cycles with the model
|
||||
packages and keep importing this module cheap."""
|
||||
from .models.ltx2 import (
|
||||
build_ltx2_base_card,
|
||||
build_ltx2_base_program,
|
||||
build_ltx2_card,
|
||||
build_ltx2_program,
|
||||
)
|
||||
from .models.wan_causal import build_wan_causal_card, build_wan_causal_program
|
||||
from .models.wan21 import (
|
||||
build_wan21_card,
|
||||
build_wan22_a14b_card,
|
||||
build_wan22_ti2v_card,
|
||||
build_wan_t2v_program,
|
||||
)
|
||||
return [
|
||||
ModelEntry(("Wan-AI/Wan2.1-T2V-1.3B-Diffusers",), build_wan21_card, build_wan_t2v_program),
|
||||
ModelEntry(("Wan-AI/Wan2.2-TI2V-5B-Diffusers",), build_wan22_ti2v_card, build_wan_t2v_program),
|
||||
ModelEntry(("Wan-AI/Wan2.2-T2V-A14B-Diffusers",), build_wan22_a14b_card, build_wan_t2v_program),
|
||||
ModelEntry(("wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers",), build_wan_causal_card, build_wan_causal_program),
|
||||
# LTX-2 two-stage distilled (ships a spatial_upsampler: base -> upsample -> refine):
|
||||
ModelEntry(("FastVideo/LTX2-Distilled-Diffusers",), build_ltx2_card, build_ltx2_program),
|
||||
# LTX-2 single-stage (no upsampler) — the non-distilled base (many-step) AND the 2.3-distilled
|
||||
# (few-step), which is also single-stage; pass a small num_inference_steps for the distilled one:
|
||||
ModelEntry(("Davids048/LTX2-Base-Diffusers", "FastVideo/LTX-2.3-Distilled-Diffusers"),
|
||||
build_ltx2_base_card, build_ltx2_base_program),
|
||||
]
|
||||
|
||||
|
||||
def _short(p: str) -> str:
|
||||
return p.rstrip("/").split("/")[-1].lower()
|
||||
|
||||
|
||||
def read_arch_signature(root: str) -> dict:
|
||||
"""Read a diffusers checkpoint's configs into an architecture signature (the FALLBACK dispatch
|
||||
keys): pipeline / transformer / VAE class names + the fields that split same-class variants
|
||||
(VAE ``z_dim``, a second transformer / ``boundary_ratio`` for MoE, a ``spatial_upsampler``).
|
||||
``root`` is a local diffusers dir or a snapshot of just the ``*.json`` configs."""
|
||||
def _load(*parts: str) -> dict:
|
||||
p = os.path.join(root, *parts)
|
||||
if os.path.exists(p):
|
||||
with open(p) as f:
|
||||
return json.load(f)
|
||||
return {}
|
||||
|
||||
mi = _load("model_index.json")
|
||||
tcfg = _load("transformer", "config.json")
|
||||
vcfg = _load("vae", "config.json")
|
||||
return {
|
||||
"pipeline": mi.get("_class_name"),
|
||||
"boundary_ratio": mi.get("boundary_ratio"),
|
||||
"has_transformer_2": os.path.isdir(os.path.join(root, "transformer_2")),
|
||||
"has_spatial_upsampler": os.path.isdir(os.path.join(root, "spatial_upsampler")) or "spatial_upsampler" in mi,
|
||||
"transformer_cls": tcfg.get("_class_name"),
|
||||
"in_channels": tcfg.get("in_channels"),
|
||||
"vae_z_dim": vcfg.get("z_dim", vcfg.get("latent_channels")),
|
||||
}
|
||||
|
||||
|
||||
def select_by_architecture(sig: dict):
|
||||
"""FALLBACK: map an architecture signature -> (build_card, build_program) by class names, so a local
|
||||
path / renamed repo / new distilled variant of a known arch still resolves with no registry entry."""
|
||||
tr, pipe = sig.get("transformer_cls"), sig.get("pipeline")
|
||||
if tr == "LTX2Transformer3DModel":
|
||||
from .models.ltx2 import (
|
||||
build_ltx2_base_card,
|
||||
build_ltx2_base_program,
|
||||
build_ltx2_card,
|
||||
build_ltx2_program,
|
||||
)
|
||||
if sig.get("has_spatial_upsampler"):
|
||||
return build_ltx2_card, build_ltx2_program # two-stage (base -> upsample -> refine)
|
||||
return build_ltx2_base_card, build_ltx2_base_program # single-stage
|
||||
if tr == "CausalWanTransformer3DModel":
|
||||
from .models.wan_causal import build_wan_causal_card, build_wan_causal_program
|
||||
return build_wan_causal_card, build_wan_causal_program
|
||||
if tr == "WanTransformer3DModel":
|
||||
from .models.wan21 import (
|
||||
build_wan21_card,
|
||||
build_wan22_a14b_card,
|
||||
build_wan22_ti2v_card,
|
||||
build_wan_t2v_program,
|
||||
)
|
||||
if pipe == "WanDMDPipeline": # FastWan: detected by pipeline class -> precise, not a load crash
|
||||
raise ValueError(
|
||||
"v2 registry: WanDMD/FastWan is not supported via the generic Wan path — its checkpoint's "
|
||||
"to_gate_compress param mapping differs from the generic WanTransformer3DModel load. See "
|
||||
"examples/inference/basic/V2_PORTING_STATUS.md.")
|
||||
if sig.get("has_transformer_2") or sig.get("boundary_ratio"):
|
||||
return build_wan22_a14b_card, build_wan_t2v_program # Wan2.2 MoE (two experts)
|
||||
if sig.get("vae_z_dim") == 48:
|
||||
return build_wan22_ti2v_card, build_wan_t2v_program # Wan2.2-TI2V-5B (z_dim=48 VAE)
|
||||
return build_wan21_card, build_wan_t2v_program # Wan2.1
|
||||
raise ValueError(
|
||||
f"v2 registry: unsupported architecture (transformer={tr!r}, pipeline={pipe!r}). Supported "
|
||||
f"transformers: WanTransformer3DModel / CausalWanTransformer3DModel / LTX2Transformer3DModel. "
|
||||
f"See examples/inference/basic/V2_PORTING_STATUS.md.")
|
||||
|
||||
|
||||
def resolve(model_path: str, root: str | None = None):
|
||||
"""Resolve ``model_path`` (an HF repo id or a local checkpoint dir) -> ``(build_card, build_program)``.
|
||||
|
||||
PRIMARY: an exact HF-id, then a short repo-name match, in the explicit registry. FALLBACK: when
|
||||
``model_path`` isn't registered, architecture inference from ``root`` (a local dir or a downloaded
|
||||
``*.json`` config snapshot). Raises ``ValueError`` if unregistered with ``root=None`` (so the caller
|
||||
knows to fetch the configs), or if the architecture is unsupported. Shared by every v2 entrypoint."""
|
||||
entries = _entries()
|
||||
for e in entries: # 1. exact HF id
|
||||
if model_path in e.hf_ids:
|
||||
return e.build_card, e.build_program
|
||||
short = _short(model_path) # 2. short repo name (renamed / forked copy)
|
||||
for e in entries:
|
||||
if any(_short(h) == short for h in e.hf_ids):
|
||||
return e.build_card, e.build_program
|
||||
if root is None: # 3. architecture inference (needs the configs)
|
||||
raise ValueError(
|
||||
f"v2 registry: {model_path!r} is not registered; pass the checkpoint root for architecture "
|
||||
f"inference, or add a ModelEntry in v2/registry.py.")
|
||||
return select_by_architecture(read_arch_signature(root))
|
||||
+21
-83
@@ -9,85 +9,19 @@ It consumes the OFFICIAL ``fastvideo.api`` typed configs (``GeneratorConfig`` /
|
||||
|
||||
All torch / fastvideo / huggingface imports are lazy (inside methods) so ``import v2`` stays CPU-clean.
|
||||
|
||||
Model dispatch is registry-style: ``from_config`` reads the checkpoint's architecture (the pipeline /
|
||||
transformer / VAE class names) and picks the v2 card — no hardcoded repo-id table (mirrors fastvideo's
|
||||
``get_pipeline_config_cls_from_name``). Wired families: Wan2.1, Wan2.2-TI2V-5B, Wan2.2-A14B (MoE),
|
||||
SF-causal Wan, LTX-2. Multi-GPU / CPU-FSDP offload / VSA are accepted for API parity but the bring-up
|
||||
runs single-GPU, resident, on the TORCH_SDPA backend (no fastvideo-kernel) — see GPU_BRINGUP.md.
|
||||
Model dispatch lives in the SHARED ``v2/registry.py`` (so a CLI / server resolve identically): an
|
||||
explicit HF-id registry PRIMARY + architecture inference FALLBACK, mirroring fastvideo's
|
||||
``fastvideo/registry.py``. Wired: Wan2.1, Wan2.2-TI2V-5B, Wan2.2-A14B (MoE), SF-causal Wan, LTX-2
|
||||
(two-stage distilled + single-stage base/2.3). Multi-GPU / CPU-FSDP offload / VSA are accepted for API
|
||||
parity but the bring-up runs single-GPU, resident, on TORCH_SDPA (no fastvideo-kernel) — see GPU_BRINGUP.md.
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import Any
|
||||
|
||||
def _read_arch_signature(root: str) -> dict:
|
||||
"""Read a diffusers checkpoint's configs into an architecture signature. Dispatch keys are the
|
||||
*class names* the model declares (pipeline / transformer / VAE) plus the few config fields that
|
||||
distinguish same-class variants (VAE ``z_dim``, a second transformer / ``boundary_ratio`` for MoE)
|
||||
— the same information fastvideo's registry (``get_pipeline_config_cls_from_name``) keys on, not
|
||||
the repo id. ``root`` is a local diffusers dir (or a snapshot of just the ``*.json`` configs)."""
|
||||
import json
|
||||
import os
|
||||
|
||||
def _load(*parts: str) -> dict:
|
||||
p = os.path.join(root, *parts)
|
||||
if os.path.exists(p):
|
||||
with open(p) as f:
|
||||
return json.load(f)
|
||||
return {}
|
||||
|
||||
mi, tcfg, vcfg = _load("model_index.json"), _load("transformer", "config.json"), _load("vae", "config.json")
|
||||
return {
|
||||
"pipeline": mi.get("_class_name"),
|
||||
"boundary_ratio": mi.get("boundary_ratio"),
|
||||
"has_transformer_2": os.path.isdir(os.path.join(root, "transformer_2")),
|
||||
"has_spatial_upsampler": os.path.isdir(os.path.join(root, "spatial_upsampler")) or "spatial_upsampler" in mi,
|
||||
"transformer_cls": tcfg.get("_class_name"),
|
||||
"in_channels": tcfg.get("in_channels"),
|
||||
"vae_z_dim": vcfg.get("z_dim", vcfg.get("latent_channels")),
|
||||
}
|
||||
|
||||
|
||||
def _select_builders(sig: dict):
|
||||
"""Map an architecture signature -> (build_card, build_program). Which v2 card a checkpoint uses is
|
||||
determined by its transformer/pipeline/VAE classes — exactly like fastvideo's registry — so a local
|
||||
path, a renamed repo, or a new distilled variant of a known arch all resolve correctly with no
|
||||
hardcoded HF-id table. New families are added here by class, next to the card that handles them."""
|
||||
tr, pipe = sig.get("transformer_cls"), sig.get("pipeline")
|
||||
if tr == "LTX2Transformer3DModel":
|
||||
from .models.ltx2 import (
|
||||
build_ltx2_base_card,
|
||||
build_ltx2_base_program,
|
||||
build_ltx2_card,
|
||||
build_ltx2_program,
|
||||
)
|
||||
if sig.get("has_spatial_upsampler"):
|
||||
return build_ltx2_card, build_ltx2_program # distilled two-stage (base→upsample→refine)
|
||||
return build_ltx2_base_card, build_ltx2_base_program # single-stage base model
|
||||
if tr == "CausalWanTransformer3DModel":
|
||||
from .models.wan_causal import build_wan_causal_card, build_wan_causal_program
|
||||
return build_wan_causal_card, build_wan_causal_program
|
||||
if tr == "WanTransformer3DModel":
|
||||
from .models.wan21 import (
|
||||
build_wan21_card,
|
||||
build_wan22_a14b_card,
|
||||
build_wan22_ti2v_card,
|
||||
build_wan_t2v_program,
|
||||
)
|
||||
if pipe == "WanDMDPipeline": # FastWan: detected by pipeline class -> precise, not a load crash
|
||||
raise ValueError(
|
||||
"v2 VideoGenerator: WanDMD/FastWan is not supported via the generic Wan path — its "
|
||||
"checkpoint's to_gate_compress param mapping differs from the generic WanTransformer3DModel "
|
||||
"load. See examples/inference/basic/V2_PORTING_STATUS.md.")
|
||||
if sig.get("has_transformer_2") or sig.get("boundary_ratio"):
|
||||
return build_wan22_a14b_card, build_wan_t2v_program # Wan2.2 MoE (two experts)
|
||||
if sig.get("vae_z_dim") == 48:
|
||||
return build_wan22_ti2v_card, build_wan_t2v_program # Wan2.2-TI2V-5B (z_dim=48 VAE)
|
||||
return build_wan21_card, build_wan_t2v_program # Wan2.1
|
||||
raise ValueError(
|
||||
f"v2 VideoGenerator: unsupported architecture (transformer={tr!r}, pipeline={pipe!r}). Supported "
|
||||
f"transformers: WanTransformer3DModel / CausalWanTransformer3DModel / LTX2Transformer3DModel. "
|
||||
f"See examples/inference/basic/V2_PORTING_STATUS.md.")
|
||||
# Model dispatch (HF-id registry primary + architecture-inference fallback) lives in the shared
|
||||
# ``v2/registry.py`` so every entrypoint (this VideoGenerator, a CLI, the server) resolves identically.
|
||||
|
||||
|
||||
class VideoGenerator:
|
||||
@@ -130,17 +64,21 @@ class VideoGenerator:
|
||||
warnings.warn("v2 VideoGenerator bring-up is single-GPU; ignoring num_gpus>1.", stacklevel=2)
|
||||
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
from .registry import resolve
|
||||
rev = getattr(config, "revision", None)
|
||||
# Registry-style dispatch on the checkpoint's architecture (class names), not the repo id. Read
|
||||
# just the configs first (a local dir, or a cheap ``*.json`` snapshot) so an unsupported arch is
|
||||
# rejected BEFORE the full (possibly 100GB+) weight download; then fetch the weights (cached).
|
||||
if os.path.isdir(model_path):
|
||||
root = model_path
|
||||
else:
|
||||
root = snapshot_download(model_path, revision=rev, allow_patterns=["*.json", "**/*.json"])
|
||||
build_card, build_program = _select_builders(_read_arch_signature(root))
|
||||
if not os.path.isdir(model_path):
|
||||
root = snapshot_download(model_path, revision=rev) # full weights (idempotent, reuses cache)
|
||||
local = os.path.isdir(model_path)
|
||||
# Resolve card+program via the shared registry (v2/registry.py): a registered HF-id resolves with
|
||||
# no download; an unregistered id / local path falls back to architecture inference, which needs
|
||||
# the configs (the local dir, or a cheap ``*.json`` snapshot fetched BEFORE the full — possibly
|
||||
# 100GB+ — weight download, so an unsupported arch is rejected early).
|
||||
try:
|
||||
build_card, build_program = resolve(model_path)
|
||||
except ValueError:
|
||||
cfg_root = model_path if local else snapshot_download(
|
||||
model_path, revision=rev, allow_patterns=["*.json", "**/*.json"])
|
||||
build_card, build_program = resolve(model_path, cfg_root)
|
||||
root = model_path if local else snapshot_download(model_path, revision=rev) # full weights (cached)
|
||||
card, program = build_card(), build_program()
|
||||
|
||||
from .models.wan21 import stamp_wan21_checkpoints # transformer/vae/text_encoder subfolder layout
|
||||
|
||||
Reference in New Issue
Block a user