diff --git a/examples/inference/basic/V2_PORTING_STATUS.md b/examples/inference/basic/V2_PORTING_STATUS.md index bae23567c..e50528d7b 100644 --- a/examples/inference/basic/V2_PORTING_STATUS.md +++ b/examples/inference/basic/V2_PORTING_STATUS.md @@ -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//`) 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_.py` mirroring the official example; generate to confirm. diff --git a/v2/registry.py b/v2/registry.py new file mode 100644 index 000000000..1fe43df29 --- /dev/null +++ b/v2/registry.py @@ -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)) diff --git a/v2/video_generator.py b/v2/video_generator.py index 0a55509d2..b1ca0c28a 100644 --- a/v2/video_generator.py +++ b/v2/video_generator.py @@ -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