Spaces:
Running on Zero
Running on Zero
| """Model loading + DiT hyperparameter inference. | |
| Modules are constructed with class names preserved for direct | |
| ``load_state_dict`` compatibility, and DiT hyperparameters are inferred from | |
| the state dict itself (no model-config file). The returned dict keys are | |
| those the runner's model API consumes: dit, vae, text_encoder, tokenizer, | |
| plucker_adapter, scheduler, skeleton_embedding. | |
| The skeleton controller is chosen from ``skeleton_mode`` ("lightweight" -> | |
| LightweightSkeletonPatchEmbedding, anything else -> VacePatchEmbedding). | |
| """ | |
| import logging | |
| from pathlib import Path | |
| from typing import Any | |
| import torch | |
| logger = logging.getLogger(__name__) | |
| def build_module(cls, state_dict: dict[str, torch.Tensor], *, strict: bool = True, **ctor_kwargs): | |
| module = cls(**ctor_kwargs).to(torch.device("cpu")).eval() | |
| module.load_state_dict(state_dict, strict=strict) | |
| return module | |
| def load_state(path: Path) -> dict[str, Any]: | |
| state = torch.load(path, map_location="cpu", weights_only=False) | |
| if isinstance(state, dict) and "state_dict" in state and isinstance(state["state_dict"], dict): | |
| return state["state_dict"] | |
| if isinstance(state, dict): | |
| return state | |
| raise TypeError(f"Unsupported weight format: {type(state)} @ {path}") | |
| def infer_dit_hparams( | |
| dit_state_dict: dict[str, torch.Tensor], | |
| ) -> tuple[int, int, int, int, int, int, tuple[int, ...], int, int]: | |
| dim = int(dit_state_dict["patch_embedding.weight"].shape[0]) | |
| in_dim = int(dit_state_dict["patch_embedding.weight"].shape[1]) | |
| patch_size = tuple(int(x) for x in dit_state_dict["patch_embedding.weight"].shape[2:5]) | |
| freq_dim = int(dit_state_dict["time_embedding.0.weight"].shape[1]) | |
| text_dim = int(dit_state_dict["text_embedding.0.weight"].shape[1]) | |
| ffn_dim = int(dit_state_dict["blocks.0.ffn.0.weight"].shape[0]) | |
| out_dim = int( | |
| dit_state_dict["head.head.weight"].shape[0] | |
| // (patch_size[0] * patch_size[1] * patch_size[2]) | |
| ) | |
| num_layers = max(int(key.split(".")[1]) for key in dit_state_dict if key.startswith("blocks.")) + 1 | |
| num_heads = max(1, dim // 128) | |
| return dim, in_dim, ffn_dim, out_dim, text_dim, freq_dim, patch_size, num_heads, num_layers | |
| def load_models( | |
| weights_dir: str | Path, | |
| device: torch.device, | |
| dtype: torch.dtype, | |
| *, | |
| skeleton_mode: str = "lightweight", | |
| ) -> dict[str, Any]: | |
| """Load all model modules onto ``device``/``dtype`` from converted weights. | |
| Expected layout of ``weights_dir``: | |
| dit.pt vae.pt text_encoder.pt plucker_adapter.pt | |
| vace_patch_embedding.pt tokenizer/ (HF tokenizer directory) | |
| ``vae.pt`` is loaded with ``strict=False`` (VAE state dicts carry extra | |
| keys across checkpoints); everything else is strict. | |
| """ | |
| from uranus.modules import ( | |
| FlowMatchScheduler, | |
| HuggingfaceTokenizer, | |
| LightweightSkeletonPatchEmbedding, | |
| PluckerAdapter, | |
| SpatialTemporalWanModel, | |
| VacePatchEmbedding, | |
| WanTextEncoder, | |
| WanVideoVAE, | |
| ) | |
| weights_dir = Path(weights_dir).expanduser().resolve() | |
| if not weights_dir.is_dir(): | |
| raise FileNotFoundError(f"Converted weights directory not found: {weights_dir}") | |
| logger.info("Loading converted weights from %s", weights_dir) | |
| weights = { | |
| "dit": load_state(weights_dir / "dit.pt"), | |
| "vae": load_state(weights_dir / "vae.pt"), | |
| "text_encoder": load_state(weights_dir / "text_encoder.pt"), | |
| "plucker_adapter": load_state(weights_dir / "plucker_adapter.pt"), | |
| "vace_patch_embedding": load_state(weights_dir / "vace_patch_embedding.pt"), | |
| } | |
| tokenizer_path = weights_dir / "tokenizer" | |
| if not tokenizer_path.is_dir(): | |
| raise FileNotFoundError(f"Tokenizer directory not found under converted weights: {tokenizer_path}") | |
| use_lightweight_skeleton_controller = skeleton_mode == "lightweight" | |
| dim, in_dim, ffn_dim, out_dim, text_dim, freq_dim, patch_size, num_heads, num_layers = ( | |
| infer_dit_hparams(weights["dit"]) | |
| ) | |
| dit = build_module( | |
| SpatialTemporalWanModel, | |
| weights["dit"], | |
| strict=True, | |
| dim=dim, | |
| in_dim=in_dim, | |
| ffn_dim=ffn_dim, | |
| out_dim=out_dim, | |
| text_dim=text_dim, | |
| freq_dim=freq_dim, | |
| eps=1e-6, | |
| patch_size=patch_size, | |
| num_heads=num_heads, | |
| num_layers=num_layers, | |
| has_image_input=False, | |
| ) | |
| vae = build_module(WanVideoVAE, weights["vae"], strict=False, z_dim=16) | |
| text_encoder = build_module(WanTextEncoder, weights["text_encoder"], strict=True) | |
| tokenizer = HuggingfaceTokenizer(name=str(tokenizer_path), seq_len=512, clean="whitespace") | |
| plucker_adapter = build_module( | |
| PluckerAdapter, weights["plucker_adapter"], strict=True, in_dim=24, out_dim=dim | |
| ) | |
| if use_lightweight_skeleton_controller: | |
| skeleton_embedding = build_module( | |
| LightweightSkeletonPatchEmbedding, | |
| weights["vace_patch_embedding"], | |
| strict=True, | |
| in_dim=768, | |
| dim=dim, | |
| ) | |
| else: | |
| skeleton_embedding = build_module( | |
| VacePatchEmbedding, | |
| weights["vace_patch_embedding"], | |
| strict=True, | |
| in_dim=2 * vae.z_dim + 64, | |
| dim=dim, | |
| ) | |
| scheduler = FlowMatchScheduler() | |
| logger.info("Converted weights loaded successfully") | |
| return { | |
| "dit": dit.to(device=device, dtype=dtype).eval(), | |
| "vae": vae.to(device=device, dtype=dtype).eval(), | |
| "text_encoder": text_encoder.to(device=device, dtype=dtype).eval(), | |
| "tokenizer": tokenizer, | |
| "plucker_adapter": plucker_adapter.to(device=device, dtype=dtype).eval(), | |
| "scheduler": scheduler, | |
| "skeleton_embedding": skeleton_embedding.to(device=device, dtype=dtype).eval(), | |
| } | |