Spaces:
Running on Zero
Running on Zero
| """Weight plumbing for the ActionRoPE Space. | |
| The upstream ``actionrope.models.wan.load_pipeline`` loads the *base* Wan2.2-TI2V-5B | |
| DiT shards (3 files, ~20 GB) and then overwrites them with the ActionRoPE | |
| checkpoint. The ActionRoPE checkpoint is itself a full DiffSynth-native DiT state | |
| dict (827 tensors: blocks.* / head.* / patch_embedding / text_embedding / | |
| time_embedding / time_projection / arope_mask_embedding), so the base shards are | |
| redundant. This module builds ``diffsynth.models.wan_video_dit.WanModel`` directly | |
| from the registered TI2V-5B config and strict-loads the ActionRoPE checkpoint | |
| (exactly what upstream ``load_checkpoint`` + ``TrainTask.build_module`` do), loads | |
| the Wan2.2 VAE the same way, and returns a minimal pipeline object plus a | |
| ``Generator`` whose ``__call__`` is the upstream one — the paper's sampling path, | |
| unchanged. | |
| """ | |
| from __future__ import annotations | |
| import os | |
| def download_space_weights() -> tuple[str, str, str]: | |
| """Fetch the three runtime weight files (ActionRoPE DiT, its NULL-prompt table, the Wan2.2 VAE). | |
| Returns absolute (ckpt_path, null_prompt_path, vae_path). | |
| """ | |
| from huggingface_hub import hf_hub_download | |
| ckpt_path = hf_hub_download( | |
| "teawhite/ActionRoPE", "wan2.2-ti2v-5b/actionrope_cns.safetensors", repo_type="model" | |
| ) | |
| null_prompt_path = hf_hub_download( | |
| "teawhite/ActionRoPE", "wan2.2-ti2v-5b/umt5_null_prompt.pt", repo_type="model" | |
| ) | |
| vae_path = hf_hub_download("Wan-AI/Wan2.2-TI2V-5B", "Wan2.2_VAE.pth", repo_type="model") | |
| return ckpt_path, null_prompt_path, vae_path | |
| class _VAEPipeline: | |
| """BasePipeline holding only the VAE. | |
| Subclassing (rather than cherry-picking bound methods) keeps every helper | |
| ``wan.py`` reaches for — ``preprocess_image``, ``vae_output_to_video`` (which | |
| internally calls ``vae_output_to_image``), ``torch_dtype``, ``device``. | |
| """ | |
| def __init__(self, vae, device, torch_dtype): | |
| from diffsynth.diffusion.base_pipeline import BasePipeline | |
| class _Pipe(BasePipeline): | |
| def __init__(self): | |
| super().__init__(device=device, torch_dtype=torch_dtype) | |
| self.vae = vae | |
| self._pipe = _Pipe() | |
| # wan.py accesses pipe.vae / pipe.device / pipe.torch_dtype / helpers on the object itself | |
| self.vae = vae | |
| self.device = device | |
| self.torch_dtype = torch_dtype | |
| def __getattr__(self, name): | |
| # delegate any BasePipeline helper (preprocess_image, vae_output_to_*) to the real pipeline object | |
| return getattr(self.__dict__["_pipe"], name) | |
| def encode(self, videos, device, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): | |
| return self.vae.encode(videos, device=device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) | |
| def decode(self, hidden_states, device, tiled=False, tile_size=(30, 52), tile_stride=(15, 26)): | |
| return self.vae.decode(hidden_states, device=device, tiled=tiled, tile_size=tile_size, tile_stride=tile_stride) | |
| def build_pipe(ckpt_path: str, vae_path: str, device, torch_dtype) -> tuple: | |
| """Build (pipeline, dit, meta): DiT from the ActionRoPE ckpt, VAE from Wan2.2_VAE.pth. | |
| No base-DiT download, no text encoder — the DiT state dict is complete in the | |
| ActionRoPE checkpoint, and prompts come from the shipped ``umt5_null_prompt.pt`` | |
| embedding table. | |
| """ | |
| from safetensors import safe_open | |
| from diffsynth.core.loader.file import load_state_dict | |
| from diffsynth.models.wan_video_dit import WanModel | |
| from diffsynth.models.wan_video_vae import WanVideoVAE38 | |
| from diffsynth.utils.state_dict_converters.wan_video_vae import WanVideoVAEStateDictConverter | |
| from actionrope.checkpoint import read_metadata | |
| from actionrope.models.wan_dit import install_mask_channel | |
| # Registered DiffSynth config for Wan-AI/Wan2.2-TI2V-5B | |
| # (diffsynth/configs/model_configs.py, model_hash 1f5ab7703c6fc803fdded85ff040c316) | |
| ti2v_5b_kwargs = { | |
| "has_image_input": False, "patch_size": [1, 2, 2], "in_dim": 48, "dim": 3072, "ffn_dim": 14336, | |
| "freq_dim": 256, "text_dim": 4096, "out_dim": 48, "num_heads": 24, "num_layers": 30, "eps": 1e-06, | |
| "seperated_timestep": True, "require_clip_embedding": False, "require_vae_embedding": False, | |
| "fuse_vae_embedding_in_latents": True, | |
| } | |
| meta = read_metadata(ckpt_path) | |
| sd = {} | |
| with safe_open(ckpt_path, framework="pt", device="cpu") as f: | |
| for k in f.keys(): | |
| sd[k] = f.get_tensor(k) | |
| has_mask = any(k.startswith("arope_mask_embedding.") for k in sd) | |
| dit = WanModel(**ti2v_5b_kwargs) | |
| if has_mask: | |
| install_mask_channel(dit) | |
| dit.load_state_dict(sd, strict=True) # raises if the checkpoint is not the full TI2V-5B DiT | |
| del sd | |
| dit = dit.to(device=device, dtype=torch_dtype).eval().requires_grad_(False) | |
| # RoPE tables are plain tensors (not buffers): move them once (mirrors TrainTask.on_prepared) | |
| dit.freqs = tuple(t.to(device) for t in dit.freqs) | |
| vae_sd = WanVideoVAEStateDictConverter(load_state_dict(vae_path, torch_dtype=torch_dtype, device="cpu")) | |
| vae = WanVideoVAE38() | |
| vae.load_state_dict(vae_sd, strict=True) | |
| del vae_sd | |
| vae = vae.to(device=device, dtype=torch_dtype).eval().requires_grad_(False) | |
| vae.mean = vae.mean.to(device) | |
| vae.std = vae.std.to(device) | |
| print(f"[arope_loader] DiT {sum(p.numel() for p in dit.parameters()) / 1e9:.2f} B params, " | |
| f"mask_channel={has_mask}, metadata={dict(sorted(meta.items()))}") | |
| return _VAEPipeline(vae, device, torch_dtype), dit, meta | |
| class Generator: | |
| """Upstream ``actionrope.models.wan.Generator.__call__`` on a directly-built pipe. | |
| The upstream ``__init__`` calls ``load_pipeline(model_dir)`` which needs the base | |
| DiT shards; everything else in it is replicated here verbatim. | |
| """ | |
| def __init__(self, ckpt_path: str, vae_path: str, null_prompt_path: str, | |
| device="cuda", torch_dtype=None): | |
| import torch | |
| from actionrope.checkpoint import meta_bool, meta_float | |
| from actionrope.models.wan import SHIFT, PromptEmbeddings | |
| from actionrope.models.wan_dit import MASK_EMBEDDING | |
| torch_dtype = torch_dtype or torch.bfloat16 | |
| self.device = device | |
| self.pipe, self.dit, self.meta = build_pipe(ckpt_path, vae_path, device, torch_dtype) | |
| self.mask_channel = hasattr(self.dit, MASK_EMBEDDING) | |
| self.shift = meta_float(self.meta, "cns_shift", SHIFT) | |
| self.cns = ( | |
| { | |
| "shift_known": meta_float(self.meta, "cns_shift_known"), | |
| "alpha": meta_float(self.meta, "cns_alpha", 1.0), | |
| } | |
| if meta_bool(self.meta, "cns") | |
| else None | |
| ) | |
| self.prompts = PromptEmbeddings("", device, null_prompt_path) | |
| # The full paper sampling path, verbatim from actionrope.models.wan.Generator. | |
| from actionrope.models.wan import Generator as _Upstream # noqa: F401 (readability marker) | |
| def __call__(self, first_frame, frame_offset_px, prompt=None, seed=0, steps=50, shift=None, cns=None): | |
| import numpy as np | |
| from actionrope.models import wan as wan_module | |
| from actionrope.prompts import NULL_PROMPT | |
| z0 = wan_module.encode_first_frame(self.pipe, first_frame) | |
| context = self.prompts(prompt if prompt is not None else NULL_PROMPT) | |
| offset = wan_module.cell_offsets(frame_offset_px).unsqueeze(0).to(self.device) | |
| mask = wan_module.known_input(offset) if self.mask_channel else None | |
| use_cns = self.cns is not None if cns is None else bool(cns) | |
| if use_cns and self.cns is None: | |
| raise ValueError("cns=True needs a checkpoint trained with CNS") | |
| latents = wan_module.sample( | |
| self.dit, z0, context, offset, mask, | |
| steps=steps, shift=self.shift if shift is None else shift, seed=seed, | |
| cns=self.cns if use_cns else None, | |
| ) | |
| return wan_module.decode(self.pipe, latents), float(wan_module.FPS) |