From b2da3b2b4d1d992265b3b38c0d7da044aaeb800d Mon Sep 17 00:00:00 2001 From: SolitaryThinker Date: Wed, 17 Jun 2026 23:30:17 +0000 Subject: [PATCH] [feat] v2 wan21: backend-aware latent geometry + checkpoint stamping - latent_shape(req, model): real Wan geometry (16ch; (T-1)//4+1, H/8, W/8) on the cuda backend; toy stand-in stays on accel/cpu. - stamp_wan21_checkpoints(card, model_root): map a root (local dir or HF id) onto the 3 components' ComponentSpec.checkpoint; build_wan21_card(checkpoint_root=...) optional. --- v2/models/wan21/__init__.py | 5 +++-- v2/models/wan21/card.py | 28 ++++++++++++++++++++++++++-- v2/models/wan21/loop.py | 18 +++++++++++++++--- 3 files changed, 44 insertions(+), 7 deletions(-) diff --git a/v2/models/wan21/__init__.py b/v2/models/wan21/__init__.py index ae3a02bb1..79674a586 100644 --- a/v2/models/wan21/__init__.py +++ b/v2/models/wan21/__init__.py @@ -1,8 +1,9 @@ """Wan2.1-1.3B — bidirectional video diffusion (T2V).""" from __future__ import annotations -from .card import build_wan21_card +from .card import build_wan21_card, stamp_wan21_checkpoints from .loop import WanDenoiseLoop, latent_shape from .program import build_wan_t2v_program -__all__ = ["build_wan21_card", "build_wan_t2v_program", "WanDenoiseLoop", "latent_shape"] +__all__ = ["build_wan21_card", "stamp_wan21_checkpoints", "build_wan_t2v_program", + "WanDenoiseLoop", "latent_shape"] diff --git a/v2/models/wan21/card.py b/v2/models/wan21/card.py index 5172d4c97..4c7c80163 100644 --- a/v2/models/wan21/card.py +++ b/v2/models/wan21/card.py @@ -29,7 +29,8 @@ from ..backend import ToyDiT, ToyTextEncoder, ToyVAE, _seed_from from .loop import WanDenoiseLoop -def build_wan21_card(model_id: str = "wan2.1-1.3b", *, cfg_policy=None, flow_shift: float = 3.0) -> ModelCard: +def build_wan21_card(model_id: str = "wan2.1-1.3b", *, cfg_policy=None, flow_shift: float = 3.0, + checkpoint_root: str | None = None) -> ModelCard: seed = _seed_from(model_id) cost = CostModel(kind=WorkUnitKind.DIFFUSION_STEP, base_seconds=1e-4, per_unit_seconds=1e-7) cfg = cfg_policy or ClassicCFG() @@ -82,4 +83,27 @@ def build_wan21_card(model_id: str = "wan2.1-1.3b", *, cfg_policy=None, flow_shi ParallelPlan(axes={"sp": 2, "cfgp": 2}, mesh_order=["cfgp", "sp"])], default_plan=ParallelPlan.single()), ) - return card.validate() + card.validate() + if checkpoint_root: + stamp_wan21_checkpoints(card, checkpoint_root) + return card + + +# Wan2.1 diffusers checkpoint layout: each component's weights live in a subfolder of the model root. +_WAN21_SUBFOLDERS = {"transformer": "transformer", "vae": "vae", "text_encoder": "text_encoder"} + + +def stamp_wan21_checkpoints(card: ModelCard, model_root: str) -> ModelCard: + """GPU deploy-time: point each component's ``ComponentSpec.checkpoint`` at its weights subfolder + under ``model_root`` (a local diffusers dir, or an HF id resolved via ``snapshot_download``). The + toy cards leave ``checkpoint`` empty; the real torch adapters require it (BRINGUP risk A). The + adapter derives the pipeline-config model root from ``dirname(checkpoint)`` and loads the sibling + ``tokenizer/`` directory too. Mutates and returns the card.""" + import os + if not os.path.isdir(model_root): + from huggingface_hub import snapshot_download + model_root = snapshot_download(model_root) + for component_id, subfolder in _WAN21_SUBFOLDERS.items(): + if component_id in card.components: + card.components[component_id].checkpoint = os.path.join(model_root, subfolder) + return card diff --git a/v2/models/wan21/loop.py b/v2/models/wan21/loop.py index 3bbc569d3..3799ab31f 100644 --- a/v2/models/wan21/loop.py +++ b/v2/models/wan21/loop.py @@ -26,9 +26,21 @@ from ...request.streams import StreamChunk from ..backend import LATENT_CHANNELS -def latent_shape(req) -> tuple[int, int, int, int]: - """Map request geometry to a small (deterministic) latent for the CPU toy backend.""" +# Real Wan2.1 VAE (AutoencoderKLWan) compression: z_dim=16, 4x temporal, 8x spatial. +WAN_LATENT_CHANNELS = 16 +WAN_TEMPORAL_RATIO = 4 +WAN_SPATIAL_RATIO = 8 + + +def latent_shape(req, model=None) -> tuple[int, int, int, int]: + """Latent geometry for the denoise loop. On the real (GPU) backend the DiT/VAE require the true Wan + geometry (16 channels; (num_frames-1)//4+1 temporal, height/8, width/8); the CPU toy uses a tiny + deterministic stand-in. ``model is None`` (e.g. cost-estimation in tests) keeps the toy shape.""" d = req.diffusion + if model is not None and getattr(getattr(model, "platform", None), "device", "cpu") == "cuda": + t = (max(1, d.num_frames) - 1) // WAN_TEMPORAL_RATIO + 1 + return (WAN_LATENT_CHANNELS, max(1, t), + max(1, d.height // WAN_SPATIAL_RATIO), max(1, d.width // WAN_SPATIAL_RATIO)) t = max(1, d.num_frames // 40) h = max(2, d.height // 120) w = max(2, d.width // 120) @@ -49,7 +61,7 @@ class WanDenoiseLoop: rng = np.random.default_rng(seed) sig = self.flow_shift.build_schedule(req.diffusion.num_steps, req.diffusion.height, req.diffusion.width, sigmas=req.diffusion.sigmas or None) - shape = latent_shape(req) + shape = latent_shape(req, model) x = (rng.standard_normal(shape) * float(sig[0])).astype("float32") st = LoopState(loop_id=self.loop_id, instance_id=model.card.model_id, request_id=req.request_id, profile=ctx.profile, rng=rng, seed=seed,