Add pre-quantized NF4 weights and complete inference package

- NF4 quantized diffusion models (high_noise + low_noise) via bitsandbytes
- T5-XXL text encoder (models_t5_umt5-xxl-enc-bf16.pth)
- VAE encoder/decoder (Wan2.1_VAE.pth)
- Tokenizer files from google/umt5-xxl
- Complete wan/ module for inference
- generate_prequant.py - main inference script
- load_prequant.py - pre-quantized weight loader

Total package: ~30GB (vs ~85GB for full precision)
Fits in 32GB VRAM with model swapping between GPU/CPU

.gitattributes

@@ -3,9 +3,9 @@
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-
 # Keep these as text
 *.json text
 *.md text
 *.py text
 *.txt text
+tokenizer/*.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,113 @@
+# LingBot-World NF4 Quantized
+
+Pre-quantized NF4 weights for LingBot-World video generation model. This is a complete, self-contained package - no additional downloads required.
+
+## Features
+
+- **4-bit NF4 quantization** via bitsandbytes - fits in 32GB VRAM
+- **Pre-quantized weights** - no runtime quantization overhead
+- **Complete package** - includes T5 encoder, VAE, and diffusion models
+
+## Quick Start
+
+```bash
+# Clone the repo
+git clone https://huggingface.co/cahlen/lingbot-world-base-cam-nf4
+cd lingbot-world-base-cam-nf4
+
+# Install dependencies
+pip install -r requirements.txt
+
+# Generate a video
+python generate_prequant.py \
+    --image your_image.jpg \
+    --prompt "A cinematic video of the scene" \
+    --frame_num 81 \
+    --output output.mp4
+```
+
+## Model Contents
+
+| File | Size | Description |
+|------|------|-------------|
+| `high_noise_model_bnb_nf4/model.safetensors` | ~9.6GB | NF4 quantized diffusion model (high noise) |
+| `low_noise_model_bnb_nf4/model.safetensors` | ~9.6GB | NF4 quantized diffusion model (low noise) |
+| `models_t5_umt5-xxl-enc-bf16.pth` | ~10.6GB | T5-XXL text encoder |
+| `Wan2.1_VAE.pth` | ~485MB | VAE encoder/decoder |
+
+**Total size: ~30GB** (vs ~85GB for full precision models)
+
+## Usage
+
+### Basic Generation
+
+```bash
+python generate_prequant.py \
+    --image input.jpg \
+    --prompt "Your prompt here" \
+    --frame_num 81 \
+    --size "480*832" \
+    --output output.mp4
+```
+
+### Parameters
+
+| Parameter | Default | Description |
+|-----------|---------|-------------|
+| `--image` | required | Input image path |
+| `--prompt` | required | Text prompt describing the video |
+| `--frame_num` | 81 | Number of frames (81 = ~5 seconds at 16fps) |
+| `--size` | "480*832" | Output resolution (height*width) |
+| `--sampling_steps` | 40 | Diffusion sampling steps |
+| `--guide_scale` | 5.0 | Classifier-free guidance scale |
+| `--seed` | -1 | Random seed (-1 for random) |
+| `--output` | "output.mp4" | Output video path |
+
+### With Camera Control
+
+```bash
+python generate_prequant.py \
+    --image input.jpg \
+    --prompt "Your prompt" \
+    --action_path /path/to/camera_poses/ \
+    --frame_num 81
+```
+
+Camera pose directory should contain:
+- `poses.npy`: Shape `[num_frames, 4, 4]` - camera transformation matrices
+- `intrinsics.npy`: Shape `[num_frames, 4]` - `[fx, fy, cx, cy]`
+
+## Requirements
+
+- Python 3.10+
+- CUDA 11.8+ (tested with CUDA 12.x)
+- ~32GB VRAM (RTX 4090, RTX 5090, A100, etc.)
+
+## Quantization Details
+
+The diffusion models are quantized using bitsandbytes NF4 with double quantization:
+
+```json
+{
+  "format": "bnb_nf4",
+  "double_quant": true,
+  "compute_dtype": "bfloat16",
+  "blocksize": 64
+}
+```
+
+This achieves ~3.9x compression while maintaining generation quality.
+
+## License
+
+This model is based on [LingBot-World](https://github.com/robbyant/lingbot-world) and follows its license terms.
+
+## Citation
+
+```bibtex
+@misc{lingbot-world-nf4,
+  title={LingBot-World NF4 Quantized},
+  year={2025},
+  url={https://huggingface.co/cahlen/lingbot-world-base-cam-nf4}
+}
+```

Wan2.1_VAE.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:38071ab59bd94681c686fa51d75a1968f64e470262043be31f7a094e442fd981
+size 507609880

generate_prequant.py

@@ -0,0 +1,425 @@
+#!/usr/bin/env python3
+"""
+Generate videos using PRE-QUANTIZED bitsandbytes NF4 models.
+
+Unlike generate_bnb.py which re-quantizes at runtime, this script loads
+pre-quantized weights directly. No base model weights are needed.
+
+Prerequisites:
+    - Pre-quantized models in {ckpt_dir}/{high,low}_noise_model_bnb_nf4/
+    - Each should contain model.safetensors (or model.pt) + config.json
+
+Usage:
+    python generate_prequant.py \
+        --image examples/00/image.jpg \
+        --prompt "A cinematic video of the scene" \
+        --frame_num 81 \
+        --size 480*832
+"""
+
+import argparse
+import gc
+import logging
+import os
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torchvision.transforms.functional as TF
+from PIL import Image
+from tqdm import tqdm
+
+sys.path.insert(0, str(Path(__file__).parent))
+
+from einops import rearrange
+
+from load_prequant import load_quantized_model
+from wan.configs.wan_i2v_A14B import i2v_A14B as cfg
+from wan.modules.t5 import T5EncoderModel
+from wan.modules.vae2_1 import Wan2_1_VAE
+from wan.utils.cam_utils import (
+    compute_relative_poses,
+    get_Ks_transformed,
+    get_plucker_embeddings,
+    interpolate_camera_poses,
+)
+from wan.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class WanI2V_PreQuant:
+    """Image-to-video pipeline using pre-quantized NF4 models."""
+
+    def __init__(
+        self,
+        checkpoint_dir: str,
+        device_id: int = 0,
+        t5_cpu: bool = True,
+    ):
+        self.device = torch.device(f"cuda:{device_id}")
+        self.config = cfg
+        self.t5_cpu = t5_cpu
+
+        self.num_train_timesteps = cfg.num_train_timesteps
+        self.boundary = cfg.boundary
+        self.param_dtype = cfg.param_dtype
+        self.vae_stride = cfg.vae_stride
+        self.patch_size = cfg.patch_size
+        self.sample_neg_prompt = cfg.sample_neg_prompt
+
+        # Load T5 encoder (not quantized)
+        logger.info("Loading T5 encoder...")
+        # Use local tokenizer if available, otherwise fall back to HuggingFace
+        local_tokenizer = os.path.join(checkpoint_dir, "tokenizer")
+        tokenizer_path = local_tokenizer if os.path.isdir(local_tokenizer) else cfg.t5_tokenizer
+        self.text_encoder = T5EncoderModel(
+            text_len=cfg.text_len,
+            dtype=cfg.t5_dtype,
+            device=torch.device("cpu"),
+            checkpoint_path=os.path.join(checkpoint_dir, cfg.t5_checkpoint),
+            tokenizer_path=tokenizer_path,
+            shard_fn=None,
+        )
+
+        # Load VAE (not quantized)
+        logger.info("Loading VAE...")
+        self.vae = Wan2_1_VAE(
+            vae_pth=os.path.join(checkpoint_dir, cfg.vae_checkpoint),
+            device=self.device,
+        )
+
+        # Load PRE-QUANTIZED diffusion models
+        logger.info("Loading pre-quantized NF4 diffusion models...")
+
+        low_noise_dir = os.path.join(
+            checkpoint_dir, cfg.low_noise_checkpoint + "_bnb_nf4"
+        )
+        high_noise_dir = os.path.join(
+            checkpoint_dir, cfg.high_noise_checkpoint + "_bnb_nf4"
+        )
+
+        # Verify directories exist
+        for d in [low_noise_dir, high_noise_dir]:
+            if not os.path.isdir(d):
+                raise FileNotFoundError(
+                    f"Pre-quantized model not found: {d}\n"
+                    "Run: python scripts/quantize_and_package.py first"
+                )
+
+        # Load to CPU first, we'll swap to GPU as needed
+        self.low_noise_model = load_quantized_model(low_noise_dir, device="cpu")
+        self.high_noise_model = load_quantized_model(high_noise_dir, device="cpu")
+
+        logger.info("Model loading complete!")
+
+    def _prepare_model_for_timestep(self, t, boundary):
+        """Prepare and return the required model for the current timestep."""
+        if t.item() >= boundary:
+            required_model_name = "high_noise_model"
+            offload_model_name = "low_noise_model"
+        else:
+            required_model_name = "low_noise_model"
+            offload_model_name = "high_noise_model"
+
+        required_model = getattr(self, required_model_name)
+        offload_model = getattr(self, offload_model_name)
+
+        # Offload unused model to CPU
+        try:
+            if next(offload_model.parameters()).device.type == "cuda":
+                offload_model.to("cpu")
+                torch.cuda.empty_cache()
+        except StopIteration:
+            pass
+
+        # Load required model to GPU
+        try:
+            if next(required_model.parameters()).device.type == "cpu":
+                required_model.to(self.device)
+        except StopIteration:
+            pass
+
+        return required_model
+
+    def generate(
+        self,
+        input_prompt: str,
+        img: Image.Image,
+        action_path: str = None,
+        max_area: int = 720 * 1280,
+        frame_num: int = 81,
+        shift: float = 5.0,
+        sampling_steps: int = 40,
+        guide_scale: float = 5.0,
+        n_prompt: str = "",
+        seed: int = -1,
+    ):
+        """Generate video from image and text prompt."""
+        if action_path is not None:
+            c2ws = np.load(os.path.join(action_path, "poses.npy"))
+            len_c2ws = ((len(c2ws) - 1) // 4) * 4 + 1
+            frame_num = min(frame_num, len_c2ws)
+            c2ws = c2ws[:frame_num]
+
+        guide_scale = (
+            (guide_scale, guide_scale)
+            if isinstance(guide_scale, float)
+            else guide_scale
+        )
+        img_tensor = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device)
+
+        F = frame_num
+        h, w = img_tensor.shape[1:]
+        aspect_ratio = h / w
+        lat_h = round(
+            np.sqrt(max_area * aspect_ratio)
+            // self.vae_stride[1]
+            // self.patch_size[1]
+            * self.patch_size[1]
+        )
+        lat_w = round(
+            np.sqrt(max_area / aspect_ratio)
+            // self.vae_stride[2]
+            // self.patch_size[2]
+            * self.patch_size[2]
+        )
+        h = lat_h * self.vae_stride[1]
+        w = lat_w * self.vae_stride[2]
+        lat_f = (F - 1) // self.vae_stride[0] + 1
+        max_seq_len = (
+            lat_f * lat_h * lat_w // (self.patch_size[1] * self.patch_size[2])
+        )
+
+        seed = seed if seed >= 0 else random.randint(0, sys.maxsize)
+        seed_g = torch.Generator(device=self.device)
+        seed_g.manual_seed(seed)
+        noise = torch.randn(
+            16,
+            (F - 1) // self.vae_stride[0] + 1,
+            lat_h,
+            lat_w,
+            dtype=torch.float32,
+            generator=seed_g,
+            device=self.device,
+        )
+
+        msk = torch.ones(1, F, lat_h, lat_w, device=self.device)
+        msk[:, 1:] = 0
+        msk = torch.concat(
+            [torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1
+        )
+        msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)
+        msk = msk.transpose(1, 2)[0]
+
+        if n_prompt == "":
+            n_prompt = self.sample_neg_prompt
+
+        # Encode text
+        if not self.t5_cpu:
+            self.text_encoder.model.to(self.device)
+            context = self.text_encoder([input_prompt], self.device)
+            context_null = self.text_encoder([n_prompt], self.device)
+            self.text_encoder.model.cpu()
+        else:
+            context = self.text_encoder([input_prompt], torch.device("cpu"))
+            context_null = self.text_encoder([n_prompt], torch.device("cpu"))
+            context = [t.to(self.device) for t in context]
+            context_null = [t.to(self.device) for t in context_null]
+
+        # Camera preparation
+        dit_cond_dict = None
+        if action_path is not None:
+            Ks = torch.from_numpy(
+                np.load(os.path.join(action_path, "intrinsics.npy"))
+            ).float()
+            Ks = get_Ks_transformed(Ks, 480, 832, h, w, h, w)
+            Ks = Ks[0]
+
+            len_c2ws = len(c2ws)
+            c2ws_infer = interpolate_camera_poses(
+                src_indices=np.linspace(0, len_c2ws - 1, len_c2ws),
+                src_rot_mat=c2ws[:, :3, :3],
+                src_trans_vec=c2ws[:, :3, 3],
+                tgt_indices=np.linspace(
+                    0, len_c2ws - 1, int((len_c2ws - 1) // 4) + 1
+                ),
+            )
+            c2ws_infer = compute_relative_poses(c2ws_infer, framewise=True)
+            Ks = Ks.repeat(len(c2ws_infer), 1)
+
+            c2ws_infer = c2ws_infer.to(self.device)
+            Ks = Ks.to(self.device)
+            c2ws_plucker_emb = get_plucker_embeddings(c2ws_infer, Ks, h, w)
+            c2ws_plucker_emb = rearrange(
+                c2ws_plucker_emb,
+                "f (h c1) (w c2) c -> (f h w) (c c1 c2)",
+                c1=int(h // lat_h),
+                c2=int(w // lat_w),
+            )
+            c2ws_plucker_emb = c2ws_plucker_emb[None, ...]
+            c2ws_plucker_emb = rearrange(
+                c2ws_plucker_emb,
+                "b (f h w) c -> b c f h w",
+                f=lat_f,
+                h=lat_h,
+                w=lat_w,
+            ).to(self.param_dtype)
+            dit_cond_dict = {"c2ws_plucker_emb": c2ws_plucker_emb.chunk(1, dim=0)}
+
+        # Encode image
+        y = self.vae.encode(
+            [
+                torch.concat(
+                    [
+                        torch.nn.functional.interpolate(
+                            img_tensor[None].cpu(), size=(h, w), mode="bicubic"
+                        ).transpose(0, 1),
+                        torch.zeros(3, F - 1, h, w),
+                    ],
+                    dim=1,
+                ).to(self.device)
+            ]
+        )[0]
+        y = torch.concat([msk, y])
+
+        # Diffusion sampling
+        with torch.amp.autocast("cuda", dtype=self.param_dtype), torch.no_grad():
+            boundary = self.boundary * self.num_train_timesteps
+
+            sample_scheduler = FlowUniPCMultistepScheduler(
+                num_train_timesteps=self.num_train_timesteps,
+                shift=1,
+                use_dynamic_shifting=False,
+            )
+            sample_scheduler.set_timesteps(sampling_steps, device=self.device, shift=shift)
+            timesteps = sample_scheduler.timesteps
+
+            latent = noise
+
+            arg_c = {
+                "context": [context[0]],
+                "seq_len": max_seq_len,
+                "y": [y],
+                "dit_cond_dict": dit_cond_dict,
+            }
+
+            arg_null = {
+                "context": context_null,
+                "seq_len": max_seq_len,
+                "y": [y],
+                "dit_cond_dict": dit_cond_dict,
+            }
+
+            torch.cuda.empty_cache()
+
+            # Pre-load first model
+            first_model_name = (
+                "high_noise_model" if timesteps[0].item() >= boundary else "low_noise_model"
+            )
+            getattr(self, first_model_name).to(self.device)
+            logger.info(f"Loaded {first_model_name} to GPU")
+
+            for _, t in enumerate(tqdm(timesteps, desc="Sampling")):
+                latent_model_input = [latent.to(self.device)]
+                timestep = torch.stack([t]).to(self.device)
+
+                model = self._prepare_model_for_timestep(t, boundary)
+                sample_guide_scale = (
+                    guide_scale[1] if t.item() >= boundary else guide_scale[0]
+                )
+
+                noise_pred_cond = model(latent_model_input, t=timestep, **arg_c)[0]
+                torch.cuda.empty_cache()
+                noise_pred_uncond = model(latent_model_input, t=timestep, **arg_null)[0]
+                torch.cuda.empty_cache()
+                noise_pred = noise_pred_uncond + sample_guide_scale * (
+                    noise_pred_cond - noise_pred_uncond
+                )
+
+                temp_x0 = sample_scheduler.step(
+                    noise_pred.unsqueeze(0),
+                    t,
+                    latent.unsqueeze(0),
+                    return_dict=False,
+                    generator=seed_g,
+                )[0]
+                latent = temp_x0.squeeze(0)
+
+            # Offload models
+            self.low_noise_model.cpu()
+            self.high_noise_model.cpu()
+            torch.cuda.empty_cache()
+
+            # Decode video
+            videos = self.vae.decode([latent])
+
+        del noise, latent
+        gc.collect()
+        torch.cuda.synchronize()
+
+        return videos[0]
+
+
+def save_video(frames: torch.Tensor, output_path: str, fps: int = 16):
+    """Save video frames to file."""
+    import imageio
+
+    frames = ((frames + 1) / 2 * 255).clamp(0, 255).byte()
+    frames = frames.permute(1, 2, 3, 0).cpu().numpy()
+
+    imageio.mimwrite(output_path, frames, fps=fps, codec="libx264")
+    logger.info(f"Saved video to {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Generate videos with pre-quantized NF4 models"
+    )
+    # Default to current directory (for self-contained HuggingFace repo)
+    script_dir = str(Path(__file__).parent)
+    parser.add_argument("--ckpt_dir", type=str, default=script_dir)
+    parser.add_argument("--image", type=str, required=True, help="Input image path")
+    parser.add_argument("--prompt", type=str, required=True, help="Text prompt")
+    parser.add_argument(
+        "--action_path", type=str, default=None, help="Camera control path"
+    )
+    parser.add_argument("--size", type=str, default="480*832", help="Output resolution")
+    parser.add_argument("--frame_num", type=int, default=81)
+    parser.add_argument("--sampling_steps", type=int, default=40)
+    parser.add_argument("--guide_scale", type=float, default=5.0)
+    parser.add_argument("--seed", type=int, default=-1)
+    parser.add_argument("--output", type=str, default="output.mp4")
+    parser.add_argument("--t5_cpu", action="store_true", default=True)
+    args = parser.parse_args()
+
+    h, w = map(int, args.size.split("*"))
+    max_area = h * w
+
+    img = Image.open(args.image).convert("RGB")
+
+    pipeline = WanI2V_PreQuant(
+        checkpoint_dir=args.ckpt_dir,
+        t5_cpu=args.t5_cpu,
+    )
+
+    logger.info("Generating video...")
+    video = pipeline.generate(
+        input_prompt=args.prompt,
+        img=img,
+        action_path=args.action_path,
+        max_area=max_area,
+        frame_num=args.frame_num,
+        sampling_steps=args.sampling_steps,
+        guide_scale=args.guide_scale,
+        seed=args.seed,
+    )
+
+    save_video(video, args.output)
+
+
+if __name__ == "__main__":
+    main()

high_noise_model_bnb_nf4/config.json

@@ -0,0 +1,15 @@
+{
+  "_class_name": "WanModel",
+  "_diffusers_version": "0.36.0",
+  "_name_or_path": "lingbot-world-base-cam",
+  "model_type": "i2v",
+  "text_len": 512,
+  "in_dim": 36,
+  "dim": 5120,
+  "ffn_dim": 13824,
+  "freq_dim": 256,
+  "out_dim": 16,
+  "num_heads": 40,
+  "num_layers": 40,
+  "eps": 1e-06
+}

high_noise_model_bnb_nf4/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9319e7bccffcc02baf23f1ed4f341423c73ee5f4409cf5700602ecacca87fcb6
+size 9577265260

high_noise_model_bnb_nf4/quantization_meta.json

@@ -0,0 +1,54 @@
+{
+  "source": {
+    "ckpt_dir": "lingbot-world-base-cam",
+    "subfolder": "high_noise_model"
+  },
+  "model_config": {
+    "model_type": "i2v",
+    "patch_size": [
+      1,
+      2,
+      2
+    ],
+    "text_len": 512,
+    "in_dim": 36,
+    "dim": 5120,
+    "ffn_dim": 13824,
+    "freq_dim": 256,
+    "text_dim": 4096,
+    "out_dim": 16,
+    "num_heads": 40,
+    "num_layers": 40,
+    "window_size": [
+      -1,
+      -1
+    ],
+    "qk_norm": true,
+    "cross_attn_norm": true,
+    "eps": 1e-06
+  },
+  "quant": {
+    "format": "bnb_nf4",
+    "double_quant": true,
+    "compute_dtype": "bfloat16",
+    "quant_type": "nf4"
+  },
+  "files": {
+    "weights": "model.safetensors",
+    "config": "config.json"
+  },
+  "sizes": {
+    "original_bytes": 37088665728,
+    "quantized_bytes": 9577265260,
+    "compression_ratio": 3.87
+  },
+  "versions": {
+    "python": "3.10.12",
+    "torch": "2.10.0+cu130",
+    "cuda": "13.0",
+    "bitsandbytes": "0.49.1",
+    "platform": "Linux-6.8.0-94-generic-x86_64-with-glibc2.35",
+    "git_commit": "fd9c95769b72"
+  },
+  "created_at": "2026-02-03T06:16:38.006296Z"
+}

load_prequant.py

@@ -0,0 +1,412 @@
+#!/usr/bin/env python3
+"""
+Load pre-quantized bitsandbytes NF4 models without needing base weights.
+
+This module provides utilities to load pre-quantized WanModel weights
+directly from safetensors/pt files, without downloading or loading
+the original FP16/BF16 base model weights.
+
+Usage:
+    from load_prequant import load_quantized_model
+
+    model = load_quantized_model("lingbot-world-base-cam/high_noise_model_bnb_nf4")
+"""
+
+import json
+import os
+from pathlib import Path
+from typing import Optional, Dict, Any, Tuple
+from collections import defaultdict
+
+import torch
+import torch.nn as nn
+import bitsandbytes as bnb
+from bitsandbytes.functional import QuantState
+
+# Add parent to path for wan imports
+import sys
+sys.path.insert(0, str(Path(__file__).parent))
+
+from wan.modules.model import WanModel
+
+
+def replace_linears_with_bnb_nf4(
+    model: nn.Module,
+    compute_dtype: torch.dtype = torch.bfloat16,
+    compress_statistics: bool = True,
+    quant_type: str = "nf4",
+) -> Tuple[int, Dict[str, Tuple[int, int]]]:
+    """
+    Replace all nn.Linear layers with empty bnb.nn.Linear4bit layers.
+
+    This creates the structure needed to load pre-quantized weights.
+    The layers are created without weights - they will be populated
+    by load_state_dict afterwards.
+
+    Args:
+        model: The model to modify in-place
+        compute_dtype: Compute dtype for the quantized layers
+        compress_statistics: Whether to use double quantization
+        quant_type: Quantization type ('nf4' or 'fp4')
+
+    Returns:
+        Tuple of (num_replaced, dict mapping layer_name to (in_features, out_features))
+    """
+    replaced = 0
+    layer_shapes = {}
+
+    # Collect all linear layers first to avoid modifying during iteration
+    linear_layers = []
+    for name, module in model.named_modules():
+        if isinstance(module, nn.Linear):
+            linear_layers.append((name, module))
+
+    for name, module in linear_layers:
+        # Get parent module
+        parent_name = '.'.join(name.split('.')[:-1])
+        child_name = name.split('.')[-1]
+
+        if parent_name:
+            parent = model.get_submodule(parent_name)
+        else:
+            parent = model
+
+        # Store original shape for reconstruction
+        layer_shapes[name] = (module.in_features, module.out_features)
+
+        # Create empty NF4 linear layer with same shape
+        nf4_linear = bnb.nn.Linear4bit(
+            module.in_features,
+            module.out_features,
+            bias=module.bias is not None,
+            compute_dtype=compute_dtype,
+            compress_statistics=compress_statistics,
+            quant_type=quant_type,
+        )
+
+        # Replace the layer (weights will be loaded from state_dict)
+        setattr(parent, child_name, nf4_linear)
+        replaced += 1
+
+    return replaced, layer_shapes
+
+
+def build_model_from_config(config: Dict[str, Any]) -> WanModel:
+    """
+    Build a WanModel instance from config dictionary.
+
+    Args:
+        config: Dictionary with model configuration
+
+    Returns:
+        Uninitialized WanModel instance
+    """
+    # Extract config values with defaults matching WanModel.__init__
+    model = WanModel(
+        model_type=config.get("model_type", "i2v"),
+        patch_size=tuple(config.get("patch_size", (1, 2, 2))),
+        text_len=config.get("text_len", 512),
+        in_dim=config.get("in_dim", 16),
+        dim=config.get("dim", 2048),
+        ffn_dim=config.get("ffn_dim", 8192),
+        freq_dim=config.get("freq_dim", 256),
+        text_dim=config.get("text_dim", 4096),
+        out_dim=config.get("out_dim", 16),
+        num_heads=config.get("num_heads", 16),
+        num_layers=config.get("num_layers", 32),
+        window_size=tuple(config.get("window_size", (-1, -1))),
+        qk_norm=config.get("qk_norm", True),
+        cross_attn_norm=config.get("cross_attn_norm", True),
+        eps=config.get("eps", 1e-6),
+    )
+
+    return model
+
+
+def reconstruct_params4bit_from_components(
+    weight_components: Dict[str, torch.Tensor],
+    device: str = "cuda",
+) -> bnb.nn.Params4bit:
+    """
+    Reconstruct a Params4bit object from serialized components using QuantState.from_dict.
+
+    This uses bitsandbytes' own deserialization method for correctness.
+
+    Args:
+        weight_components: Dict with keys like 'weight', 'absmax', 'quant_map',
+                          'nested_absmax', 'nested_quant_map', 'quant_state_data'
+        device: Device to load to
+
+    Returns:
+        Reconstructed Params4bit
+    """
+    # Build the dict that QuantState.from_dict expects
+    qs_dict = {
+        "absmax": weight_components["absmax"],
+        "quant_map": weight_components["quant_map"],
+    }
+
+    # Add nested quantization components if present (double quantization)
+    if "nested_absmax" in weight_components:
+        qs_dict["nested_absmax"] = weight_components["nested_absmax"]
+        qs_dict["nested_quant_map"] = weight_components["nested_quant_map"]
+
+    # Add the packed quant_state data (contains shape, dtype, etc.)
+    if "quant_state_data" in weight_components:
+        qs_dict["quant_state.bitsandbytes__nf4"] = weight_components["quant_state_data"]
+
+    # Use bitsandbytes' own deserialization
+    quant_state = QuantState.from_dict(qs_dict, device=torch.device(device))
+
+    # Get quantized weight and move to device
+    quantized_weight = weight_components["weight"].to(device)
+
+    # Create Params4bit with the quantized data
+    param = bnb.nn.Params4bit(
+        data=quantized_weight,
+        requires_grad=False,
+        quant_state=quant_state,
+        bnb_quantized=True,  # Already quantized, don't re-quantize on .to()
+    )
+
+    return param
+
+
+def load_quantized_state(
+    model: nn.Module,
+    weights_path: str,
+    layer_shapes: Dict[str, Tuple[int, int]],
+    device: str = "cpu",
+) -> nn.Module:
+    """
+    Load quantized weights into a model with bnb.Linear4bit layers.
+
+    This function handles the special bitsandbytes serialization format
+    where weights are decomposed into quantized data + metadata tensors.
+
+    Uses QuantState.from_dict for proper deserialization.
+
+    Args:
+        model: Model with Linear4bit layers already in place
+        weights_path: Path to model.safetensors or model.pt
+        layer_shapes: Dict mapping layer names to (in_features, out_features)
+        device: Device to load quantized weights to
+
+    Returns:
+        Model with loaded weights
+    """
+    if weights_path.endswith(".safetensors"):
+        from safetensors.torch import load_file
+        sd = load_file(weights_path)
+    else:
+        sd = torch.load(weights_path, map_location="cpu", weights_only=False)
+
+    # Group keys by their base weight name
+    weight_components = defaultdict(dict)
+    other_keys = {}
+
+    # Quantization-related suffixes
+    quant_suffixes = [".absmax", ".quant_map", ".nested_absmax", ".nested_quant_map", ".quant_state.bitsandbytes__nf4"]
+
+    for key, tensor in sd.items():
+        base_key = None
+        component = None
+
+        if ".weight.absmax" in key:
+            base_key = key.replace(".weight.absmax", "")
+            component = "absmax"
+        elif ".weight.quant_map" in key:
+            base_key = key.replace(".weight.quant_map", "")
+            component = "quant_map"
+        elif ".weight.nested_absmax" in key:
+            base_key = key.replace(".weight.nested_absmax", "")
+            component = "nested_absmax"
+        elif ".weight.nested_quant_map" in key:
+            base_key = key.replace(".weight.nested_quant_map", "")
+            component = "nested_quant_map"
+        elif ".weight.quant_state.bitsandbytes__nf4" in key:
+            base_key = key.replace(".weight.quant_state.bitsandbytes__nf4", "")
+            component = "quant_state_data"
+        elif key.endswith(".weight"):
+            # Check if this is a quantized linear weight or regular weight
+            potential_base = key[:-7]  # Remove ".weight"
+            has_quant_metadata = any(f"{potential_base}.weight{suffix}" in sd for suffix in quant_suffixes)
+
+            if has_quant_metadata:
+                base_key = potential_base
+                component = "weight"
+            else:
+                other_keys[key] = tensor
+                continue
+        else:
+            other_keys[key] = tensor
+            continue
+
+        if base_key and component:
+            weight_components[base_key][component] = tensor
+
+    # Load quantized weights into model
+    loaded_count = 0
+    for name, module in model.named_modules():
+        if isinstance(module, bnb.nn.Linear4bit):
+            if name in weight_components and name in layer_shapes:
+                components = weight_components[name]
+
+                if "weight" in components:
+                    # Use bitsandbytes' own deserialization via QuantState.from_dict
+                    param = reconstruct_params4bit_from_components(components, device=device)
+                    module.weight = param
+                    loaded_count += 1
+
+            # Load bias if present
+            bias_key = f"{name}.bias"
+            if bias_key in other_keys and module.bias is not None:
+                module.bias.data.copy_(other_keys[bias_key].to(device))
+
+    # Load non-quantized weights (embeddings, norms, biases, etc.)
+    non_linear_sd = {}
+    for key, tensor in other_keys.items():
+        non_linear_sd[key] = tensor
+
+    if non_linear_sd:
+        missing, unexpected = model.load_state_dict(non_linear_sd, strict=False)
+        expected_missing = {f"{name}.weight" for name in layer_shapes.keys()}
+        critical_missing = [k for k in missing if k not in expected_missing and not k.endswith("freqs")]
+        if critical_missing:
+            print(f"Warning: Missing non-quantized keys: {critical_missing[:10]}...")
+
+    print(f"  Loaded {loaded_count} quantized linear layers")
+    return model
+
+
+def load_quantized_model(
+    model_dir: str,
+    device: str = "cuda",
+    compute_dtype: torch.dtype = torch.bfloat16,
+) -> WanModel:
+    """
+    Load a pre-quantized WanModel from a directory.
+
+    This function:
+    1. Reads config.json to get model architecture
+    2. Builds an empty WanModel
+    3. Replaces Linear layers with bnb.Linear4bit
+    4. Loads the pre-quantized weights with proper reconstruction
+    5. Moves model to device
+
+    Args:
+        model_dir: Directory containing config.json and model.safetensors/model.pt
+        device: Device to load model to
+        compute_dtype: Compute dtype for quantized layers
+
+    Returns:
+        Loaded and ready WanModel
+    """
+    model_dir = Path(model_dir)
+
+    # Load config
+    config_path = model_dir / "config.json"
+    if not config_path.exists():
+        raise FileNotFoundError(f"Config not found: {config_path}")
+
+    with open(config_path, "r") as f:
+        config = json.load(f)
+
+    # Check for quantization metadata (optional but recommended)
+    meta_path = model_dir / "quantization_meta.json"
+    if meta_path.exists():
+        with open(meta_path, "r") as f:
+            meta = json.load(f)
+        quant_config = meta.get("quant", {})
+        compute_dtype_str = quant_config.get("compute_dtype", "bfloat16")
+        compute_dtype = getattr(torch, compute_dtype_str, torch.bfloat16)
+
+    # Find weights file (prefer safetensors)
+    safetensors_path = model_dir / "model.safetensors"
+    pt_path = model_dir / "model.pt"
+
+    if safetensors_path.exists():
+        weights_path = str(safetensors_path)
+    elif pt_path.exists():
+        weights_path = str(pt_path)
+    else:
+        raise FileNotFoundError(
+            f"No weights found in {model_dir}. "
+            "Expected model.safetensors or model.pt"
+        )
+
+    print(f"Loading pre-quantized model from {model_dir}")
+    print(f"  Config: {config_path}")
+    print(f"  Weights: {weights_path}")
+
+    # Build model from config (creates initialized weights we'll replace)
+    model = build_model_from_config(config)
+
+    # Replace Linear → Linear4bit (empty, ready for state_dict)
+    replaced, layer_shapes = replace_linears_with_bnb_nf4(model, compute_dtype=compute_dtype)
+    print(f"  Replaced {replaced} linear layers with bnb.Linear4bit")
+
+    # Load quantized weights with proper reconstruction
+    # This loads quantized weights directly to the target device
+    model = load_quantized_state(model, weights_path, layer_shapes, device=device)
+
+    # Move non-quantized parts to device and set eval mode
+    model.to(device)
+    model.eval()
+    model.requires_grad_(False)
+
+    print(f"  Model ready on {device}")
+
+    return model
+
+
+def verify_quantized_model(model: nn.Module) -> Dict[str, Any]:
+    """
+    Verify that a model has been properly quantized.
+
+    Args:
+        model: Model to verify
+
+    Returns:
+        Dictionary with verification results
+    """
+    total_params = 0
+    quantized_params = 0
+    linear4bit_count = 0
+    regular_linear_count = 0
+
+    for name, module in model.named_modules():
+        if isinstance(module, bnb.nn.Linear4bit):
+            linear4bit_count += 1
+            if hasattr(module.weight, 'quant_state') and module.weight.quant_state is not None:
+                quantized_params += module.weight.numel()
+        elif isinstance(module, nn.Linear):
+            regular_linear_count += 1
+
+    for param in model.parameters():
+        total_params += param.numel()
+
+    return {
+        "total_params": total_params,
+        "quantized_params": quantized_params,
+        "linear4bit_count": linear4bit_count,
+        "regular_linear_count": regular_linear_count,
+        "is_quantized": linear4bit_count > 0 and regular_linear_count == 0,
+    }
+
+
+if __name__ == "__main__":
+    import argparse
+
+    parser = argparse.ArgumentParser(description="Test loading pre-quantized model")
+    parser.add_argument("model_dir", type=str, help="Path to quantized model directory")
+    parser.add_argument("--device", type=str, default="cuda", help="Device to load to")
+    args = parser.parse_args()
+
+    model = load_quantized_model(args.model_dir, device=args.device)
+
+    info = verify_quantized_model(model)
+    print(f"\nVerification:")
+    print(f"  Linear4bit layers: {info['linear4bit_count']}")
+    print(f"  Regular Linear layers: {info['regular_linear_count']}")
+    print(f"  Is properly quantized: {info['is_quantized']}")

low_noise_model_bnb_nf4/config.json

@@ -0,0 +1,15 @@
+{
+  "_class_name": "WanModel",
+  "_diffusers_version": "0.36.0",
+  "_name_or_path": "lingbot-world-base-cam",
+  "model_type": "i2v",
+  "text_len": 512,
+  "in_dim": 36,
+  "dim": 5120,
+  "ffn_dim": 13824,
+  "freq_dim": 256,
+  "out_dim": 16,
+  "num_heads": 40,
+  "num_layers": 40,
+  "eps": 1e-06
+}

low_noise_model_bnb_nf4/model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd898a2770300c34c7092b64a6a6ed94ca881223473c65f37fbbb04d88e3bb56
+size 9577265273

low_noise_model_bnb_nf4/quantization_meta.json

@@ -0,0 +1,54 @@
+{
+  "source": {
+    "ckpt_dir": "lingbot-world-base-cam",
+    "subfolder": "low_noise_model"
+  },
+  "model_config": {
+    "model_type": "i2v",
+    "patch_size": [
+      1,
+      2,
+      2
+    ],
+    "text_len": 512,
+    "in_dim": 36,
+    "dim": 5120,
+    "ffn_dim": 13824,
+    "freq_dim": 256,
+    "text_dim": 4096,
+    "out_dim": 16,
+    "num_heads": 40,
+    "num_layers": 40,
+    "window_size": [
+      -1,
+      -1
+    ],
+    "qk_norm": true,
+    "cross_attn_norm": true,
+    "eps": 1e-06
+  },
+  "quant": {
+    "format": "bnb_nf4",
+    "double_quant": true,
+    "compute_dtype": "bfloat16",
+    "quant_type": "nf4"
+  },
+  "files": {
+    "weights": "model.safetensors",
+    "config": "config.json"
+  },
+  "sizes": {
+    "original_bytes": 37088665728,
+    "quantized_bytes": 9577265273,
+    "compression_ratio": 3.87
+  },
+  "versions": {
+    "python": "3.10.12",
+    "torch": "2.10.0+cu130",
+    "cuda": "13.0",
+    "bitsandbytes": "0.49.1",
+    "platform": "Linux-6.8.0-94-generic-x86_64-with-glibc2.35",
+    "git_commit": "fd9c95769b72"
+  },
+  "created_at": "2026-02-03T06:19:46.213374Z"
+}

models_t5_umt5-xxl-enc-bf16.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7cace0da2b446bbbbc57d031ab6cf163a3d59b366da94e5afe36745b746fd81d
+size 11361920418

requirements.txt

@@ -0,0 +1,15 @@
+torch>=2.0.0
+torchvision
+torchaudio
+bitsandbytes>=0.43.0
+safetensors
+accelerate
+transformers
+diffusers
+einops
+imageio
+imageio-ffmpeg
+pillow
+numpy
+tqdm
+easydict

tokenizer/special_tokens_map.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:456b58fd240a06c743a7c2cf8008bec501240d68ebd1fc4018ea569505fea270
+size 7079

tokenizer/tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:20a46ac256746594ed7e1e3ef733b83fbc5a6f0922aa7480eda961743de080ef
+size 16837459

tokenizer/tokenizer_config.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:89bef11d01c80a229079666dd470d536aba87b8eaa00b51cc481ba625be59269
+size 61786

wan/__init__.py

@@ -0,0 +1,2 @@
+from . import configs, distributed, modules
+from .image2video import WanI2V

wan/configs/__init__.py

@@ -0,0 +1,36 @@
+import copy
+import os
+
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+from .wan_i2v_A14B import i2v_A14B
+
+WAN_CONFIGS = {
+    'i2v-A14B': i2v_A14B,
+}
+
+SIZE_CONFIGS = {
+    '720*1280': (720, 1280),
+    '1280*720': (1280, 720),
+    '480*832': (480, 832),
+    '832*480': (832, 480),
+    '704*1280': (704, 1280),
+    '1280*704': (1280, 704),
+    '1024*704': (1024, 704),
+    '704*1024': (704, 1024),
+}
+
+MAX_AREA_CONFIGS = {
+    '720*1280': 720 * 1280,
+    '1280*720': 1280 * 720,
+    '480*832': 480 * 832,
+    '832*480': 832 * 480,
+    '704*1280': 704 * 1280,
+    '1280*704': 1280 * 704,
+    '1024*704': 1024 * 704,
+    '704*1024': 704 * 1024,
+}
+
+SUPPORTED_SIZES = {
+    'i2v-A14B': ('720*1280', '1280*720', '480*832', '832*480')
+}

wan/configs/shared_config.py

@@ -0,0 +1,19 @@
+import torch
+from easydict import EasyDict
+
+#------------------------ Wan shared config ------------------------#
+wan_shared_cfg = EasyDict()
+
+# t5
+wan_shared_cfg.t5_model = 'umt5_xxl'
+wan_shared_cfg.t5_dtype = torch.bfloat16
+wan_shared_cfg.text_len = 512
+
+# transformer
+wan_shared_cfg.param_dtype = torch.bfloat16
+
+# inference
+wan_shared_cfg.num_train_timesteps = 1000
+wan_shared_cfg.sample_fps = 16
+wan_shared_cfg.sample_neg_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
+wan_shared_cfg.frame_num = 81

wan/configs/wan_i2v_A14B.py

@@ -0,0 +1,37 @@
+import torch
+from easydict import EasyDict
+
+from .shared_config import wan_shared_cfg
+
+#------------------------ Wan I2V A14B ------------------------#
+
+i2v_A14B = EasyDict(__name__='Config: Wan I2V A14B')
+i2v_A14B.update(wan_shared_cfg)
+
+i2v_A14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+i2v_A14B.t5_tokenizer = 'google/umt5-xxl'
+
+# vae
+i2v_A14B.vae_checkpoint = 'Wan2.1_VAE.pth'
+i2v_A14B.vae_stride = (4, 8, 8)
+
+# transformer
+i2v_A14B.patch_size = (1, 2, 2)
+i2v_A14B.dim = 5120
+i2v_A14B.ffn_dim = 13824
+i2v_A14B.freq_dim = 256
+i2v_A14B.num_heads = 40
+i2v_A14B.num_layers = 40
+i2v_A14B.window_size = (-1, -1)
+i2v_A14B.qk_norm = True
+i2v_A14B.cross_attn_norm = True
+i2v_A14B.eps = 1e-6
+i2v_A14B.low_noise_checkpoint = 'low_noise_model'
+i2v_A14B.high_noise_checkpoint = 'high_noise_model'
+
+# inference
+i2v_A14B.sample_shift = 10.0
+i2v_A14B.sample_steps = 70
+i2v_A14B.boundary = 0.947
+i2v_A14B.sample_guide_scale = (5.0, 5.0)  # low noise, high noise
+i2v_A14B.sample_neg_prompt = '画面突变，色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走，镜头晃动，画面闪烁，模糊，噪点，水印，签名，文字，变形，扭曲，液化，不合逻辑的结构，卡顿，PPT幻灯片感，过暗，欠曝，低对比度，霓虹灯光感，过度锐化，3D渲染感，人物，行人，游客，身体，皮肤，肢体，面部特征，汽车，电线'

wan/distributed/fsdp.py

@@ -0,0 +1,44 @@
+import gc
+from functools import partial
+
+import torch
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
+from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy
+from torch.distributed.utils import _free_storage
+
+
+def shard_model(
+    model,
+    device_id,
+    param_dtype=torch.bfloat16,
+    reduce_dtype=torch.float32,
+    buffer_dtype=torch.float32,
+    process_group=None,
+    sharding_strategy=ShardingStrategy.FULL_SHARD,
+    sync_module_states=True,
+    use_lora=False
+):
+    model = FSDP(
+        module=model,
+        process_group=process_group,
+        sharding_strategy=sharding_strategy,
+        auto_wrap_policy=partial(
+            lambda_auto_wrap_policy, lambda_fn=lambda m: m in model.blocks),
+        mixed_precision=MixedPrecision(
+            param_dtype=param_dtype,
+            reduce_dtype=reduce_dtype,
+            buffer_dtype=buffer_dtype),
+        device_id=device_id,
+        sync_module_states=sync_module_states,
+        use_orig_params=True if use_lora else False)
+    return model
+
+
+def free_model(model):
+    for m in model.modules():
+        if isinstance(m, FSDP):
+            _free_storage(m._handle.flat_param.data)
+    del model
+    gc.collect()
+    torch.cuda.empty_cache()

wan/distributed/sequence_parallel.py

@@ -0,0 +1,213 @@
+import torch
+import torch.cuda.amp as amp
+import torch.nn.functional as torch_F
+from einops import rearrange
+
+from ..modules.model import sinusoidal_embedding_1d
+from .ulysses import distributed_attention
+from .util import gather_forward, get_rank, get_world_size
+
+
+def pad_freqs(original_tensor, target_len):
+    seq_len, s1, s2 = original_tensor.shape
+    pad_size = target_len - seq_len
+    padding_tensor = torch.ones(
+        pad_size,
+        s1,
+        s2,
+        dtype=original_tensor.dtype,
+        device=original_tensor.device)
+    padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0)
+    return padded_tensor
+
+
+@torch.amp.autocast('cuda', enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+    """
+    x:          [B, L, N, C].
+    grid_sizes: [B, 3].
+    freqs:      [M, C // 2].
+    """
+    s, n, c = x.size(1), x.size(2), x.size(3) // 2
+    # split freqs
+    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+    # loop over samples
+    output = []
+    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+        seq_len = f * h * w
+
+        # precompute multipliers
+        x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(
+            s, n, -1, 2))
+        freqs_i = torch.cat([
+            freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+        ],
+                            dim=-1).reshape(seq_len, 1, -1)
+
+        # apply rotary embedding
+        sp_size = get_world_size()
+        sp_rank = get_rank()
+        freqs_i = pad_freqs(freqs_i, s * sp_size)
+        s_per_rank = s
+        freqs_i_rank = freqs_i[(sp_rank * s_per_rank):((sp_rank + 1) *
+                                                       s_per_rank), :, :]
+        x_i = torch.view_as_real(x_i * freqs_i_rank).flatten(2)
+        x_i = torch.cat([x_i, x[i, s:]])
+
+        # append to collection
+        output.append(x_i)
+    return torch.stack(output).float()
+
+
+def sp_dit_forward(
+    self,
+    x,
+    t,
+    context,
+    seq_len,
+    y=None,
+    dit_cond_dict=None,
+):
+    """
+    x:              A list of videos each with shape [C, T, H, W].
+    t:              [B].
+    context:        A list of text embeddings each with shape [L, C].
+    """
+    if self.model_type == 'i2v':
+        assert y is not None
+    # params
+    device = self.patch_embedding.weight.device
+    if self.freqs.device != device:
+        self.freqs = self.freqs.to(device)
+
+    if y is not None:
+        x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+    # embeddings
+    x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+    grid_sizes = torch.stack(
+        [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+    x = [u.flatten(2).transpose(1, 2) for u in x]
+    seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+    assert seq_lens.max() <= seq_len
+    x = torch.cat([
+        torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)
+        for u in x
+    ])
+
+    # time embeddings
+    if t.dim() == 1:
+        t = t.expand(t.size(0), seq_len)
+    with torch.amp.autocast('cuda', dtype=torch.float32):
+        bt = t.size(0)
+        t = t.flatten()
+        e = self.time_embedding(
+            sinusoidal_embedding_1d(self.freq_dim,
+                                    t).unflatten(0, (bt, seq_len)).float())
+        e0 = self.time_projection(e).unflatten(2, (6, self.dim))
+        assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+    # context
+    context_lens = None
+    context = self.text_embedding(
+        torch.stack([
+            torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+            for u in context
+        ]))
+
+    # cam
+    if dit_cond_dict is not None and "c2ws_plucker_emb" in dit_cond_dict:
+        c2ws_plucker_emb = dit_cond_dict["c2ws_plucker_emb"]
+        c2ws_plucker_emb = [
+            rearrange(
+                i,
+                '1 c (f c1) (h c2) (w c3) -> 1 (f h w) (c c1 c2 c3)',
+                c1=self.patch_size[0],
+                c2=self.patch_size[1],
+                c3=self.patch_size[2],
+            ) for i in c2ws_plucker_emb
+        ]
+        c2ws_plucker_emb = torch.cat(c2ws_plucker_emb,
+                                     dim=1)  # [1, (L1+...+Ln), C]
+        c2ws_plucker_emb = self.patch_embedding_wancamctrl(c2ws_plucker_emb)
+        c2ws_hidden_states = self.c2ws_hidden_states_layer2(
+            torch_F.silu(self.c2ws_hidden_states_layer1(c2ws_plucker_emb)))
+        c2ws_plucker_emb = c2ws_plucker_emb + c2ws_hidden_states
+
+        cam_len = c2ws_plucker_emb.size(1)
+        if cam_len < seq_len:
+            pad_len = seq_len - cam_len
+            pad = c2ws_plucker_emb.new_zeros(
+                c2ws_plucker_emb.size(0), pad_len, c2ws_plucker_emb.size(2))
+            c2ws_plucker_emb = torch.cat([c2ws_plucker_emb, pad], dim=1)
+        elif cam_len > seq_len:
+            c2ws_plucker_emb = c2ws_plucker_emb[:, :seq_len, :]
+
+        if get_world_size() > 1:
+            c2ws_plucker_emb = torch.chunk(
+                c2ws_plucker_emb, get_world_size(), dim=1)[get_rank()]
+        dit_cond_dict = dict(dit_cond_dict)
+        dit_cond_dict["c2ws_plucker_emb"] = c2ws_plucker_emb
+
+    # Context Parallel
+    x = torch.chunk(x, get_world_size(), dim=1)[get_rank()]
+    e = torch.chunk(e, get_world_size(), dim=1)[get_rank()]
+    e0 = torch.chunk(e0, get_world_size(), dim=1)[get_rank()]
+
+    # arguments
+    kwargs = dict(
+        e=e0,
+        seq_lens=seq_lens,
+        grid_sizes=grid_sizes,
+        freqs=self.freqs,
+        context=context,
+        context_lens=context_lens,
+        dit_cond_dict=dit_cond_dict)
+
+    for block in self.blocks:
+        x = block(x, **kwargs)
+
+    # head
+    x = self.head(x, e)
+
+    # Context Parallel
+    x = gather_forward(x, dim=1)
+
+    # unpatchify
+    x = self.unpatchify(x, grid_sizes)
+    return [u.float() for u in x]
+
+
+def sp_attn_forward(self, x, seq_lens, grid_sizes, freqs, dtype=torch.bfloat16):
+    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+    half_dtypes = (torch.float16, torch.bfloat16)
+
+    def half(x):
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    # query, key, value function
+    def qkv_fn(x):
+        q = self.norm_q(self.q(x)).view(b, s, n, d)
+        k = self.norm_k(self.k(x)).view(b, s, n, d)
+        v = self.v(x).view(b, s, n, d)
+        return q, k, v
+
+    q, k, v = qkv_fn(x)
+    q = rope_apply(q, grid_sizes, freqs)
+    k = rope_apply(k, grid_sizes, freqs)
+
+    x = distributed_attention(
+        half(q),
+        half(k),
+        half(v),
+        seq_lens,
+        window_size=self.window_size,
+    )
+
+    # output
+    x = x.flatten(2)
+    x = self.o(x)
+    return x

wan/distributed/ulysses.py

@@ -0,0 +1,46 @@
+import torch
+import torch.distributed as dist
+
+from ..modules.attention import flash_attention
+from .util import all_to_all
+
+
+def distributed_attention(
+        q,
+        k,
+        v,
+        seq_lens,
+        window_size=(-1, -1),
+):
+    """
+    Performs distributed attention based on DeepSpeed Ulysses attention mechanism.
+    please refer to https://arxiv.org/pdf/2309.14509
+
+    Args:
+        q:           [B, Lq // p, Nq, C1].
+        k:           [B, Lk // p, Nk, C1].
+        v:           [B, Lk // p, Nk, C2]. Nq must be divisible by Nk.
+        seq_lens:    [B], length of each sequence in batch
+        window_size: (left right). If not (-1, -1), apply sliding window local attention.
+    """
+    if not dist.is_initialized():
+        raise ValueError("distributed group should be initialized.")
+    b = q.shape[0]
+
+    # gather q/k/v sequence
+    q = all_to_all(q, scatter_dim=2, gather_dim=1)
+    k = all_to_all(k, scatter_dim=2, gather_dim=1)
+    v = all_to_all(v, scatter_dim=2, gather_dim=1)
+
+    # apply attention
+    x = flash_attention(
+        q,
+        k,
+        v,
+        k_lens=seq_lens,
+        window_size=window_size,
+    )
+
+    # scatter q/k/v sequence
+    x = all_to_all(x, scatter_dim=1, gather_dim=2)
+    return x

wan/distributed/util.py

@@ -0,0 +1,50 @@
+import torch
+import torch.distributed as dist
+
+
+def init_distributed_group():
+    """r initialize sequence parallel group.
+    """
+    if not dist.is_initialized():
+        dist.init_process_group(backend='nccl')
+
+
+def get_rank():
+    return dist.get_rank()
+
+
+def get_world_size():
+    return dist.get_world_size()
+
+
+def all_to_all(x, scatter_dim, gather_dim, group=None, **kwargs):
+    """
+    `scatter` along one dimension and `gather` along another.
+    """
+    world_size = get_world_size()
+    if world_size > 1:
+        inputs = [u.contiguous() for u in x.chunk(world_size, dim=scatter_dim)]
+        outputs = [torch.empty_like(u) for u in inputs]
+        dist.all_to_all(outputs, inputs, group=group, **kwargs)
+        x = torch.cat(outputs, dim=gather_dim).contiguous()
+    return x
+
+
+def all_gather(tensor):
+    world_size = dist.get_world_size()
+    if world_size == 1:
+        return [tensor]
+    tensor_list = [torch.empty_like(tensor) for _ in range(world_size)]
+    torch.distributed.all_gather(tensor_list, tensor)
+    return tensor_list
+
+
+def gather_forward(input, dim):
+    # skip if world_size == 1
+    world_size = dist.get_world_size()
+    if world_size == 1:
+        return input
+
+    # gather sequence
+    output = all_gather(input)
+    return torch.cat(output, dim=dim).contiguous()

wan/image2video.py

@@ -0,0 +1,489 @@
+import gc
+import logging
+import math
+import os
+import random
+import sys
+import types
+from contextlib import contextmanager
+from functools import partial
+
+import numpy as np
+import torch
+import torch.cuda.amp as amp
+import torch.distributed as dist
+import torchvision.transforms.functional as TF
+from tqdm import tqdm
+
+from .distributed.fsdp import shard_model
+from .distributed.sequence_parallel import sp_attn_forward, sp_dit_forward
+from .distributed.util import get_world_size
+from .modules.model import WanModel
+from .modules.t5 import T5EncoderModel
+from .modules.vae2_1 import Wan2_1_VAE
+from .utils.fm_solvers import (
+    FlowDPMSolverMultistepScheduler,
+    get_sampling_sigmas,
+    retrieve_timesteps,
+)
+from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from .utils.cam_utils import (
+    compute_relative_poses,
+    interpolate_camera_poses,
+    get_plucker_embeddings,
+    get_Ks_transformed,
+)
+from einops import rearrange
+
+
+class WanI2V:
+
+    def __init__(
+        self,
+        config,
+        checkpoint_dir,
+        device_id=0,
+        rank=0,
+        t5_fsdp=False,
+        dit_fsdp=False,
+        use_sp=False,
+        t5_cpu=False,
+        init_on_cpu=True,
+        convert_model_dtype=False,
+    ):
+        r"""
+        Initializes the image-to-video generation model components.
+
+        Args:
+            config (EasyDict):
+                Object containing model parameters initialized from config.py
+            checkpoint_dir (`str`):
+                Path to directory containing model checkpoints
+            device_id (`int`,  *optional*, defaults to 0):
+                Id of target GPU device
+            rank (`int`,  *optional*, defaults to 0):
+                Process rank for distributed training
+            t5_fsdp (`bool`, *optional*, defaults to False):
+                Enable FSDP sharding for T5 model
+            dit_fsdp (`bool`, *optional*, defaults to False):
+                Enable FSDP sharding for DiT model
+            use_sp (`bool`, *optional*, defaults to False):
+                Enable distribution strategy of sequence parallel.
+            t5_cpu (`bool`, *optional*, defaults to False):
+                Whether to place T5 model on CPU. Only works without t5_fsdp.
+            init_on_cpu (`bool`, *optional*, defaults to True):
+                Enable initializing Transformer Model on CPU. Only works without FSDP or USP.
+            convert_model_dtype (`bool`, *optional*, defaults to False):
+                Convert DiT model parameters dtype to 'config.param_dtype'.
+                Only works without FSDP.
+        """
+        self.device = torch.device(f"cuda:{device_id}")
+        self.config = config
+        self.rank = rank
+        self.t5_cpu = t5_cpu
+        self.init_on_cpu = init_on_cpu
+
+        self.num_train_timesteps = config.num_train_timesteps
+        self.boundary = config.boundary
+        self.param_dtype = config.param_dtype
+
+        if t5_fsdp or dit_fsdp or use_sp:
+            self.init_on_cpu = False
+
+        shard_fn = partial(shard_model, device_id=device_id)
+        self.text_encoder = T5EncoderModel(
+            text_len=config.text_len,
+            dtype=config.t5_dtype,
+            device=torch.device('cpu'),
+            checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),
+            tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),
+            shard_fn=shard_fn if t5_fsdp else None,
+        )
+
+        self.vae_stride = config.vae_stride
+        self.patch_size = config.patch_size
+        self.vae = Wan2_1_VAE(
+            vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
+            device=self.device)
+
+        logging.info(f"Creating WanModel from {checkpoint_dir}")
+        # Load models with low_cpu_mem_usage to allow CPU loading for single GPU setups
+        self.low_noise_model = WanModel.from_pretrained(
+            checkpoint_dir, subfolder=config.low_noise_checkpoint, torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True)
+        self.low_noise_model = self._configure_model(
+            model=self.low_noise_model,
+            use_sp=use_sp,
+            dit_fsdp=dit_fsdp,
+            shard_fn=shard_fn,
+            convert_model_dtype=convert_model_dtype)
+
+        self.high_noise_model = WanModel.from_pretrained(
+            checkpoint_dir, subfolder=config.high_noise_checkpoint, torch_dtype=torch.bfloat16,
+            low_cpu_mem_usage=True)
+        self.high_noise_model = self._configure_model(
+            model=self.high_noise_model,
+            use_sp=use_sp,
+            dit_fsdp=dit_fsdp,
+            shard_fn=shard_fn,
+            convert_model_dtype=convert_model_dtype)
+        if use_sp:
+            self.sp_size = get_world_size()
+        else:
+            self.sp_size = 1
+
+        self.sample_neg_prompt = config.sample_neg_prompt
+
+    def _configure_model(self, model, use_sp, dit_fsdp, shard_fn,
+                         convert_model_dtype):
+        """
+        Configures a model object. This includes setting evaluation modes,
+        applying distributed parallel strategy, and handling device placement.
+
+        Args:
+            model (torch.nn.Module):
+                The model instance to configure.
+            use_sp (`bool`):
+                Enable distribution strategy of sequence parallel.
+            dit_fsdp (`bool`):
+                Enable FSDP sharding for DiT model.
+            shard_fn (callable):
+                The function to apply FSDP sharding.
+            convert_model_dtype (`bool`):
+                Convert DiT model parameters dtype to 'config.param_dtype'.
+                Only works without FSDP.
+
+        Returns:
+            torch.nn.Module:
+                The configured model.
+        """
+        model.eval().requires_grad_(False)
+
+        if use_sp:
+            for block in model.blocks:
+                block.self_attn.forward = types.MethodType(
+                    sp_attn_forward, block.self_attn)
+            model.forward = types.MethodType(sp_dit_forward, model)
+
+        if dist.is_initialized():
+            dist.barrier()
+
+        if dit_fsdp:
+            model = shard_fn(model)
+        else:
+            if convert_model_dtype:
+                model.to(self.param_dtype)
+            if not self.init_on_cpu:
+                model.to(self.device)
+
+        return model
+
+    def _prepare_model_for_timestep(self, t, boundary, offload_model):
+        r"""
+        Prepares and returns the required model for the current timestep.
+
+        Args:
+            t (torch.Tensor):
+                current timestep.
+            boundary (`int`):
+                The timestep threshold. If `t` is at or above this value,
+                the `high_noise_model` is considered as the required model.
+            offload_model (`bool`):
+                A flag intended to control the offloading behavior.
+
+        Returns:
+            torch.nn.Module:
+                The active model on the target device for the current timestep.
+        """
+        if t.item() >= boundary:
+            required_model_name = 'high_noise_model'
+            offload_model_name = 'low_noise_model'
+        else:
+            required_model_name = 'low_noise_model'
+            offload_model_name = 'high_noise_model'
+        if offload_model or self.init_on_cpu:
+            if next(getattr(
+                    self,
+                    offload_model_name).parameters()).device.type == 'cuda':
+                getattr(self, offload_model_name).to('cpu')
+            if next(getattr(
+                    self,
+                    required_model_name).parameters()).device.type == 'cpu':
+                getattr(self, required_model_name).to(self.device)
+        return getattr(self, required_model_name)
+
+    def generate(self,
+                 input_prompt,
+                 img,
+                 action_path=None,
+                 max_area=720 * 1280,
+                 frame_num=81,
+                 shift=5.0,
+                 sample_solver='unipc',
+                 sampling_steps=40,
+                 guide_scale=5.0,
+                 n_prompt="",
+                 seed=-1,
+                 offload_model=True):
+        r"""
+        Generates video frames from input image and text prompt using diffusion process.
+
+        Args:
+            input_prompt (`str`):
+                Text prompt for content generation.
+            img (PIL.Image.Image):
+                Input image tensor. Shape: [3, H, W]
+            max_area (`int`, *optional*, defaults to 720*1280):
+                Maximum pixel area for latent space calculation. Controls video resolution scaling
+            frame_num (`int`, *optional*, defaults to 81):
+                How many frames to sample from a video. The number should be 4n+1
+            shift (`float`, *optional*, defaults to 5.0):
+                Noise schedule shift parameter. Affects temporal dynamics
+                [NOTE]: If you want to generate a 480p video, it is recommended to set the shift value to 3.0.
+            sample_solver (`str`, *optional*, defaults to 'unipc'):
+                Solver used to sample the video.
+            sampling_steps (`int`, *optional*, defaults to 40):
+                Number of diffusion sampling steps. Higher values improve quality but slow generation
+            guide_scale (`float` or tuple[`float`], *optional*, defaults 5.0):
+                Classifier-free guidance scale. Controls prompt adherence vs. creativity.
+                If tuple, the first guide_scale will be used for low noise model and
+                the second guide_scale will be used for high noise model.
+            n_prompt (`str`, *optional*, defaults to ""):
+                Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt`
+            seed (`int`, *optional*, defaults to -1):
+                Random seed for noise generation. If -1, use random seed
+            offload_model (`bool`, *optional*, defaults to True):
+                If True, offloads models to CPU during generation to save VRAM
+
+        Returns:
+            torch.Tensor:
+                Generated video frames tensor. Dimensions: (C, N H, W) where:
+                - C: Color channels (3 for RGB)
+                - N: Number of frames (81)
+                - H: Frame height (from max_area)
+                - W: Frame width from max_area)
+        """
+        if action_path is not None:
+            c2ws = np.load(os.path.join(action_path, "poses.npy")) # opencv coordinate
+            len_c2ws = ((len(c2ws) - 1) // 4) * 4 + 1
+            frame_num = min(frame_num, len_c2ws)
+            c2ws = c2ws[:frame_num]
+
+        # preprocess
+        guide_scale = (guide_scale, guide_scale) if isinstance(
+            guide_scale, float) else guide_scale
+        img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device)
+
+        F = frame_num
+        h, w = img.shape[1:]
+        aspect_ratio = h / w
+        lat_h = round(
+            np.sqrt(max_area * aspect_ratio) // self.vae_stride[1] //
+            self.patch_size[1] * self.patch_size[1])
+        lat_w = round(
+            np.sqrt(max_area / aspect_ratio) // self.vae_stride[2] //
+            self.patch_size[2] * self.patch_size[2])
+        h = lat_h * self.vae_stride[1]
+        w = lat_w * self.vae_stride[2]
+        lat_f = (F - 1) // self.vae_stride[0] + 1
+        max_seq_len = lat_f * lat_h * lat_w // (
+            self.patch_size[1] * self.patch_size[2])
+        max_seq_len = int(math.ceil(max_seq_len / self.sp_size)) * self.sp_size
+
+        seed = seed if seed >= 0 else random.randint(0, sys.maxsize)
+        seed_g = torch.Generator(device=self.device)
+        seed_g.manual_seed(seed)
+        noise = torch.randn(
+            16,
+            (F - 1) // self.vae_stride[0] + 1,
+            lat_h,
+            lat_w,
+            dtype=torch.float32,
+            generator=seed_g,
+            device=self.device)
+
+        msk = torch.ones(1, F, lat_h, lat_w, device=self.device)
+        msk[:, 1:] = 0
+        msk = torch.concat([
+            torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]
+        ],
+                           dim=1)
+        msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)
+        msk = msk.transpose(1, 2)[0]
+
+        if n_prompt == "":
+            n_prompt = self.sample_neg_prompt
+
+        # preprocess
+        if not self.t5_cpu:
+            self.text_encoder.model.to(self.device)
+            context = self.text_encoder([input_prompt], self.device)
+            context_null = self.text_encoder([n_prompt], self.device)
+            if offload_model:
+                self.text_encoder.model.cpu()
+        else:
+            context = self.text_encoder([input_prompt], torch.device('cpu'))
+            context_null = self.text_encoder([n_prompt], torch.device('cpu'))
+            context = [t.to(self.device) for t in context]
+            context_null = [t.to(self.device) for t in context_null]
+
+        # cam preparation (only if action_path is provided)
+        dit_cond_dict = None
+        if action_path is not None:
+            Ks = torch.from_numpy(np.load(os.path.join(action_path, "intrinsics.npy"))).float()
+
+            # The provided intrinsics are for original image size (480p). We need to transform them according to the new image size (h, w).
+            Ks = get_Ks_transformed(Ks,
+                                    height_org=480,
+                                    width_org=832,
+                                    height_resize=h,
+                                    width_resize=w,
+                                    height_final=h,
+                                    width_final=w)
+            Ks = Ks[0]
+            
+            len_c2ws = len(c2ws)
+            c2ws_infer = interpolate_camera_poses(
+                src_indices=np.linspace(0, len_c2ws - 1, len_c2ws),
+                src_rot_mat=c2ws[:, :3, :3],
+                src_trans_vec=c2ws[:, :3, 3],
+                tgt_indices=np.linspace(0, len_c2ws - 1, int((len_c2ws - 1) // 4) + 1),
+            )
+            c2ws_infer = compute_relative_poses(c2ws_infer, framewise=True)
+            Ks = Ks.repeat(len(c2ws_infer), 1)
+
+            c2ws_infer = c2ws_infer.to(self.device)
+            Ks = Ks.to(self.device)
+            c2ws_plucker_emb = get_plucker_embeddings(c2ws_infer, Ks, h, w)
+            c2ws_plucker_emb = rearrange(
+                c2ws_plucker_emb,
+                'f (h c1) (w c2) c -> (f h w) (c c1 c2)',
+                c1=int(h // lat_h),
+                c2=int(w // lat_w),
+            )
+            c2ws_plucker_emb = c2ws_plucker_emb[None, ...] # [b, f*h*w, c]
+            c2ws_plucker_emb = rearrange(c2ws_plucker_emb, 'b (f h w) c -> b c f h w', f=lat_f, h=lat_h, w=lat_w).to(self.param_dtype)
+            dit_cond_dict = {
+                "c2ws_plucker_emb": c2ws_plucker_emb.chunk(1, dim=0),
+            }
+
+        y = self.vae.encode([
+            torch.concat([
+                torch.nn.functional.interpolate(
+                    img[None].cpu(), size=(h, w), mode='bicubic').transpose(
+                        0, 1),
+                torch.zeros(3, F - 1, h, w)
+            ],
+                         dim=1).to(self.device)
+        ])[0]
+        y = torch.concat([msk, y])
+
+        @contextmanager
+        def noop_no_sync():
+            yield
+
+        no_sync_low_noise = getattr(self.low_noise_model, 'no_sync',
+                                    noop_no_sync)
+        no_sync_high_noise = getattr(self.high_noise_model, 'no_sync',
+                                     noop_no_sync)
+
+        # evaluation mode
+        with (
+                torch.amp.autocast('cuda', dtype=self.param_dtype),
+                torch.no_grad(),
+                no_sync_low_noise(),
+                no_sync_high_noise(),
+        ):
+            boundary = self.boundary * self.num_train_timesteps
+
+            if sample_solver == 'unipc':
+                sample_scheduler = FlowUniPCMultistepScheduler(
+                    num_train_timesteps=self.num_train_timesteps,
+                    shift=1,
+                    use_dynamic_shifting=False)
+                sample_scheduler.set_timesteps(
+                    sampling_steps, device=self.device, shift=shift)
+                timesteps = sample_scheduler.timesteps
+            elif sample_solver == 'dpm++':
+                sample_scheduler = FlowDPMSolverMultistepScheduler(
+                    num_train_timesteps=self.num_train_timesteps,
+                    shift=1,
+                    use_dynamic_shifting=False)
+                sampling_sigmas = get_sampling_sigmas(sampling_steps, shift)
+                timesteps, _ = retrieve_timesteps(
+                    sample_scheduler,
+                    device=self.device,
+                    sigmas=sampling_sigmas)
+            else:
+                raise NotImplementedError("Unsupported solver.")
+
+            # sample videos
+            latent = noise
+
+            arg_c = {
+                'context': [context[0]],
+                'seq_len': max_seq_len,
+                'y': [y],
+                'dit_cond_dict': dit_cond_dict,
+            }
+
+            arg_null = {
+                'context': context_null,
+                'seq_len': max_seq_len,
+                'y': [y],
+                'dit_cond_dict': dit_cond_dict,
+            }
+
+            if offload_model:
+                torch.cuda.empty_cache()
+
+            for _, t in enumerate(tqdm(timesteps)):
+                latent_model_input = [latent.to(self.device)]
+                timestep = [t]
+
+                timestep = torch.stack(timestep).to(self.device)
+
+                model = self._prepare_model_for_timestep(
+                    t, boundary, offload_model)
+                sample_guide_scale = guide_scale[1] if t.item(
+                ) >= boundary else guide_scale[0]
+
+                noise_pred_cond = model(
+                    latent_model_input, t=timestep, **arg_c)[0]
+                if offload_model:
+                    torch.cuda.empty_cache()
+                noise_pred_uncond = model(
+                    latent_model_input, t=timestep, **arg_null)[0]
+                if offload_model:
+                    torch.cuda.empty_cache()
+                noise_pred = noise_pred_uncond + sample_guide_scale * (
+                    noise_pred_cond - noise_pred_uncond)
+
+                temp_x0 = sample_scheduler.step(
+                    noise_pred.unsqueeze(0),
+                    t,
+                    latent.unsqueeze(0),
+                    return_dict=False,
+                    generator=seed_g)[0]
+                latent = temp_x0.squeeze(0)
+
+                x0 = [latent]
+                del latent_model_input, timestep
+
+            if offload_model:
+                self.low_noise_model.cpu()
+                self.high_noise_model.cpu()
+                torch.cuda.empty_cache()
+
+            if self.rank == 0:
+                videos = self.vae.decode(x0)
+
+        del noise, latent, x0
+        del sample_scheduler
+        if offload_model:
+            gc.collect()
+            torch.cuda.synchronize()
+        if dist.is_initialized():
+            dist.barrier()
+
+        return videos[0] if self.rank == 0 else None

wan/modules/__init__.py

@@ -0,0 +1,18 @@
+from .attention import flash_attention
+from .model import WanModel
+from .t5 import T5Decoder, T5Encoder, T5EncoderModel, T5Model
+from .tokenizers import HuggingfaceTokenizer
+from .vae2_1 import Wan2_1_VAE
+from .vae2_2 import Wan2_2_VAE
+
+__all__ = [
+    'Wan2_1_VAE',
+    'Wan2_2_VAE',
+    'WanModel',
+    'T5Model',
+    'T5Encoder',
+    'T5Decoder',
+    'T5EncoderModel',
+    'HuggingfaceTokenizer',
+    'flash_attention',
+]

wan/modules/animate/__init__.py

@@ -0,0 +1,4 @@
+
+from .model_animate import WanAnimateModel
+from .clip import CLIPModel
+__all__ = ['WanAnimateModel', 'CLIPModel']

wan/modules/animate/animate_utils.py

@@ -0,0 +1,143 @@
+
+import torch
+import numbers
+from peft import LoraConfig
+
+
+def get_loraconfig(transformer, rank=128, alpha=128, init_lora_weights="gaussian"):
+    target_modules = []
+    for name, module in transformer.named_modules():
+        if "blocks" in name and "face" not in name and "modulation" not in name and isinstance(module, torch.nn.Linear):
+            target_modules.append(name)
+
+    transformer_lora_config = LoraConfig(
+        r=rank,
+        lora_alpha=alpha,
+        init_lora_weights=init_lora_weights,
+        target_modules=target_modules,
+    )
+    return transformer_lora_config
+
+
+
+class TensorList(object):
+
+    def __init__(self, tensors):
+        """
+        tensors: a list of torch.Tensor objects. No need to have uniform shape.
+        """
+        assert isinstance(tensors, (list, tuple))
+        assert all(isinstance(u, torch.Tensor) for u in tensors)
+        assert len(set([u.ndim for u in tensors])) == 1
+        assert len(set([u.dtype for u in tensors])) == 1
+        assert len(set([u.device for u in tensors])) == 1
+        self.tensors = tensors
+    
+    def to(self, *args, **kwargs):
+        return TensorList([u.to(*args, **kwargs) for u in self.tensors])
+    
+    def size(self, dim):
+        assert dim == 0, 'only support get the 0th size'
+        return len(self.tensors)
+    
+    def pow(self, *args, **kwargs):
+        return TensorList([u.pow(*args, **kwargs) for u in self.tensors])
+    
+    def squeeze(self, dim):
+        assert dim != 0
+        if dim > 0:
+            dim -= 1
+        return TensorList([u.squeeze(dim) for u in self.tensors])
+    
+    def type(self, *args, **kwargs):
+        return TensorList([u.type(*args, **kwargs) for u in self.tensors])
+    
+    def type_as(self, other):
+        assert isinstance(other, (torch.Tensor, TensorList))
+        if isinstance(other, torch.Tensor):
+            return TensorList([u.type_as(other) for u in self.tensors])
+        else:
+            return TensorList([u.type(other.dtype) for u in self.tensors])
+    
+    @property
+    def dtype(self):
+        return self.tensors[0].dtype
+    
+    @property
+    def device(self):
+        return self.tensors[0].device
+    
+    @property
+    def ndim(self):
+        return 1 + self.tensors[0].ndim
+    
+    def __getitem__(self, index):
+        return self.tensors[index]
+    
+    def __len__(self):
+        return len(self.tensors)
+    
+    def __add__(self, other):
+        return self._apply(other, lambda u, v: u + v)
+    
+    def __radd__(self, other):
+        return self._apply(other, lambda u, v: v + u)
+    
+    def __sub__(self, other):
+        return self._apply(other, lambda u, v: u - v)
+    
+    def __rsub__(self, other):
+        return self._apply(other, lambda u, v: v - u)
+    
+    def __mul__(self, other):
+        return self._apply(other, lambda u, v: u * v)
+    
+    def __rmul__(self, other):
+        return self._apply(other, lambda u, v: v * u)
+    
+    def __floordiv__(self, other):
+        return self._apply(other, lambda u, v: u // v)
+    
+    def __truediv__(self, other):
+        return self._apply(other, lambda u, v: u / v)
+    
+    def __rfloordiv__(self, other):
+        return self._apply(other, lambda u, v: v // u)
+    
+    def __rtruediv__(self, other):
+        return self._apply(other, lambda u, v: v / u)
+    
+    def __pow__(self, other):
+        return self._apply(other, lambda u, v: u ** v)
+    
+    def __rpow__(self, other):
+        return self._apply(other, lambda u, v: v ** u)
+    
+    def __neg__(self):
+        return TensorList([-u for u in self.tensors])
+    
+    def __iter__(self):
+        for tensor in self.tensors:
+            yield tensor
+    
+    def __repr__(self):
+        return 'TensorList: \n' + repr(self.tensors)
+
+    def _apply(self, other, op):
+        if isinstance(other, (list, tuple, TensorList)) or (
+            isinstance(other, torch.Tensor) and (
+                other.numel() > 1 or other.ndim > 1
+            )
+        ):
+            assert len(other) == len(self.tensors)
+            return TensorList([op(u, v) for u, v in zip(self.tensors, other)])
+        elif isinstance(other, numbers.Number) or (
+            isinstance(other, torch.Tensor) and (
+                other.numel() == 1 and other.ndim <= 1
+            )
+        ):
+            return TensorList([op(u, other) for u in self.tensors])
+        else:
+            raise TypeError(
+                f'unsupported operand for *: "TensorList" and "{type(other)}"'
+            )

wan/modules/animate/clip.py

@@ -0,0 +1,542 @@
+# Modified from ``https://github.com/openai/CLIP'' and ``https://github.com/mlfoundations/open_clip''
+
+import logging
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as T
+
+from ..attention import flash_attention
+from ..tokenizers import HuggingfaceTokenizer
+from .xlm_roberta import XLMRoberta
+
+__all__ = [
+    'XLMRobertaCLIP',
+    'clip_xlm_roberta_vit_h_14',
+    'CLIPModel',
+]
+
+
+def pos_interpolate(pos, seq_len):
+    if pos.size(1) == seq_len:
+        return pos
+    else:
+        src_grid = int(math.sqrt(pos.size(1)))
+        tar_grid = int(math.sqrt(seq_len))
+        n = pos.size(1) - src_grid * src_grid
+        return torch.cat([
+            pos[:, :n],
+            F.interpolate(
+                pos[:, n:].float().reshape(1, src_grid, src_grid, -1).permute(
+                    0, 3, 1, 2),
+                size=(tar_grid, tar_grid),
+                mode='bicubic',
+                align_corners=False).flatten(2).transpose(1, 2)
+        ],
+                         dim=1)
+
+
+class QuickGELU(nn.Module):
+
+    def forward(self, x):
+        return x * torch.sigmoid(1.702 * x)
+
+
+class LayerNorm(nn.LayerNorm):
+
+    def forward(self, x):
+        return super().forward(x.float()).type_as(x)
+
+
+class SelfAttention(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 num_heads,
+                 causal=False,
+                 attn_dropout=0.0,
+                 proj_dropout=0.0):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.causal = causal
+        self.attn_dropout = attn_dropout
+        self.proj_dropout = proj_dropout
+
+        # layers
+        self.to_qkv = nn.Linear(dim, dim * 3)
+        self.proj = nn.Linear(dim, dim)
+
+    def forward(self, x):
+        """
+        x:   [B, L, C].
+        """
+        b, s, c, n, d = *x.size(), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q, k, v = self.to_qkv(x).view(b, s, 3, n, d).unbind(2)
+
+        # compute attention
+        p = self.attn_dropout if self.training else 0.0
+        x = flash_attention(q, k, v, dropout_p=p, causal=self.causal, version=2)
+        x = x.reshape(b, s, c)
+
+        # output
+        x = self.proj(x)
+        x = F.dropout(x, self.proj_dropout, self.training)
+        return x
+
+
+class SwiGLU(nn.Module):
+
+    def __init__(self, dim, mid_dim):
+        super().__init__()
+        self.dim = dim
+        self.mid_dim = mid_dim
+
+        # layers
+        self.fc1 = nn.Linear(dim, mid_dim)
+        self.fc2 = nn.Linear(dim, mid_dim)
+        self.fc3 = nn.Linear(mid_dim, dim)
+
+    def forward(self, x):
+        x = F.silu(self.fc1(x)) * self.fc2(x)
+        x = self.fc3(x)
+        return x
+
+
+class AttentionBlock(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 mlp_ratio,
+                 num_heads,
+                 post_norm=False,
+                 causal=False,
+                 activation='quick_gelu',
+                 attn_dropout=0.0,
+                 proj_dropout=0.0,
+                 norm_eps=1e-5):
+        assert activation in ['quick_gelu', 'gelu', 'swi_glu']
+        super().__init__()
+        self.dim = dim
+        self.mlp_ratio = mlp_ratio
+        self.num_heads = num_heads
+        self.post_norm = post_norm
+        self.causal = causal
+        self.norm_eps = norm_eps
+
+        # layers
+        self.norm1 = LayerNorm(dim, eps=norm_eps)
+        self.attn = SelfAttention(dim, num_heads, causal, attn_dropout,
+                                  proj_dropout)
+        self.norm2 = LayerNorm(dim, eps=norm_eps)
+        if activation == 'swi_glu':
+            self.mlp = SwiGLU(dim, int(dim * mlp_ratio))
+        else:
+            self.mlp = nn.Sequential(
+                nn.Linear(dim, int(dim * mlp_ratio)),
+                QuickGELU() if activation == 'quick_gelu' else nn.GELU(),
+                nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout))
+
+    def forward(self, x):
+        if self.post_norm:
+            x = x + self.norm1(self.attn(x))
+            x = x + self.norm2(self.mlp(x))
+        else:
+            x = x + self.attn(self.norm1(x))
+            x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class AttentionPool(nn.Module):
+
+    def __init__(self,
+                 dim,
+                 mlp_ratio,
+                 num_heads,
+                 activation='gelu',
+                 proj_dropout=0.0,
+                 norm_eps=1e-5):
+        assert dim % num_heads == 0
+        super().__init__()
+        self.dim = dim
+        self.mlp_ratio = mlp_ratio
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.proj_dropout = proj_dropout
+        self.norm_eps = norm_eps
+
+        # layers
+        gain = 1.0 / math.sqrt(dim)
+        self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))
+        self.to_q = nn.Linear(dim, dim)
+        self.to_kv = nn.Linear(dim, dim * 2)
+        self.proj = nn.Linear(dim, dim)
+        self.norm = LayerNorm(dim, eps=norm_eps)
+        self.mlp = nn.Sequential(
+            nn.Linear(dim, int(dim * mlp_ratio)),
+            QuickGELU() if activation == 'quick_gelu' else nn.GELU(),
+            nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout))
+
+    def forward(self, x):
+        """
+        x:  [B, L, C].
+        """
+        b, s, c, n, d = *x.size(), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.to_q(self.cls_embedding).view(1, 1, n, d).expand(b, -1, -1, -1)
+        k, v = self.to_kv(x).view(b, s, 2, n, d).unbind(2)
+
+        # compute attention
+        x = flash_attention(q, k, v, version=2)
+        x = x.reshape(b, 1, c)
+
+        # output
+        x = self.proj(x)
+        x = F.dropout(x, self.proj_dropout, self.training)
+
+        # mlp
+        x = x + self.mlp(self.norm(x))
+        return x[:, 0]
+
+
+class VisionTransformer(nn.Module):
+
+    def __init__(self,
+                 image_size=224,
+                 patch_size=16,
+                 dim=768,
+                 mlp_ratio=4,
+                 out_dim=512,
+                 num_heads=12,
+                 num_layers=12,
+                 pool_type='token',
+                 pre_norm=True,
+                 post_norm=False,
+                 activation='quick_gelu',
+                 attn_dropout=0.0,
+                 proj_dropout=0.0,
+                 embedding_dropout=0.0,
+                 norm_eps=1e-5):
+        if image_size % patch_size != 0:
+            print(
+                '[WARNING] image_size is not divisible by patch_size',
+                flush=True)
+        assert pool_type in ('token', 'token_fc', 'attn_pool')
+        out_dim = out_dim or dim
+        super().__init__()
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.num_patches = (image_size // patch_size)**2
+        self.dim = dim
+        self.mlp_ratio = mlp_ratio
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.pool_type = pool_type
+        self.post_norm = post_norm
+        self.norm_eps = norm_eps
+
+        # embeddings
+        gain = 1.0 / math.sqrt(dim)
+        self.patch_embedding = nn.Conv2d(
+            3,
+            dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=not pre_norm)
+        if pool_type in ('token', 'token_fc'):
+            self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))
+        self.pos_embedding = nn.Parameter(gain * torch.randn(
+            1, self.num_patches +
+            (1 if pool_type in ('token', 'token_fc') else 0), dim))
+        self.dropout = nn.Dropout(embedding_dropout)
+
+        # transformer
+        self.pre_norm = LayerNorm(dim, eps=norm_eps) if pre_norm else None
+        self.transformer = nn.Sequential(*[
+            AttentionBlock(dim, mlp_ratio, num_heads, post_norm, False,
+                           activation, attn_dropout, proj_dropout, norm_eps)
+            for _ in range(num_layers)
+        ])
+        self.post_norm = LayerNorm(dim, eps=norm_eps)
+
+        # head
+        if pool_type == 'token':
+            self.head = nn.Parameter(gain * torch.randn(dim, out_dim))
+        elif pool_type == 'token_fc':
+            self.head = nn.Linear(dim, out_dim)
+        elif pool_type == 'attn_pool':
+            self.head = AttentionPool(dim, mlp_ratio, num_heads, activation,
+                                      proj_dropout, norm_eps)
+
+    def forward(self, x, interpolation=False, use_31_block=False):
+        b = x.size(0)
+
+        # embeddings
+        x = self.patch_embedding(x).flatten(2).permute(0, 2, 1)
+        if self.pool_type in ('token', 'token_fc'):
+            x = torch.cat([self.cls_embedding.expand(b, -1, -1), x], dim=1)
+        if interpolation:
+            e = pos_interpolate(self.pos_embedding, x.size(1))
+        else:
+            e = self.pos_embedding
+        x = self.dropout(x + e)
+        if self.pre_norm is not None:
+            x = self.pre_norm(x)
+
+        # transformer
+        if use_31_block:
+            x = self.transformer[:-1](x)
+            return x
+        else:
+            x = self.transformer(x)
+            return x
+
+
+class XLMRobertaWithHead(XLMRoberta):
+
+    def __init__(self, **kwargs):
+        self.out_dim = kwargs.pop('out_dim')
+        super().__init__(**kwargs)
+
+        # head
+        mid_dim = (self.dim + self.out_dim) // 2
+        self.head = nn.Sequential(
+            nn.Linear(self.dim, mid_dim, bias=False), nn.GELU(),
+            nn.Linear(mid_dim, self.out_dim, bias=False))
+
+    def forward(self, ids):
+        # xlm-roberta
+        x = super().forward(ids)
+
+        # average pooling
+        mask = ids.ne(self.pad_id).unsqueeze(-1).to(x)
+        x = (x * mask).sum(dim=1) / mask.sum(dim=1)
+
+        # head
+        x = self.head(x)
+        return x
+
+
+class XLMRobertaCLIP(nn.Module):
+
+    def __init__(self,
+                 embed_dim=1024,
+                 image_size=224,
+                 patch_size=14,
+                 vision_dim=1280,
+                 vision_mlp_ratio=4,
+                 vision_heads=16,
+                 vision_layers=32,
+                 vision_pool='token',
+                 vision_pre_norm=True,
+                 vision_post_norm=False,
+                 activation='gelu',
+                 vocab_size=250002,
+                 max_text_len=514,
+                 type_size=1,
+                 pad_id=1,
+                 text_dim=1024,
+                 text_heads=16,
+                 text_layers=24,
+                 text_post_norm=True,
+                 text_dropout=0.1,
+                 attn_dropout=0.0,
+                 proj_dropout=0.0,
+                 embedding_dropout=0.0,
+                 norm_eps=1e-5):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.image_size = image_size
+        self.patch_size = patch_size
+        self.vision_dim = vision_dim
+        self.vision_mlp_ratio = vision_mlp_ratio
+        self.vision_heads = vision_heads
+        self.vision_layers = vision_layers
+        self.vision_pre_norm = vision_pre_norm
+        self.vision_post_norm = vision_post_norm
+        self.activation = activation
+        self.vocab_size = vocab_size
+        self.max_text_len = max_text_len
+        self.type_size = type_size
+        self.pad_id = pad_id
+        self.text_dim = text_dim
+        self.text_heads = text_heads
+        self.text_layers = text_layers
+        self.text_post_norm = text_post_norm
+        self.norm_eps = norm_eps
+
+        # models
+        self.visual = VisionTransformer(
+            image_size=image_size,
+            patch_size=patch_size,
+            dim=vision_dim,
+            mlp_ratio=vision_mlp_ratio,
+            out_dim=embed_dim,
+            num_heads=vision_heads,
+            num_layers=vision_layers,
+            pool_type=vision_pool,
+            pre_norm=vision_pre_norm,
+            post_norm=vision_post_norm,
+            activation=activation,
+            attn_dropout=attn_dropout,
+            proj_dropout=proj_dropout,
+            embedding_dropout=embedding_dropout,
+            norm_eps=norm_eps)
+        self.textual = XLMRobertaWithHead(
+            vocab_size=vocab_size,
+            max_seq_len=max_text_len,
+            type_size=type_size,
+            pad_id=pad_id,
+            dim=text_dim,
+            out_dim=embed_dim,
+            num_heads=text_heads,
+            num_layers=text_layers,
+            post_norm=text_post_norm,
+            dropout=text_dropout)
+        self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([]))
+
+    def forward(self, imgs, txt_ids):
+        """
+        imgs:       [B, 3, H, W] of torch.float32.
+        - mean:     [0.48145466, 0.4578275, 0.40821073]
+        - std:      [0.26862954, 0.26130258, 0.27577711]
+        txt_ids:    [B, L] of torch.long.
+                    Encoded by data.CLIPTokenizer.
+        """
+        xi = self.visual(imgs)
+        xt = self.textual(txt_ids)
+        return xi, xt
+
+    def param_groups(self):
+        groups = [{
+            'params': [
+                p for n, p in self.named_parameters()
+                if 'norm' in n or n.endswith('bias')
+            ],
+            'weight_decay': 0.0
+        }, {
+            'params': [
+                p for n, p in self.named_parameters()
+                if not ('norm' in n or n.endswith('bias'))
+            ]
+        }]
+        return groups
+
+
+def _clip(pretrained=False,
+          pretrained_name=None,
+          model_cls=XLMRobertaCLIP,
+          return_transforms=False,
+          return_tokenizer=False,
+          tokenizer_padding='eos',
+          dtype=torch.float32,
+          device='cpu',
+          **kwargs):
+    # init a model on device
+    with torch.device(device):
+        model = model_cls(**kwargs)
+
+    # set device
+    model = model.to(dtype=dtype, device=device)
+    output = (model,)
+
+    # init transforms
+    if return_transforms:
+        # mean and std
+        if 'siglip' in pretrained_name.lower():
+            mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]
+        else:
+            mean = [0.48145466, 0.4578275, 0.40821073]
+            std = [0.26862954, 0.26130258, 0.27577711]
+
+        # transforms
+        transforms = T.Compose([
+            T.Resize((model.image_size, model.image_size),
+                     interpolation=T.InterpolationMode.BICUBIC),
+            T.ToTensor(),
+            T.Normalize(mean=mean, std=std)
+        ])
+        output += (transforms,)
+    return output[0] if len(output) == 1 else output
+
+
+def clip_xlm_roberta_vit_h_14(
+        pretrained=False,
+        pretrained_name='open-clip-xlm-roberta-large-vit-huge-14',
+        **kwargs):
+    cfg = dict(
+        embed_dim=1024,
+        image_size=224,
+        patch_size=14,
+        vision_dim=1280,
+        vision_mlp_ratio=4,
+        vision_heads=16,
+        vision_layers=32,
+        vision_pool='token',
+        activation='gelu',
+        vocab_size=250002,
+        max_text_len=514,
+        type_size=1,
+        pad_id=1,
+        text_dim=1024,
+        text_heads=16,
+        text_layers=24,
+        text_post_norm=True,
+        text_dropout=0.1,
+        attn_dropout=0.0,
+        proj_dropout=0.0,
+        embedding_dropout=0.0)
+    cfg.update(**kwargs)
+    return _clip(pretrained, pretrained_name, XLMRobertaCLIP, **cfg)
+
+
+class CLIPModel:
+
+    def __init__(self, dtype, device, checkpoint_path, tokenizer_path):
+        self.dtype = dtype
+        self.device = device
+        self.checkpoint_path = checkpoint_path
+        self.tokenizer_path = tokenizer_path
+
+        # init model
+        self.model, self.transforms = clip_xlm_roberta_vit_h_14(
+            pretrained=False,
+            return_transforms=True,
+            return_tokenizer=False,
+            dtype=dtype,
+            device=device)
+        self.model = self.model.eval().requires_grad_(False)
+        logging.info(f'loading {checkpoint_path}')
+        self.model.load_state_dict(
+            torch.load(checkpoint_path, map_location='cpu'))
+
+        # init tokenizer
+        self.tokenizer = HuggingfaceTokenizer(
+            name=tokenizer_path,
+            seq_len=self.model.max_text_len - 2,
+            clean='whitespace')
+
+    def visual(self, videos):
+        # preprocess
+        size = (self.model.image_size,) * 2
+        videos = torch.cat([
+            F.interpolate(
+                u.transpose(0, 1),
+                size=size,
+                mode='bicubic',
+                align_corners=False) for u in videos
+        ])
+        videos = self.transforms.transforms[-1](videos.mul_(0.5).add_(0.5))
+
+        # forward
+        with torch.cuda.amp.autocast(dtype=self.dtype):
+            out = self.model.visual(videos, use_31_block=True)
+            return out

wan/modules/animate/face_blocks.py

@@ -0,0 +1,383 @@
+
+from torch import nn
+import torch
+from typing import Tuple, Optional
+from einops import rearrange
+import torch.nn.functional as F
+import math
+from ...distributed.util import gather_forward, get_rank, get_world_size
+
+
+try:
+    from flash_attn import flash_attn_qkvpacked_func, flash_attn_func
+except ImportError:
+    flash_attn_func = None
+
+MEMORY_LAYOUT = {
+    "flash": (
+        lambda x: x.view(x.shape[0] * x.shape[1], *x.shape[2:]),
+        lambda x: x,
+    ),
+    "torch": (
+        lambda x: x.transpose(1, 2),
+        lambda x: x.transpose(1, 2),
+    ),
+    "vanilla": (
+        lambda x: x.transpose(1, 2),
+        lambda x: x.transpose(1, 2),
+    ),
+}
+
+
+def attention(
+    q,
+    k,
+    v,
+    mode="flash",
+    drop_rate=0,
+    attn_mask=None,
+    causal=False,
+    max_seqlen_q=None,
+    batch_size=1,
+):
+    """
+    Perform QKV self attention.
+
+    Args:
+        q (torch.Tensor): Query tensor with shape [b, s, a, d], where a is the number of heads.
+        k (torch.Tensor): Key tensor with shape [b, s1, a, d]
+        v (torch.Tensor): Value tensor with shape [b, s1, a, d]
+        mode (str): Attention mode. Choose from 'self_flash', 'cross_flash', 'torch', and 'vanilla'.
+        drop_rate (float): Dropout rate in attention map. (default: 0)
+        attn_mask (torch.Tensor): Attention mask with shape [b, s1] (cross_attn), or [b, a, s, s1] (torch or vanilla).
+            (default: None)
+        causal (bool): Whether to use causal attention. (default: False)
+        cu_seqlens_q (torch.Tensor): dtype torch.int32. The cumulative sequence lengths of the sequences in the batch,
+            used to index into q.
+        cu_seqlens_kv (torch.Tensor): dtype torch.int32. The cumulative sequence lengths of the sequences in the batch,
+            used to index into kv.
+        max_seqlen_q (int): The maximum sequence length in the batch of q.
+        max_seqlen_kv (int): The maximum sequence length in the batch of k and v.
+
+    Returns:
+        torch.Tensor: Output tensor after self attention with shape [b, s, ad]
+    """
+    pre_attn_layout, post_attn_layout = MEMORY_LAYOUT[mode]
+
+    if mode == "torch":
+        if attn_mask is not None and attn_mask.dtype != torch.bool:
+            attn_mask = attn_mask.to(q.dtype)
+        x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal)
+
+    elif mode == "flash":
+        x = flash_attn_func(
+            q,
+            k,
+            v,
+        )
+        x = x.view(batch_size, max_seqlen_q, x.shape[-2], x.shape[-1])  # reshape x to [b, s, a, d]
+    elif mode == "vanilla":
+        scale_factor = 1 / math.sqrt(q.size(-1))
+
+        b, a, s, _ = q.shape
+        s1 = k.size(2)
+        attn_bias = torch.zeros(b, a, s, s1, dtype=q.dtype, device=q.device)
+        if causal:
+            # Only applied to self attention
+            assert attn_mask is None, "Causal mask and attn_mask cannot be used together"
+            temp_mask = torch.ones(b, a, s, s, dtype=torch.bool, device=q.device).tril(diagonal=0)
+            attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+            attn_bias.to(q.dtype)
+
+        if attn_mask is not None:
+            if attn_mask.dtype == torch.bool:
+                attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
+            else:
+                attn_bias += attn_mask
+
+        attn = (q @ k.transpose(-2, -1)) * scale_factor
+        attn += attn_bias
+        attn = attn.softmax(dim=-1)
+        attn = torch.dropout(attn, p=drop_rate, train=True)
+        x = attn @ v
+    else:
+        raise NotImplementedError(f"Unsupported attention mode: {mode}")
+
+    x = post_attn_layout(x)
+    b, s, a, d = x.shape
+    out = x.reshape(b, s, -1)
+    return out
+
+
+class CausalConv1d(nn.Module):
+
+    def __init__(self, chan_in, chan_out, kernel_size=3, stride=1, dilation=1, pad_mode="replicate", **kwargs):
+        super().__init__()
+
+        self.pad_mode = pad_mode
+        padding = (kernel_size - 1, 0)  # T
+        self.time_causal_padding = padding
+
+        self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs)
+
+    def forward(self, x):
+        x = F.pad(x, self.time_causal_padding, mode=self.pad_mode)
+        return self.conv(x)
+
+
+
+class FaceEncoder(nn.Module):
+    def __init__(self, in_dim: int, hidden_dim: int, num_heads=int, dtype=None, device=None):
+        factory_kwargs = {"dtype": dtype, "device": device}
+        super().__init__()
+
+        self.num_heads = num_heads
+        self.conv1_local = CausalConv1d(in_dim, 1024 * num_heads, 3, stride=1)
+        self.norm1 = nn.LayerNorm(hidden_dim // 8, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+        self.act = nn.SiLU()
+        self.conv2 = CausalConv1d(1024, 1024, 3, stride=2)
+        self.conv3 = CausalConv1d(1024, 1024, 3, stride=2)
+
+        self.out_proj = nn.Linear(1024, hidden_dim)
+        self.norm1 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+
+        self.norm2 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+
+        self.norm3 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+
+        self.padding_tokens = nn.Parameter(torch.zeros(1, 1, 1, hidden_dim))
+
+    def forward(self, x):
+        
+        x = rearrange(x, "b t c -> b c t")
+        b, c, t = x.shape
+
+        x = self.conv1_local(x)
+        x = rearrange(x, "b (n c) t -> (b n) t c", n=self.num_heads)
+        
+        x = self.norm1(x)
+        x = self.act(x)
+        x = rearrange(x, "b t c -> b c t")
+        x = self.conv2(x)
+        x = rearrange(x, "b c t -> b t c")
+        x = self.norm2(x)
+        x = self.act(x)
+        x = rearrange(x, "b t c -> b c t")
+        x = self.conv3(x)
+        x = rearrange(x, "b c t -> b t c")
+        x = self.norm3(x)
+        x = self.act(x)
+        x = self.out_proj(x)
+        x = rearrange(x, "(b n) t c -> b t n c", b=b)
+        padding = self.padding_tokens.repeat(b, x.shape[1], 1, 1)
+        x = torch.cat([x, padding], dim=-2)
+        x_local = x.clone()
+
+        return x_local
+
+
+
+class RMSNorm(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        elementwise_affine=True,
+        eps: float = 1e-6,
+        device=None,
+        dtype=None,
+    ):
+        """
+        Initialize the RMSNorm normalization layer.
+
+        Args:
+            dim (int): The dimension of the input tensor.
+            eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
+
+        Attributes:
+            eps (float): A small value added to the denominator for numerical stability.
+            weight (nn.Parameter): Learnable scaling parameter.
+
+        """
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+        self.eps = eps
+        if elementwise_affine:
+            self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))
+
+    def _norm(self, x):
+        """
+        Apply the RMSNorm normalization to the input tensor.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The normalized tensor.
+
+        """
+        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+    def forward(self, x):
+        """
+        Forward pass through the RMSNorm layer.
+
+        Args:
+            x (torch.Tensor): The input tensor.
+
+        Returns:
+            torch.Tensor: The output tensor after applying RMSNorm.
+
+        """
+        output = self._norm(x.float()).type_as(x)
+        if hasattr(self, "weight"):
+            output = output * self.weight
+        return output
+
+
+def get_norm_layer(norm_layer):
+    """
+    Get the normalization layer.
+
+    Args:
+        norm_layer (str): The type of normalization layer.
+
+    Returns:
+        norm_layer (nn.Module): The normalization layer.
+    """
+    if norm_layer == "layer":
+        return nn.LayerNorm
+    elif norm_layer == "rms":
+        return RMSNorm
+    else:
+        raise NotImplementedError(f"Norm layer {norm_layer} is not implemented")
+
+
+class FaceAdapter(nn.Module):
+    def __init__(
+        self,
+        hidden_dim: int,
+        heads_num: int,
+        qk_norm: bool = True,
+        qk_norm_type: str = "rms",
+        num_adapter_layers: int = 1,
+        dtype=None,
+        device=None,
+    ):
+
+        factory_kwargs = {"dtype": dtype, "device": device}
+        super().__init__()
+        self.hidden_size = hidden_dim
+        self.heads_num = heads_num
+        self.fuser_blocks = nn.ModuleList(
+            [
+                FaceBlock(
+                    self.hidden_size,
+                    self.heads_num,
+                    qk_norm=qk_norm,
+                    qk_norm_type=qk_norm_type,
+                    **factory_kwargs,
+                )
+                for _ in range(num_adapter_layers)
+            ]
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        motion_embed: torch.Tensor,
+        idx: int,
+        freqs_cis_q: Tuple[torch.Tensor, torch.Tensor] = None,
+        freqs_cis_k: Tuple[torch.Tensor, torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        return self.fuser_blocks[idx](x, motion_embed, freqs_cis_q, freqs_cis_k)
+
+
+
+class FaceBlock(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        heads_num: int,
+        qk_norm: bool = True,
+        qk_norm_type: str = "rms",
+        qk_scale: float = None,
+        dtype: Optional[torch.dtype] = None,
+        device: Optional[torch.device] = None,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+
+        self.deterministic = False
+        self.hidden_size = hidden_size
+        self.heads_num = heads_num
+        head_dim = hidden_size // heads_num
+        self.scale = qk_scale or head_dim**-0.5
+       
+        self.linear1_kv = nn.Linear(hidden_size, hidden_size * 2, **factory_kwargs)
+        self.linear1_q = nn.Linear(hidden_size, hidden_size, **factory_kwargs)
+
+        self.linear2 = nn.Linear(hidden_size, hidden_size, **factory_kwargs)
+
+        qk_norm_layer = get_norm_layer(qk_norm_type)
+        self.q_norm = (
+            qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity()
+        )
+        self.k_norm = (
+            qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity()
+        )
+
+        self.pre_norm_feat = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+
+        self.pre_norm_motion = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        motion_vec: torch.Tensor,
+        motion_mask: Optional[torch.Tensor] = None,
+        use_context_parallel=False,
+    ) -> torch.Tensor:
+        
+        B, T, N, C = motion_vec.shape
+        T_comp = T
+
+        x_motion = self.pre_norm_motion(motion_vec)
+        x_feat = self.pre_norm_feat(x)
+
+        kv = self.linear1_kv(x_motion)
+        q = self.linear1_q(x_feat)
+
+        k, v = rearrange(kv, "B L N (K H D) -> K B L N H D", K=2, H=self.heads_num)
+        q = rearrange(q, "B S (H D) -> B S H D", H=self.heads_num)
+
+        # Apply QK-Norm if needed.
+        q = self.q_norm(q).to(v)
+        k = self.k_norm(k).to(v)
+
+        k = rearrange(k, "B L N H D -> (B L) N H D")  
+        v = rearrange(v, "B L N H D -> (B L) N H D") 
+
+        if use_context_parallel:
+            q = gather_forward(q, dim=1)
+
+        q = rearrange(q, "B (L S) H D -> (B L) S H D", L=T_comp)  
+        # Compute attention.
+        attn = attention(
+            q,
+            k,
+            v,
+            max_seqlen_q=q.shape[1],
+            batch_size=q.shape[0],
+        )
+
+        attn = rearrange(attn, "(B L) S C -> B (L S) C", L=T_comp)
+        if use_context_parallel:
+            attn = torch.chunk(attn, get_world_size(), dim=1)[get_rank()]
+
+        output = self.linear2(attn)
+
+        if motion_mask is not None:
+            output = output * rearrange(motion_mask, "B T H W -> B (T H W)").unsqueeze(-1)
+
+        return output

wan/modules/animate/model_animate.py

@@ -0,0 +1,500 @@
+
+import math
+import types
+from copy import deepcopy
+from einops import  rearrange
+from typing import List
+import numpy as np
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.loaders import PeftAdapterMixin
+
+from ...distributed.sequence_parallel import (
+    distributed_attention,
+    gather_forward,
+    get_rank,
+    get_world_size,
+)
+
+
+from ..model import (
+    Head,
+    WanAttentionBlock,
+    WanLayerNorm,
+    WanRMSNorm,
+    WanModel,
+    WanSelfAttention,
+    flash_attention,
+    rope_params,
+    sinusoidal_embedding_1d,
+    rope_apply
+)
+
+from .face_blocks import FaceEncoder, FaceAdapter
+from .motion_encoder import Generator
+
+class HeadAnimate(Head):
+
+    def forward(self, x, e):
+        """
+        Args:
+            x(Tensor): Shape [B, L1, C]
+            e(Tensor): Shape [B, L1, C]
+        """
+        assert e.dtype == torch.float32
+        with amp.autocast(dtype=torch.float32):
+            e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
+            x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))
+        return x
+
+
+class WanAnimateSelfAttention(WanSelfAttention):
+
+    def forward(self, x, seq_lens, grid_sizes, freqs):
+        """
+        Args:
+            x(Tensor): Shape [B, L, num_heads, C / num_heads]
+            seq_lens(Tensor): Shape [B]
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+
+        # query, key, value function
+        def qkv_fn(x):
+            q = self.norm_q(self.q(x)).view(b, s, n, d)
+            k = self.norm_k(self.k(x)).view(b, s, n, d)
+            v = self.v(x).view(b, s, n, d)
+            return q, k, v
+
+        q, k, v = qkv_fn(x)
+
+        x = flash_attention(
+            q=rope_apply(q, grid_sizes, freqs),
+            k=rope_apply(k, grid_sizes, freqs),
+            v=v,
+            k_lens=seq_lens,
+            window_size=self.window_size)
+
+        # output
+        x = x.flatten(2)
+        x = self.o(x)
+        return x
+
+
+class WanAnimateCrossAttention(WanSelfAttention):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        window_size=(-1, -1),
+        qk_norm=True,
+        eps=1e-6,
+        use_img_emb=True
+    ):
+        super().__init__(
+            dim,
+            num_heads,
+            window_size,
+            qk_norm,
+            eps
+        )
+        self.use_img_emb = use_img_emb
+
+        if use_img_emb:
+            self.k_img = nn.Linear(dim, dim)
+            self.v_img = nn.Linear(dim, dim)
+            self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+    def forward(self, x, context, context_lens):
+        """
+        x:              [B, L1, C].
+        context:        [B, L2, C].
+        context_lens:   [B].
+        """
+        if self.use_img_emb:
+            context_img = context[:, :257]
+            context = context[:, 257:]
+        else:
+            context = context
+
+        b, n, d = x.size(0), self.num_heads, self.head_dim
+
+        # compute query, key, value
+        q = self.norm_q(self.q(x)).view(b, -1, n, d)
+        k = self.norm_k(self.k(context)).view(b, -1, n, d)
+        v = self.v(context).view(b, -1, n, d)
+
+        if self.use_img_emb:
+            k_img = self.norm_k_img(self.k_img(context_img)).view(b, -1, n, d)
+            v_img = self.v_img(context_img).view(b, -1, n, d)
+            img_x = flash_attention(q, k_img, v_img, k_lens=None)
+        # compute attention
+        x = flash_attention(q, k, v, k_lens=context_lens)
+
+        # output
+        x = x.flatten(2)
+
+        if self.use_img_emb:
+            img_x = img_x.flatten(2)
+            x = x + img_x
+
+        x = self.o(x)
+        return x
+
+
+class WanAnimateAttentionBlock(nn.Module):
+    def __init__(self,
+                 dim,
+                 ffn_dim,
+                 num_heads,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 cross_attn_norm=True,
+                 eps=1e-6,
+                 use_img_emb=True):
+
+        super().__init__()
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.num_heads = num_heads
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+
+        # layers
+        self.norm1 = WanLayerNorm(dim, eps)
+        self.self_attn = WanAnimateSelfAttention(dim, num_heads, window_size, qk_norm, eps)
+        
+        self.norm3 = WanLayerNorm(
+            dim, eps, elementwise_affine=True
+        ) if cross_attn_norm else nn.Identity()
+            
+        self.cross_attn = WanAnimateCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps, use_img_emb=use_img_emb)
+        self.norm2 = WanLayerNorm(dim, eps)
+        self.ffn = nn.Sequential(
+            nn.Linear(dim, ffn_dim),
+            nn.GELU(approximate='tanh'),
+            nn.Linear(ffn_dim, dim)
+        )
+
+        # modulation
+        self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim ** 0.5)
+
+    def forward(
+        self,
+        x,
+        e,
+        seq_lens,
+        grid_sizes,
+        freqs,
+        context,
+        context_lens,
+    ):
+        """
+        Args:
+            x(Tensor): Shape [B, L, C]
+            e(Tensor): Shape [B, L1, 6, C]
+            seq_lens(Tensor): Shape [B], length of each sequence in batch
+            grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+            freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+        """
+        assert e.dtype == torch.float32
+        with amp.autocast(dtype=torch.float32):
+            e = (self.modulation + e).chunk(6, dim=1)
+        assert e[0].dtype == torch.float32
+
+        # self-attention
+        y = self.self_attn(
+            self.norm1(x).float() * (1 + e[1]) + e[0], seq_lens, grid_sizes, freqs
+        )
+        with amp.autocast(dtype=torch.float32):
+            x = x + y * e[2]
+
+        # cross-attention & ffn function
+        def cross_attn_ffn(x, context, context_lens, e):
+            x = x + self.cross_attn(self.norm3(x), context, context_lens)
+            y = self.ffn(self.norm2(x).float() * (1 + e[4]) + e[3])
+            with amp.autocast(dtype=torch.float32):
+                x = x + y * e[5]
+            return x
+
+        x = cross_attn_ffn(x, context, context_lens, e)
+        return x
+
+
+class MLPProj(torch.nn.Module):
+    def __init__(self, in_dim, out_dim):
+        super().__init__()
+
+        self.proj = torch.nn.Sequential(
+            torch.nn.LayerNorm(in_dim),
+            torch.nn.Linear(in_dim, in_dim),
+            torch.nn.GELU(),
+            torch.nn.Linear(in_dim, out_dim),
+            torch.nn.LayerNorm(out_dim),
+        )
+
+    def forward(self, image_embeds):
+        clip_extra_context_tokens = self.proj(image_embeds)
+        return clip_extra_context_tokens
+
+class WanAnimateModel(ModelMixin, ConfigMixin, PeftAdapterMixin):
+    _no_split_modules = ['WanAttentionBlock']
+
+    @register_to_config
+    def __init__(self,
+                 patch_size=(1, 2, 2),
+                 text_len=512,
+                 in_dim=36,
+                 dim=5120,
+                 ffn_dim=13824,
+                 freq_dim=256,
+                 text_dim=4096,
+                 out_dim=16,
+                 num_heads=40,
+                 num_layers=40,
+                 window_size=(-1, -1),
+                 qk_norm=True,
+                 cross_attn_norm=True,
+                 eps=1e-6,
+                 motion_encoder_dim=512,
+                 use_context_parallel=False,
+                 use_img_emb=True):
+
+        super().__init__()
+        self.patch_size = patch_size
+        self.text_len = text_len
+        self.in_dim = in_dim
+        self.dim = dim
+        self.ffn_dim = ffn_dim
+        self.freq_dim = freq_dim
+        self.text_dim = text_dim
+        self.out_dim = out_dim
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.window_size = window_size
+        self.qk_norm = qk_norm
+        self.cross_attn_norm = cross_attn_norm
+        self.eps = eps
+        self.motion_encoder_dim = motion_encoder_dim
+        self.use_context_parallel = use_context_parallel
+        self.use_img_emb = use_img_emb
+
+        # embeddings
+        self.patch_embedding = nn.Conv3d(
+            in_dim, dim, kernel_size=patch_size, stride=patch_size)
+
+        self.pose_patch_embedding = nn.Conv3d(
+            16, dim, kernel_size=patch_size, stride=patch_size
+        )
+
+        self.text_embedding = nn.Sequential(
+            nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
+            nn.Linear(dim, dim))
+
+        self.time_embedding = nn.Sequential(
+            nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
+        self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
+
+        # blocks
+        self.blocks = nn.ModuleList([
+            WanAnimateAttentionBlock(dim, ffn_dim, num_heads, window_size, qk_norm,
+                              cross_attn_norm, eps, use_img_emb) for _ in range(num_layers)
+        ])
+
+        # head
+        self.head = HeadAnimate(dim, out_dim, patch_size, eps)
+
+        # buffers (don't use register_buffer otherwise dtype will be changed in to())
+        assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
+        d = dim // num_heads
+        self.freqs = torch.cat([
+            rope_params(1024, d - 4 * (d // 6)),
+            rope_params(1024, 2 * (d // 6)),
+            rope_params(1024, 2 * (d // 6))
+        ], dim=1)
+
+        self.img_emb = MLPProj(1280, dim)
+        
+        # initialize weights
+        self.init_weights()
+
+        self.motion_encoder = Generator(size=512, style_dim=512, motion_dim=20)
+        self.face_adapter = FaceAdapter(
+            heads_num=self.num_heads,
+            hidden_dim=self.dim,
+            num_adapter_layers=self.num_layers // 5,
+        )
+
+        self.face_encoder = FaceEncoder(
+            in_dim=motion_encoder_dim,
+            hidden_dim=self.dim,
+            num_heads=4,
+        )
+
+    def after_patch_embedding(self, x: List[torch.Tensor], pose_latents, face_pixel_values):
+        pose_latents = [self.pose_patch_embedding(u.unsqueeze(0)) for u in pose_latents]
+        for x_, pose_latents_ in zip(x, pose_latents):
+            x_[:, :, 1:] += pose_latents_
+        
+        b,c,T,h,w = face_pixel_values.shape
+        face_pixel_values = rearrange(face_pixel_values, "b c t h w -> (b t) c h w")
+
+        encode_bs = 8
+        face_pixel_values_tmp = []
+        for i in range(math.ceil(face_pixel_values.shape[0]/encode_bs)):
+            face_pixel_values_tmp.append(self.motion_encoder.get_motion(face_pixel_values[i*encode_bs:(i+1)*encode_bs]))
+
+        motion_vec = torch.cat(face_pixel_values_tmp)
+        
+        motion_vec = rearrange(motion_vec, "(b t) c -> b t c", t=T)
+        motion_vec = self.face_encoder(motion_vec)
+
+        B, L, H, C = motion_vec.shape
+        pad_face = torch.zeros(B, 1, H, C).type_as(motion_vec)
+        motion_vec = torch.cat([pad_face, motion_vec], dim=1)
+        return x, motion_vec
+
+
+    def after_transformer_block(self, block_idx, x, motion_vec, motion_masks=None):
+        if block_idx % 5 == 0:
+            adapter_args = [x, motion_vec, motion_masks, self.use_context_parallel]
+            residual_out = self.face_adapter.fuser_blocks[block_idx // 5](*adapter_args)
+            x = residual_out + x
+        return x
+
+
+    def forward(
+        self,
+        x,
+        t,
+        clip_fea,
+        context,
+        seq_len,
+        y=None,
+        pose_latents=None, 
+        face_pixel_values=None
+    ):
+        # params
+        device = self.patch_embedding.weight.device
+        if self.freqs.device != device:
+            self.freqs = self.freqs.to(device)
+
+        if y is not None:
+            x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+        # embeddings
+        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+        x, motion_vec = self.after_patch_embedding(x, pose_latents, face_pixel_values)
+
+        grid_sizes = torch.stack(
+            [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+        x = [u.flatten(2).transpose(1, 2) for u in x]
+        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+        assert seq_lens.max() <= seq_len
+        x = torch.cat([
+            torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
+                      dim=1) for u in x
+        ])
+
+        # time embeddings
+        with amp.autocast(dtype=torch.float32):
+            e = self.time_embedding(
+                sinusoidal_embedding_1d(self.freq_dim, t).float()
+            )
+            e0 = self.time_projection(e).unflatten(1, (6, self.dim))
+            assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+        # context
+        context_lens = None
+        context = self.text_embedding(
+            torch.stack([
+                torch.cat(
+                    [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+                for u in context
+            ]))
+
+        if self.use_img_emb:
+            context_clip = self.img_emb(clip_fea) # bs x 257 x dim
+            context = torch.concat([context_clip, context], dim=1)
+
+        # arguments
+        kwargs = dict(
+            e=e0,
+            seq_lens=seq_lens,
+            grid_sizes=grid_sizes,
+            freqs=self.freqs,
+            context=context,
+            context_lens=context_lens)
+
+        if self.use_context_parallel:
+            x = torch.chunk(x, get_world_size(), dim=1)[get_rank()]
+
+        for idx, block in enumerate(self.blocks):
+            x = block(x, **kwargs)
+            x = self.after_transformer_block(idx, x, motion_vec)
+
+        # head
+        x = self.head(x, e)
+
+        if self.use_context_parallel:
+            x = gather_forward(x, dim=1)
+
+        # unpatchify
+        x = self.unpatchify(x, grid_sizes)
+        return [u.float() for u in x]
+
+
+    def unpatchify(self, x, grid_sizes):
+        r"""
+        Reconstruct video tensors from patch embeddings.
+
+        Args:
+            x (List[Tensor]):
+                List of patchified features, each with shape [L, C_out * prod(patch_size)]
+            grid_sizes (Tensor):
+                Original spatial-temporal grid dimensions before patching,
+                    shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
+
+        Returns:
+            List[Tensor]:
+                Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
+        """
+
+        c = self.out_dim
+        out = []
+        for u, v in zip(x, grid_sizes.tolist()):
+            u = u[:math.prod(v)].view(*v, *self.patch_size, c)
+            u = torch.einsum('fhwpqrc->cfphqwr', u)
+            u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
+            out.append(u)
+        return out
+
+    def init_weights(self):
+        r"""
+        Initialize model parameters using Xavier initialization.
+        """
+
+        # basic init
+        for m in self.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.xavier_uniform_(m.weight)
+                if m.bias is not None:
+                    nn.init.zeros_(m.bias)
+
+        # init embeddings
+        nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
+        for m in self.text_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=.02)
+        for m in self.time_embedding.modules():
+            if isinstance(m, nn.Linear):
+                nn.init.normal_(m.weight, std=.02)
+
+        # init output layer
+        nn.init.zeros_(self.head.head.weight)

wan/modules/animate/motion_encoder.py

@@ -0,0 +1,307 @@
+# Modified from ``https://github.com/wyhsirius/LIA``
+
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import math
+
+def custom_qr(input_tensor):
+    original_dtype = input_tensor.dtype
+    if original_dtype == torch.bfloat16:
+        q, r = torch.linalg.qr(input_tensor.to(torch.float32))
+        return q.to(original_dtype), r.to(original_dtype)
+    return torch.linalg.qr(input_tensor)
+
+def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
+	return F.leaky_relu(input + bias, negative_slope) * scale
+
+
+def upfirdn2d_native(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1):
+	_, minor, in_h, in_w = input.shape
+	kernel_h, kernel_w = kernel.shape
+
+	out = input.view(-1, minor, in_h, 1, in_w, 1)
+	out = F.pad(out, [0, up_x - 1, 0, 0, 0, up_y - 1, 0, 0])
+	out = out.view(-1, minor, in_h * up_y, in_w * up_x)
+
+	out = F.pad(out, [max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)])
+	out = out[:, :, max(-pad_y0, 0): out.shape[2] - max(-pad_y1, 0),
+		  max(-pad_x0, 0): out.shape[3] - max(-pad_x1, 0), ]
+
+	out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1])
+	w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
+	out = F.conv2d(out, w)
+	out = out.reshape(-1, minor, in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
+					  in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1, )
+	return out[:, :, ::down_y, ::down_x]
+
+
+def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
+	return upfirdn2d_native(input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1])
+
+
+def make_kernel(k):
+	k = torch.tensor(k, dtype=torch.float32)
+	if k.ndim == 1:
+		k = k[None, :] * k[:, None]
+	k /= k.sum()
+	return k
+
+
+class FusedLeakyReLU(nn.Module):
+	def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
+		super().__init__()
+		self.bias = nn.Parameter(torch.zeros(1, channel, 1, 1))
+		self.negative_slope = negative_slope
+		self.scale = scale
+
+	def forward(self, input):
+		out = fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)
+		return out
+
+
+class Blur(nn.Module):
+	def __init__(self, kernel, pad, upsample_factor=1):
+		super().__init__()
+
+		kernel = make_kernel(kernel)
+
+		if upsample_factor > 1:
+			kernel = kernel * (upsample_factor ** 2)
+
+		self.register_buffer('kernel', kernel)
+
+		self.pad = pad
+
+	def forward(self, input):
+		return upfirdn2d(input, self.kernel, pad=self.pad)
+
+
+class ScaledLeakyReLU(nn.Module):
+	def __init__(self, negative_slope=0.2):
+		super().__init__()
+
+		self.negative_slope = negative_slope
+
+	def forward(self, input):
+		return F.leaky_relu(input, negative_slope=self.negative_slope)
+
+
+class EqualConv2d(nn.Module):
+	def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True):
+		super().__init__()
+
+		self.weight = nn.Parameter(torch.randn(out_channel, in_channel, kernel_size, kernel_size))
+		self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)
+
+		self.stride = stride
+		self.padding = padding
+
+		if bias:
+			self.bias = nn.Parameter(torch.zeros(out_channel))
+		else:
+			self.bias = None
+
+	def forward(self, input):
+
+		return F.conv2d(input, self.weight * self.scale, bias=self.bias, stride=self.stride, padding=self.padding)
+
+	def __repr__(self):
+		return (
+			f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},'
+			f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})'
+		)
+
+
+class EqualLinear(nn.Module):
+	def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None):
+		super().__init__()
+
+		self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))
+
+		if bias:
+			self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))
+		else:
+			self.bias = None
+
+		self.activation = activation
+
+		self.scale = (1 / math.sqrt(in_dim)) * lr_mul
+		self.lr_mul = lr_mul
+
+	def forward(self, input):
+
+		if self.activation:
+			out = F.linear(input, self.weight * self.scale)
+			out = fused_leaky_relu(out, self.bias * self.lr_mul)
+		else:
+			out = F.linear(input, self.weight * self.scale, bias=self.bias * self.lr_mul)
+
+		return out
+
+	def __repr__(self):
+		return (f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})')
+
+
+class ConvLayer(nn.Sequential):
+	def __init__(
+			self,
+			in_channel,
+			out_channel,
+			kernel_size,
+			downsample=False,
+			blur_kernel=[1, 3, 3, 1],
+			bias=True,
+			activate=True,
+	):
+		layers = []
+
+		if downsample:
+			factor = 2
+			p = (len(blur_kernel) - factor) + (kernel_size - 1)
+			pad0 = (p + 1) // 2
+			pad1 = p // 2
+
+			layers.append(Blur(blur_kernel, pad=(pad0, pad1)))
+
+			stride = 2
+			self.padding = 0
+
+		else:
+			stride = 1
+			self.padding = kernel_size // 2
+
+		layers.append(EqualConv2d(in_channel, out_channel, kernel_size, padding=self.padding, stride=stride,
+								  bias=bias and not activate))
+
+		if activate:
+			if bias:
+				layers.append(FusedLeakyReLU(out_channel))
+			else:
+				layers.append(ScaledLeakyReLU(0.2))
+
+		super().__init__(*layers)
+
+
+class ResBlock(nn.Module):
+	def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
+		super().__init__()
+
+		self.conv1 = ConvLayer(in_channel, in_channel, 3)
+		self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)
+
+		self.skip = ConvLayer(in_channel, out_channel, 1, downsample=True, activate=False, bias=False)
+
+	def forward(self, input):
+		out = self.conv1(input)
+		out = self.conv2(out)
+
+		skip = self.skip(input)
+		out = (out + skip) / math.sqrt(2)
+
+		return out
+
+
+class EncoderApp(nn.Module):
+	def __init__(self, size, w_dim=512):
+		super(EncoderApp, self).__init__()
+
+		channels = {
+			4: 512,
+			8: 512,
+			16: 512,
+			32: 512,
+			64: 256,
+			128: 128,
+			256: 64,
+			512: 32,
+			1024: 16
+		}
+
+		self.w_dim = w_dim
+		log_size = int(math.log(size, 2))
+
+		self.convs = nn.ModuleList()
+		self.convs.append(ConvLayer(3, channels[size], 1))
+
+		in_channel = channels[size]
+		for i in range(log_size, 2, -1):
+			out_channel = channels[2 ** (i - 1)]
+			self.convs.append(ResBlock(in_channel, out_channel))
+			in_channel = out_channel
+
+		self.convs.append(EqualConv2d(in_channel, self.w_dim, 4, padding=0, bias=False))
+
+	def forward(self, x):
+
+		res = []
+		h = x
+		for conv in self.convs:
+			h = conv(h)
+			res.append(h)
+
+		return res[-1].squeeze(-1).squeeze(-1), res[::-1][2:]
+
+
+class Encoder(nn.Module):
+	def __init__(self, size, dim=512, dim_motion=20):
+		super(Encoder, self).__init__()
+
+		# appearance netmork
+		self.net_app = EncoderApp(size, dim)
+
+		# motion network
+		fc = [EqualLinear(dim, dim)]
+		for i in range(3):
+			fc.append(EqualLinear(dim, dim))
+
+		fc.append(EqualLinear(dim, dim_motion))
+		self.fc = nn.Sequential(*fc)
+
+	def enc_app(self, x):
+		h_source = self.net_app(x)
+		return h_source
+
+	def enc_motion(self, x):
+		h, _ = self.net_app(x)
+		h_motion = self.fc(h)
+		return h_motion
+
+
+class Direction(nn.Module):
+    def __init__(self, motion_dim):
+        super(Direction, self).__init__()
+        self.weight = nn.Parameter(torch.randn(512, motion_dim))
+
+    def forward(self, input):
+
+        weight = self.weight + 1e-8
+        Q, R = custom_qr(weight)
+        if input is None:
+            return Q
+        else:
+            input_diag = torch.diag_embed(input)  # alpha, diagonal matrix
+            out = torch.matmul(input_diag, Q.T)
+            out = torch.sum(out, dim=1)
+            return out
+
+
+class Synthesis(nn.Module):
+    def __init__(self, motion_dim):
+        super(Synthesis, self).__init__()
+        self.direction = Direction(motion_dim)
+
+
+class Generator(nn.Module):
+    def __init__(self, size, style_dim=512, motion_dim=20):
+        super().__init__()
+
+        self.enc = Encoder(size, style_dim, motion_dim)
+        self.dec = Synthesis(motion_dim)
+
+    def get_motion(self, img):
+        #motion_feat = self.enc.enc_motion(img)
+        motion_feat = torch.utils.checkpoint.checkpoint((self.enc.enc_motion), img, use_reentrant=True)
+        with torch.cuda.amp.autocast(dtype=torch.float32):
+            motion = self.dec.direction(motion_feat)
+        return motion