Upload 69 files

.gitattributes

@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/desi.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/desi.png filter=lfs diff=lfs merge=lfs -text
+examples/man.png filter=lfs diff=lfs merge=lfs -text
+examples/paul.mp4 filter=lfs diff=lfs merge=lfs -text

INSTALL.md

@@ -0,0 +1,55 @@
+# Installation Guide
+
+## Install with pip
+
+```bash
+pip install .
+pip install .[dev]  # Installe aussi les outils de dev
+```
+
+## Install with Poetry
+
+Ensure you have [Poetry](https://python-poetry.org/docs/#installation) installed on your system.
+
+To install all dependencies:
+
+```bash
+poetry install
+```
+
+### Handling `flash-attn` Installation Issues
+
+If `flash-attn` fails due to **PEP 517 build issues**, you can try one of the following fixes.
+
+#### No-Build-Isolation Installation (Recommended)
+```bash
+poetry run pip install --upgrade pip setuptools wheel
+poetry run pip install flash-attn --no-build-isolation
+poetry install
+```
+
+#### Install from Git (Alternative)
+```bash
+poetry run pip install git+https://github.com/Dao-AILab/flash-attention.git
+```
+
+---
+
+### Running the Model
+
+Once the installation is complete, you can run **Wan2.2** using:
+
+```bash
+poetry run python generate.py --task t2v-A14B --size '1280*720' --ckpt_dir ./Wan2.2-T2V-A14B --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage."
+```
+
+#### Test
+```bash
+bash tests/test.sh
+```
+
+#### Format
+```bash
+black .
+isort .
+```

LICENSE.txt

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

Makefile

@@ -0,0 +1,5 @@
+.PHONY: format
+
+format:
+	isort generate.py wan
+	yapf -i -r *.py generate.py wan

README.md

@@ -1,12 +1,12 @@
----
-title: Wan2.2 Animate ZEROGPU
-emoji: 😻
-colorFrom: indigo
-colorTo: purple
-sdk: gradio
-sdk_version: 5.49.1
-app_file: app.py
-pinned: false
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: Wan2.2 Animate [Local]
+emoji: 🔥
+colorFrom: pink
+colorTo: yellow
+sdk: gradio
+sdk_version: 5.49.1
+app_file: app.py
+pinned: false
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,546 @@
+import spaces
+from huggingface_hub import snapshot_download, hf_hub_download
+import os
+import subprocess
+import importlib, site
+from PIL import Image
+import uuid
+import shutil
+import time
+import cv2
+from generate import generate, load_model
+import json
+
+# Re-discover all .pth/.egg-link files
+for sitedir in site.getsitepackages():
+    site.addsitedir(sitedir)
+
+# Clear caches so importlib will pick up new modules
+importlib.invalidate_caches()
+
+def sh(cmd): subprocess.check_call(cmd, shell=True)
+
+try:
+    print("Attempting to download and build sam2...")
+
+    print("download sam")
+    sam_dir = snapshot_download(repo_id="alexnasa/sam2")
+
+    @spaces.GPU(duration=450)
+    def install_sam():
+        os.environ["TORCH_CUDA_ARCH_LIST"] = "9.0"
+        sh(f"cd {sam_dir} && python setup.py build_ext --inplace && pip install -e .")
+
+    print("install sam")
+    install_sam()
+
+    # tell Python to re-scan site-packages now that the egg-link exists
+    import importlib, site; site.addsitedir(site.getsitepackages()[0]); importlib.invalidate_caches()
+
+    flash_attention_installed = True
+    print("sam2 installed successfully.")
+
+except Exception as e:
+    print(f"⚠️ Could not install sam2: {e}")
+    print("Continuing without sam2...")
+
+import torch
+print(f"Torch version: {torch.__version__}")
+
+os.environ["PROCESSED_RESULTS"] = f"{os.getcwd()}/processed_results"
+
+import gradio as gr
+
+
+snapshot_download(repo_id="Wan-AI/Wan2.2-Animate-14B", local_dir="./Wan2.2-Animate-14B")
+wan_animate = load_model(True)
+
+
+rc_mapping = {
+    "Video → Ref Image" : False,
+    "Video ← Ref Image" : True
+}
+
+
+def preprocess_video(input_video_path, session_id=None):
+    
+    if session_id is None:
+        session_id = uuid.uuid4().hex
+
+    output_dir = os.path.join(os.environ["PROCESSED_RESULTS"], session_id)
+    os.makedirs(output_dir, exist_ok=True)
+
+    process_video_path = os.path.join(output_dir, 'input_video.mp4')
+
+    convert_video_to_30fps_and_clip(input_video_path, process_video_path, crop_width=720, crop_height=1280)
+
+    return process_video_path
+
+def extract_audio_from_video_ffmpeg(video_path, output_wav_path, sample_rate=None):
+    """
+    Extracts the audio track from a video file and saves it as a WAV file.
+
+    Args:
+        video_path (str): Path to the input video file.
+        output_wav_path (str): Path to save the extracted WAV file.
+        sample_rate (int, optional): Output sample rate (e.g., 16000). 
+                                     If None, keep the original.
+    """
+    cmd = [
+        'ffmpeg',
+        '-i', video_path,             # Input video
+        '-vn',                        # Disable video
+        '-acodec', 'pcm_s16le',       # 16-bit PCM (WAV format)
+        '-ac', '1',                   # Mono channel (use '2' for stereo)
+        '-y',                         # Overwrite output
+        '-loglevel', 'error'          # Cleaner output
+    ]
+
+    # Only add the sample rate option if explicitly specified
+    if sample_rate is not None:
+        cmd.extend(['-ar', str(sample_rate)])
+
+    cmd.append(output_wav_path)
+    
+    try:
+        subprocess.run(cmd, check=True, capture_output=True, text=True)
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(f"ffmpeg failed ({e.returncode}): {e.stderr.strip()}")
+
+    
+def combine_video_and_audio_ffmpeg(video_path, audio_path, output_video_path):
+    """
+    Combines a silent MP4 video with a WAV audio file into a single MP4 with sound.
+
+    Args:
+        video_path (str): Path to the silent video file.
+        audio_path (str): Path to the WAV audio file.
+        output_video_path (str): Path to save the output MP4 with audio.
+    """
+    cmd = [
+        'ffmpeg',
+        '-i', video_path,            # Input video
+        '-i', audio_path,            # Input audio
+        '-c:v', 'copy',              # Copy video without re-encoding
+        '-c:a', 'aac',               # Encode audio as AAC (MP4-compatible)
+        '-shortest',                 # Stop when the shortest stream ends
+        '-y',                        # Overwrite output
+        '-loglevel', 'error',
+        output_video_path
+    ]
+    
+    try:
+        subprocess.run(cmd, check=True, capture_output=True, text=True)
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(f"ffmpeg failed ({e.returncode}): {e.stderr.strip()}")
+
+
+def convert_video_to_30fps_and_clip(
+    input_video_path,
+    output_video_path,
+    duration_s=2,
+    target_fps=30,
+    crop_width=None,
+    crop_height=None
+):
+    # Get input video dimensions using ffprobe
+    if crop_width and crop_height:
+        probe_cmd = [
+            'ffprobe', '-v', 'error', '-select_streams', 'v:0',
+            '-show_entries', 'stream=width,height',
+            '-of', 'json', input_video_path
+        ]
+        probe_result = subprocess.run(probe_cmd, capture_output=True, text=True, check=True)
+        video_info = json.loads(probe_result.stdout)
+        w = video_info['streams'][0]['width']
+        h = video_info['streams'][0]['height']
+
+        # Clamp crop size to not exceed actual dimensions
+        crop_width = min(crop_width, w)
+        crop_height = min(crop_height, h)
+
+        # Center crop offsets
+        crop_x = max((w - crop_width) // 2, 0)
+        crop_y = max((h - crop_height) // 2, 0)
+        crop_filter = f"crop={crop_width}:{crop_height}:{crop_x}:{crop_y}"
+    else:
+        crop_filter = None
+
+    cmd = [
+        'ffmpeg',
+        '-i', input_video_path,
+        '-r', str(target_fps),
+        '-t', str(duration_s),
+    ]
+
+    if crop_filter:
+        cmd += ['-vf', crop_filter]
+
+    cmd += [
+        '-c:v', 'libx264',
+        '-c:a', 'aac',
+        '-strict', 'experimental',
+        '-y',
+        '-loglevel', 'error',
+        output_video_path
+    ]
+
+    try:
+        subprocess.run(cmd, check=True, capture_output=True, text=True)
+    except subprocess.CalledProcessError as e:
+        raise RuntimeError(f"ffmpeg failed ({e.returncode}): {e.stderr.strip()}")
+
+def get_frames_count(video_file):
+
+    # Get video information
+    cap = cv2.VideoCapture(video_file)
+    if not cap.isOpened():
+        error_msg = "Cannot open video file"
+        gr.Warning(error_msg)
+    
+    orig_frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    orig_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    orig_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    
+    cap.release()
+
+    return orig_frame_count
+
+def calculate_time_required(input_video, rc_bool):
+
+    frames_count = get_frames_count(input_video)
+
+    chunks = frames_count // 77 + 1
+
+
+    if rc_bool:
+        pose2d_tracking_duration_s = 75
+        iteration_per_step_s = 13
+    else:
+        pose2d_tracking_duration_s = 50
+        iteration_per_step_s = 12
+        
+    time_required = pose2d_tracking_duration_s + iteration_per_step_s * 20 * chunks
+    print(f'for frames_count:{frames_count} doing {chunks} chunks the time_required is {time_required}')
+    return time_required
+
+def update_time_required(input_video, rc_str):
+
+    if input_video is None:
+        return gr.update(value="⌚ Zero GPU Required: --")
+    
+    rc_bool = rc_mapping[rc_str]
+
+    duration_s = calculate_time_required(input_video, rc_bool)
+    duration_m = duration_s / 60
+    
+    return gr.update(value=f"⌚ Zero GPU Required: ~{duration_s}.0s ({duration_m:.1f} mins)")
+
+def get_duration(input_video, edited_frame, rc_bool, session_id, progress):
+
+    return calculate_time_required(input_video, rc_bool)
+    
+
+@spaces.GPU(duration=get_duration)
+def _animate(input_video, edited_frame, rc_bool, session_id = None, progress=gr.Progress(track_tqdm=True),):
+
+    if session_id is None:
+        session_id = uuid.uuid4().hex
+
+    output_dir = os.path.join(os.environ["PROCESSED_RESULTS"], session_id)
+    os.makedirs(output_dir, exist_ok=True)
+
+    preprocess_dir = os.path.join(output_dir, "preprocess_dir")
+    os.makedirs(preprocess_dir, exist_ok=True)
+
+    output_video_path = os.path.join(output_dir, 'result.mp4')
+
+    # --- Measure preprocess time ---
+    start_preprocess = time.time()
+
+    # w = 720
+    # h = 480
+
+    # w = 720
+    # h = 1280
+
+    w = 480
+    h = 832
+
+    # w = 480
+    # h = 720
+
+    tag_string = "retarget_flag"
+
+    if rc_bool:
+        tag_string = "replace_flag"
+
+    sh("python ./wan/modules/animate/preprocess/preprocess_data.py "
+       "--ckpt_path ./Wan2.2-Animate-14B/process_checkpoint "
+       f"--video_path {input_video} "
+       f"--refer_path {edited_frame} "
+       f"--save_path {preprocess_dir} "
+       f"--resolution_area {w} {h} --{tag_string} "
+       )
+
+    preprocess_time = time.time() - start_preprocess
+    print(f"Preprocess took {preprocess_time:.2f} seconds")
+
+    # --- Measure generate time ---
+    start_generate = time.time()
+    
+    generate(wan_animate, preprocess_dir, output_video_path, rc_bool)
+
+    generate_time = time.time() - start_generate
+    print(f"Generate took {generate_time:.2f} seconds")
+
+    # --- Optional total time ---
+    total_time = preprocess_time + generate_time
+    print(f"Total time: {total_time:.2f} seconds")
+
+    return output_video_path
+
+def animate_scene(input_video, edited_frame, rc_str, session_id = None, progress=gr.Progress(track_tqdm=True),):
+
+    if not input_video:
+        raise gr.Error("Please provide an video")
+    
+    if not edited_frame:
+        raise gr.Error("Please provide an image")
+    
+    if session_id is None:
+        session_id = uuid.uuid4().hex
+
+    rc_bool = rc_mapping[rc_str]
+    
+    output_dir = os.path.join(os.environ["PROCESSED_RESULTS"], session_id)
+    os.makedirs(output_dir, exist_ok=True)
+
+    input_audio_path = os.path.join(output_dir, 'input_audio.wav')
+    
+    extract_audio_from_video_ffmpeg(input_video, input_audio_path)
+
+    output_video_path = _animate(input_video, edited_frame, rc_bool, session_id, progress)
+
+    final_video_path = os.path.join(output_dir, 'final_result.mp4')
+
+    preprocess_dir = os.path.join(output_dir, "preprocess_dir")
+    pose_video =  os.path.join(preprocess_dir, 'src_pose.mp4') 
+
+    if rc_bool:
+        mask_video = os.path.join(preprocess_dir, 'src_mask.mp4')
+        bg_video = os.path.join(preprocess_dir, 'src_bg.mp4')
+        face_video = os.path.join(preprocess_dir, 'src_face.mp4')
+    else:
+        mask_video = os.path.join(preprocess_dir, 'src_pose.mp4') 
+        bg_video = os.path.join(preprocess_dir, 'src_pose.mp4') 
+        face_video = os.path.join(preprocess_dir, 'src_pose.mp4') 
+
+    combine_video_and_audio_ffmpeg(output_video_path, input_audio_path, final_video_path)
+
+    return final_video_path, pose_video, bg_video, mask_video, face_video
+
+css = """
+    #col-container {
+        margin: 0 auto;
+        max-width: 1600px;
+    }
+
+    #step-column {
+        padding: 20px;
+        border-radius: 8px;
+        box-shadow: var(--card-shadow);
+        margin: 10px;
+    }
+
+    #col-showcase {
+        margin: 0 auto;
+        max-width: 1100px;
+    }
+
+    .button-gradient {
+        background: linear-gradient(45deg, rgb(255, 65, 108), rgb(255, 75, 43), rgb(255, 155, 0), rgb(255, 65, 108)) 0% 0% / 400% 400%;
+        border: none;
+        padding: 14px 28px;
+        font-size: 16px;
+        font-weight: bold;
+        color: white;
+        border-radius: 10px;
+        cursor: pointer;
+        transition: 0.3s ease-in-out;
+        animation: 2s linear 0s infinite normal none running gradientAnimation;
+        box-shadow: rgba(255, 65, 108, 0.6) 0px 4px 10px;
+    }
+
+    .toggle-container {
+    display: inline-flex;
+    background-color: #ffd6ff;  /* light pink background */
+    border-radius: 9999px;
+    padding: 4px;
+    position: relative;
+    width: fit-content;
+    font-family: sans-serif;
+    }
+
+    .toggle-container input[type="radio"] {
+    display: none;
+    }
+
+    .toggle-container label {
+    position: relative;
+    z-index: 2;
+    flex: 1;
+    text-align: center;
+    font-weight: 700;
+    color: #4b2ab5; /* dark purple text for unselected */
+    padding: 6px 22px;
+    border-radius: 9999px;
+    cursor: pointer;
+    transition: color 0.25s ease;
+    }
+
+    /* Moving highlight */
+    .toggle-highlight {
+    position: absolute;
+    top: 4px;
+    left: 4px;
+    width: calc(50% - 4px);
+    height: calc(100% - 8px);
+    background-color: #4b2ab5; /* dark purple background */
+    border-radius: 9999px;
+    transition: transform 0.25s ease;
+    z-index: 1;
+    }
+
+    /* When "True" is checked */
+    #true:checked ~ label[for="true"] {
+    color: #ffd6ff; /* light pink text */
+    }
+
+    /* When "False" is checked */
+    #false:checked ~ label[for="false"] {
+    color: #ffd6ff; /* light pink text */
+    }
+
+    /* Move highlight to right side when False is checked */
+    #false:checked ~ .toggle-highlight {
+    transform: translateX(100%);
+    }
+    """
+def start_session(request: gr.Request):
+
+    return request.session_hash
+
+def cleanup(request: gr.Request):
+
+    sid = request.session_hash
+    
+    if sid:
+        d1 = os.path.join(os.environ["PROCESSED_RESULTS"], sid)
+        shutil.rmtree(d1, ignore_errors=True)
+
+with gr.Blocks(css=css, title="Wan 2.2 Animate --replace", theme=gr.themes.Ocean()) as demo:
+        
+    session_state = gr.State()
+    demo.load(start_session, outputs=[session_state])
+
+    with gr.Column(elem_id="col-container"):
+        with gr.Row():
+            gr.HTML(
+                """
+                <div style="text-align: center;">
+                    <p style="font-size:16px; display: inline; margin: 0;">
+                        <strong>Wan2.2-Animate-14B </strong>
+                    </p>
+                    <a href="https://huggingface.co/Wan-AI/Wan2.2-Animate-14B" style="display: inline-block; vertical-align: middle; margin-left: 0.5em;">
+                        [Model]
+                    </a>
+                <div style="text-align: center;">
+                    <p style="font-size:16px; display: inline; margin: 0;">
+                        HF Space By:
+                    </p>
+                    <a href="https://huggingface.co/alexnasa" style="display: inline-block; vertical-align: middle; margin-left: 0.5em;">
+                        <img src="https://img.shields.io/badge/🤗-Follow Me-yellow.svg">
+                    </a>
+                </div>
+                """
+            )
+        with gr.Row():
+            with gr.Column(elem_id="step-column"):
+                gr.HTML("""
+                <div>
+                    <span style="font-size: 24px;">1. Upload a Video</span><br>
+                </div>
+                """)
+                input_video = gr.Video(label="Input Video", height=512)
+
+
+            with gr.Column(elem_id="step-column"):
+                gr.HTML("""
+                <div>
+                    <span style="font-size: 24px;">2. Upload a Ref Image</span><br>
+                </div>
+                """)
+                edited_frame = gr.Image(label="Ref Image", type="filepath", height=512)
+                gr.HTML("""
+                <div>
+                    <span style="font-size: 24px;">3. Choose Mode</span><br>
+                </div>
+                """)
+                replace_character_string = gr.Radio(
+                    ["Video → Ref Image", "Video ← Ref Image"], value="Video → Ref Image", show_label=False
+                )
+
+            with gr.Column(elem_id="step-column"):
+                gr.HTML("""
+                <div>
+                    <span style="font-size: 24px;">4. Wan Animate it!</span><br>
+                </div>
+                """)
+                output_video = gr.Video(label="Edited Video", height=512)
+
+                time_required = gr.Text(value="⌚ Zero GPU Required: --", show_label=False)
+                action_button = gr.Button("Wan Animate 🦆", variant='primary', elem_classes="button-gradient")
+
+                with gr.Accordion("Preprocessed Data", open=False, visible=False):
+                    pose_video = gr.Video(label="Pose Video", height=512)
+                    bg_video = gr.Video(label="Background Video", height=512)
+                    face_video = gr.Video(label="Face Video", height=512)
+                    mask_video = gr.Video(label="Mask Video", height=512)
+
+        with gr.Row():
+            with gr.Column(elem_id="col-showcase"):
+
+                gr.Examples(
+                    examples=[
+
+                        [
+                            "./examples/desi.mp4", 
+                            "./examples/desi.png",     
+                            "Video ← Ref Image"
+                        ],
+
+                        [
+                            "./examples/paul.mp4", 
+                            "./examples/man.png",     
+                            "Video → Ref Image"
+                        ],
+
+
+                    ],
+                    inputs=[input_video, edited_frame, replace_character_string],
+                    outputs=[output_video, pose_video, bg_video, mask_video, face_video],
+                    fn=animate_scene,
+                    cache_examples=True,
+                )
+
+    action_button.click(fn=animate_scene, inputs=[input_video, edited_frame, replace_character_string, session_state], outputs=[output_video, pose_video, bg_video, mask_video, face_video])
+
+    input_video.upload(preprocess_video, inputs=[input_video, session_state], outputs=[input_video]).then(update_time_required, inputs=[input_video, replace_character_string], outputs=[time_required])
+    replace_character_string.change(update_time_required, inputs=[input_video, replace_character_string], outputs=[time_required])
+                   
+if __name__ == "__main__":
+    demo.queue()
+    demo.unload(cleanup)
+    demo.launch(ssr_mode=False, share=True)
+    

examples/desi.mp4

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

examples/paul.mp4

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

generate.py

@@ -0,0 +1,236 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import argparse
+import logging
+import os
+import sys
+import warnings
+from datetime import datetime
+
+warnings.filterwarnings('ignore')
+
+import random
+
+import torch
+import torch.distributed as dist
+from PIL import Image
+
+import wan
+from wan.configs import MAX_AREA_CONFIGS, SIZE_CONFIGS, SUPPORTED_SIZES, WAN_CONFIGS
+from wan.distributed.util import init_distributed_group
+from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander
+from wan.utils.utils import merge_video_audio, save_video, str2bool
+
+
+EXAMPLE_PROMPT = {
+    "t2v-A14B": {
+        "prompt":
+            "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
+    },
+    "i2v-A14B": {
+        "prompt":
+            "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside.",
+        "image":
+            "examples/i2v_input.JPG",
+    },
+    "ti2v-5B": {
+        "prompt":
+            "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
+    },
+    "animate-14B": {
+        "prompt": "视频中的人在做动作",
+        "video": "",
+        "pose": "",
+        "mask": "",
+    },
+    "s2v-14B": {
+        "prompt":
+            "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside.",
+        "image":
+            "examples/i2v_input.JPG",
+        "audio":
+            "examples/talk.wav",
+        "tts_prompt_audio":
+            "examples/zero_shot_prompt.wav",
+        "tts_prompt_text":
+            "希望你以后能够做的比我还好呦。",
+        "tts_text":
+            "收到好友从远方寄来的生日礼物，那份意外的惊喜与深深的祝福让我心中充满了甜蜜的快乐，笑容如花儿般绽放。"
+    },
+}
+
+
+def _validate_args(args):
+    # Basic check
+    assert args.ckpt_dir is not None, "Please specify the checkpoint directory."
+    assert args.task in WAN_CONFIGS, f"Unsupport task: {args.task}"
+    assert args.task in EXAMPLE_PROMPT, f"Unsupport task: {args.task}"
+
+    if args.prompt is None:
+        args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
+    if args.image is None and "image" in EXAMPLE_PROMPT[args.task]:
+        args.image = EXAMPLE_PROMPT[args.task]["image"]
+    if args.audio is None and args.enable_tts is False and "audio" in EXAMPLE_PROMPT[args.task]:
+        args.audio = EXAMPLE_PROMPT[args.task]["audio"]
+    if (args.tts_prompt_audio is None or args.tts_text is None) and args.enable_tts is True and "audio" in EXAMPLE_PROMPT[args.task]:
+        args.tts_prompt_audio = EXAMPLE_PROMPT[args.task]["tts_prompt_audio"]
+        args.tts_prompt_text = EXAMPLE_PROMPT[args.task]["tts_prompt_text"]
+        args.tts_text = EXAMPLE_PROMPT[args.task]["tts_text"]
+
+    if args.task == "i2v-A14B":
+        assert args.image is not None, "Please specify the image path for i2v."
+
+    cfg = WAN_CONFIGS[args.task]
+
+    if args.sample_steps is None:
+        args.sample_steps = cfg.sample_steps
+
+    if args.sample_shift is None:
+        args.sample_shift = cfg.sample_shift
+
+    if args.sample_guide_scale is None:
+        args.sample_guide_scale = cfg.sample_guide_scale
+
+    if args.frame_num is None:
+        args.frame_num = cfg.frame_num
+
+    args.base_seed = args.base_seed if args.base_seed >= 0 else random.randint(
+        0, sys.maxsize)
+    # Size check
+    if not 's2v' in args.task:
+        assert args.size in SUPPORTED_SIZES[
+            args.
+            task], f"Unsupport size {args.size} for task {args.task}, supported sizes are: {', '.join(SUPPORTED_SIZES[args.task])}"
+
+
+class _Args:
+    pass
+
+def _parse_args():
+    args = _Args()
+
+    # core generation options
+    args.task                    = "animate-14B"
+    # args.size                    = "1280*720"
+    args.size                    = "720*1280"
+    args.frame_num               = None
+    args.ckpt_dir                = "./Wan2.2-Animate-14B/"
+    args.offload_model           = True
+    args.ulysses_size            = 1
+    args.t5_fsdp                 = False
+    args.t5_cpu                  = False
+    args.dit_fsdp                = False
+    args.prompt                  = None
+    args.use_prompt_extend       = False
+    args.prompt_extend_method    = "local_qwen"     # ["dashscope", "local_qwen"]
+    args.prompt_extend_model     = None
+    args.prompt_extend_target_lang = "zh"           # ["zh", "en"]
+    args.base_seed               = 0
+    args.image                   = None
+    args.sample_solver           = "unipc"          # ['unipc', 'dpm++']
+    args.sample_steps            = None
+    args.sample_shift            = None
+    args.sample_guide_scale      = None
+    args.convert_model_dtype     = False
+
+    # animate
+    args.refert_num              = 1
+
+    # s2v-only
+    args.num_clip                = None
+    args.audio                   = None
+    args.enable_tts              = False
+    args.tts_prompt_audio        = None
+    args.tts_prompt_text         = None
+    args.tts_text                = None
+    args.pose_video              = None
+    args.start_from_ref          = False
+    args.infer_frames            = 80
+
+    _validate_args(args)
+    return args
+
+
+
+def _init_logging(rank):
+    # logging
+    if rank == 0:
+        # set format
+        logging.basicConfig(
+            level=logging.INFO,
+            format="[%(asctime)s] %(levelname)s: %(message)s",
+            handlers=[logging.StreamHandler(stream=sys.stdout)])
+    else:
+        logging.basicConfig(level=logging.ERROR)
+
+def load_model(use_relighting_lora = False):
+    
+    cfg = WAN_CONFIGS["animate-14B"]
+    
+    return wan.WanAnimate(
+        config=cfg,
+        checkpoint_dir="./Wan2.2-Animate-14B/",
+        device_id=0,
+        rank=0,
+        t5_fsdp=False,
+        dit_fsdp=False,
+        use_sp=False,
+        t5_cpu=False,
+        convert_model_dtype=False,
+        use_relighting_lora=use_relighting_lora
+    )
+
+def generate(wan_animate, preprocess_dir, save_file, replace_flag = False):
+    args = _parse_args()
+    rank = int(os.getenv("RANK", 0))
+    world_size = int(os.getenv("WORLD_SIZE", 1))
+    local_rank = int(os.getenv("LOCAL_RANK", 0))
+    device = local_rank
+    _init_logging(rank)
+
+    cfg = WAN_CONFIGS[args.task]
+
+    logging.info(f"Input prompt: {args.prompt}")
+    img = None
+    if args.image is not None:
+        img = Image.open(args.image).convert("RGB")
+        logging.info(f"Input image: {args.image}")
+
+    print(f'rank:{rank}')
+
+
+
+    logging.info(f"Generating video ...")
+    video = wan_animate.generate(
+        src_root_path=preprocess_dir,
+        replace_flag=replace_flag,
+        refert_num = args.refert_num,
+        clip_len=args.frame_num,
+        shift=args.sample_shift,
+        sample_solver=args.sample_solver,
+        sampling_steps=args.sample_steps,
+        guide_scale=args.sample_guide_scale,
+        seed=args.base_seed,
+        offload_model=args.offload_model)
+    if rank == 0:
+
+        save_video(
+            tensor=video[None],
+            save_file=save_file,
+            fps=cfg.sample_fps,
+            nrow=1,
+            normalize=True,
+            value_range=(-1, 1))
+        # if "s2v" in args.task:
+        #     if args.enable_tts is False:
+        #         merge_video_audio(video_path=args.save_file, audio_path=args.audio)
+        #     else:
+        #         merge_video_audio(video_path=args.save_file, audio_path="tts.wav")
+    del video
+
+    torch.cuda.synchronize()
+    if dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()
+
+    logging.info("Finished.")
+

pyproject.toml

@@ -0,0 +1,66 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "wan"
+version = "2.2.0"
+description = "Wan: Open and Advanced Large-Scale Video Generative Models"
+authors = [
+    { name = "Wan Team", email = "wan.ai@alibabacloud.com" }
+]
+license = { file = "LICENSE.txt" }
+readme = "README.md"
+requires-python = ">=3.10,<4.0"
+dependencies = [
+    "torch>=2.4.0",
+    "torchvision>=0.19.0",
+    "opencv-python>=4.9.0.80",
+    "diffusers>=0.31.0",
+    "transformers>=4.49.0",
+    "tokenizers>=0.20.3",
+    "accelerate>=1.1.1",
+    "tqdm",
+    "imageio",
+    "easydict",
+    "ftfy",
+    "dashscope",
+    "imageio-ffmpeg",
+    "flash_attn",
+    "numpy>=1.23.5,<2"
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "black",
+    "flake8",
+    "isort",
+    "mypy",
+    "huggingface-hub[cli]"
+]
+
+[project.urls]
+homepage = "https://wanxai.com"
+documentation = "https://github.com/Wan-Video/Wan2.2"
+repository = "https://github.com/Wan-Video/Wan2.2"
+huggingface = "https://huggingface.co/Wan-AI/"
+modelscope = "https://modelscope.cn/organization/Wan-AI"
+discord = "https://discord.gg/p5XbdQV7"
+
+[tool.setuptools]
+packages = ["wan"]
+
+[tool.setuptools.package-data]
+"wan" = ["**/*.py"]
+
+[tool.black]
+line-length = 88
+
+[tool.isort]
+profile = "black"
+
+[tool.mypy]
+strict = true
+
+

requirements.txt

@@ -0,0 +1,31 @@
+torch==2.8.0
+decord
+peft
+pandas
+matplotlib
+loguru
+sentencepiece
+dashscope
+ftfy
+diffusers
+opencv-python
+moviepy
+torchvision
+torchaudio
+transformers
+tokenizers
+accelerate
+tqdm
+imageio[ffmpeg]
+easydict
+imageio-ffmpeg
+numpy>=1.23.5,<2
+hydra-core
+iopath
+pytest
+pillow
+fvcore
+librosa
+flash-attn
+onnxruntime-gpu
+flash-attn-3 @ https://huggingface.co/alexnasa/flash-attn-3/resolve/main/128/flash_attn_3-3.0.0b1-cp39-abi3-linux_x86_64.whl

wan/__init__.py

@@ -0,0 +1,7 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from . import configs, distributed, modules
+from .image2video import WanI2V
+from .speech2video import WanS2V
+from .text2video import WanT2V
+from .textimage2video import WanTI2V
+from .animate import WanAnimate

wan/animate.py

@@ -0,0 +1,653 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import logging
+import math
+import os
+import cv2
+import types
+from copy import deepcopy
+from functools import partial
+from einops import rearrange
+import numpy as np
+import torch
+
+import torch.distributed as dist
+from peft import set_peft_model_state_dict
+from decord import VideoReader
+from tqdm import tqdm
+import torch.nn.functional as F
+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.animate import WanAnimateModel
+from .modules.animate import CLIPModel
+from .modules.t5 import T5EncoderModel
+from .modules.vae2_1 import Wan2_1_VAE
+from .modules.animate.animate_utils import TensorList, get_loraconfig
+from .utils.fm_solvers import (
+    FlowDPMSolverMultistepScheduler,
+    get_sampling_sigmas,
+    retrieve_timesteps,
+)
+from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+
+
+
+class WanAnimate:
+
+    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,
+        use_relighting_lora=False
+    ):
+        r"""
+        Initializes the 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.
+            use_relighting_lora (`bool`, *optional*, defaults to False):
+               Whether to use relighting lora for character replacement. 
+        """
+        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.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.clip = CLIPModel(
+            dtype=torch.float16,
+            device=self.device,
+            checkpoint_path=os.path.join(checkpoint_dir,
+                                         config.clip_checkpoint),
+            tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer))
+
+        self.vae = Wan2_1_VAE(
+            vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
+            device=self.device)
+
+        logging.info(f"Creating WanAnimate from {checkpoint_dir}")
+
+        if not dit_fsdp:
+            self.noise_model = WanAnimateModel.from_pretrained(
+                checkpoint_dir,
+                torch_dtype=self.param_dtype,
+                device_map=self.device)
+        else:
+            self.noise_model = WanAnimateModel.from_pretrained(
+                checkpoint_dir, torch_dtype=self.param_dtype)
+
+        self.noise_model = self._configure_model(
+            model=self.noise_model,
+            use_sp=use_sp,
+            dit_fsdp=dit_fsdp,
+            shard_fn=shard_fn,
+            convert_model_dtype=convert_model_dtype,
+            use_lora=use_relighting_lora,
+            checkpoint_dir=checkpoint_dir,
+            config=config
+            )
+
+        if use_sp:
+            self.sp_size = get_world_size()
+        else:
+            self.sp_size = 1
+
+        self.sample_neg_prompt = config.sample_neg_prompt
+        self.sample_prompt = config.prompt
+
+
+    def _configure_model(self, model, use_sp, dit_fsdp, shard_fn,
+                         convert_model_dtype, use_lora, checkpoint_dir, config):
+        """
+        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.use_context_parallel = True
+
+        if dist.is_initialized():
+            dist.barrier()
+
+        if use_lora:
+            logging.info("Loading Relighting Lora. ")
+            lora_config = get_loraconfig(
+                transformer=model,
+                rank=128,
+                alpha=128
+            )
+            model.add_adapter(lora_config)
+            lora_path = os.path.join(checkpoint_dir, config.lora_checkpoint)
+            peft_state_dict = torch.load(lora_path)["state_dict"]
+            set_peft_model_state_dict(model, peft_state_dict)
+
+        if dit_fsdp:
+            model = shard_fn(model, use_lora=use_lora)
+        else:
+            if convert_model_dtype:
+                model.to(self.param_dtype)
+            if not self.init_on_cpu:
+                model.to(self.device)
+
+        return model
+
+    def inputs_padding(self, array, target_len):
+        idx = 0
+        flip = False
+        target_array = []
+        while len(target_array) < target_len:
+            target_array.append(deepcopy(array[idx]))
+            if flip:
+                idx -= 1
+            else:
+                idx += 1
+            if idx == 0 or idx == len(array) - 1:
+                flip = not flip
+        return target_array[:target_len]
+
+    def get_valid_len(self, real_len, clip_len=81, overlap=1):
+        real_clip_len = clip_len - overlap
+        last_clip_num = (real_len - overlap) % real_clip_len
+        if last_clip_num == 0:
+            extra = 0
+        else:
+            extra = real_clip_len - last_clip_num
+        target_len = real_len + extra
+        return target_len
+
+
+    def get_i2v_mask(self, lat_t, lat_h, lat_w, mask_len=1, mask_pixel_values=None, device="cuda"):
+        if mask_pixel_values is None:
+            msk = torch.zeros(1, (lat_t-1) * 4 + 1, lat_h, lat_w, device=device)
+        else:
+            msk = mask_pixel_values.clone()
+        msk[:, :mask_len] = 1
+        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]
+        return msk
+
+    def padding_resize(self, img_ori, height=512, width=512, padding_color=(0, 0, 0), interpolation=cv2.INTER_LINEAR):
+        ori_height = img_ori.shape[0]
+        ori_width = img_ori.shape[1]
+        channel = img_ori.shape[2]
+
+        img_pad = np.zeros((height, width, channel))
+        if channel == 1:
+            img_pad[:, :, 0] = padding_color[0]
+        else:
+            img_pad[:, :, 0] = padding_color[0]
+            img_pad[:, :, 1] = padding_color[1]
+            img_pad[:, :, 2] = padding_color[2]
+
+        if (ori_height / ori_width) > (height / width):
+            new_width = int(height / ori_height * ori_width)
+            img = cv2.resize(img_ori, (new_width, height), interpolation=interpolation)
+            padding = int((width - new_width) / 2)
+            if len(img.shape) == 2:
+                img = img[:, :, np.newaxis]  
+            img_pad[:, padding: padding + new_width, :] = img
+        else:
+            new_height = int(width / ori_width * ori_height)
+            img = cv2.resize(img_ori, (width, new_height), interpolation=interpolation)
+            padding = int((height - new_height) / 2)
+            if len(img.shape) == 2:
+                img = img[:, :, np.newaxis]  
+            img_pad[padding: padding + new_height, :, :] = img
+
+        img_pad = np.uint8(img_pad)
+
+        return img_pad
+
+    def prepare_source(self, src_pose_path, src_face_path, src_ref_path):
+        pose_video_reader = VideoReader(src_pose_path)
+        pose_len = len(pose_video_reader)
+        pose_idxs = list(range(pose_len))
+        cond_images = pose_video_reader.get_batch(pose_idxs).asnumpy()
+
+        face_video_reader = VideoReader(src_face_path)
+        face_len = len(face_video_reader)
+        face_idxs = list(range(face_len))
+        face_images = face_video_reader.get_batch(face_idxs).asnumpy()
+        height, width = cond_images[0].shape[:2]
+        refer_images = cv2.imread(src_ref_path)[..., ::-1]
+        refer_images = self.padding_resize(refer_images, height=height, width=width)
+        return cond_images, face_images, refer_images
+    
+    def prepare_source_for_replace(self, src_bg_path, src_mask_path):
+        bg_video_reader = VideoReader(src_bg_path)
+        bg_len = len(bg_video_reader)
+        bg_idxs = list(range(bg_len))
+        bg_images = bg_video_reader.get_batch(bg_idxs).asnumpy()
+
+        mask_video_reader = VideoReader(src_mask_path)
+        mask_len = len(mask_video_reader)
+        mask_idxs = list(range(mask_len))
+        mask_images = mask_video_reader.get_batch(mask_idxs).asnumpy()
+        mask_images = mask_images[:, :, :, 0] / 255
+        return bg_images, mask_images
+
+    def generate(
+        self,
+        src_root_path,
+        replace_flag=False,
+        clip_len=77,
+        refert_num=1,
+        shift=5.0,
+        sample_solver='dpm++',
+        sampling_steps=20,
+        guide_scale=1,
+        input_prompt="",
+        n_prompt="",
+        seed=-1,
+        offload_model=True,
+    ):
+        r"""
+        Generates video frames from input image using diffusion process.
+
+        Args:
+            src_root_path ('str'):
+                Process output path
+            replace_flag (`bool`, *optional*, defaults to False):
+                Whether to use character replace.
+            clip_len (`int`, *optional*, defaults to 77):
+                How many frames to generate per clips. The number should be 4n+1
+            refert_num (`int`, *optional*, defaults to 1):
+                How many frames used for temporal guidance. Recommended to be 1 or 5.
+            shift (`float`, *optional*, defaults to 5.0):
+                Noise schedule shift parameter. 
+            sample_solver (`str`, *optional*, defaults to 'dpm++'):
+                Solver used to sample the video.
+            sampling_steps (`int`, *optional*, defaults to 20):
+                Number of diffusion sampling steps. Higher values improve quality but slow generation
+            guide_scale (`float` or tuple[`float`], *optional*, defaults 1.0):
+                Classifier-free guidance scale. We only use it for expression control. 
+                In most cases, it's not necessary and faster generation can be achieved without it. 
+                When expression adjustments are needed, you may consider using this feature.
+            input_prompt (`str`):
+                Text prompt for content generation. We don't recommend custom prompts (although they work)
+            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
+                - H: Frame height 
+                - W: Frame width 
+        """
+        assert refert_num == 1 or refert_num == 5, "refert_num should be 1 or 5."
+
+        seed_g = torch.Generator(device=self.device)
+        seed_g.manual_seed(seed)
+
+
+        if n_prompt == "":
+            n_prompt = self.sample_neg_prompt
+
+        if input_prompt == "":
+            input_prompt = self.sample_prompt
+
+        src_pose_path = os.path.join(src_root_path, "src_pose.mp4")
+        src_face_path = os.path.join(src_root_path, "src_face.mp4")
+        src_ref_path = os.path.join(src_root_path, "src_ref.png")
+
+        cond_images, face_images, refer_images = self.prepare_source(src_pose_path=src_pose_path, src_face_path=src_face_path, src_ref_path=src_ref_path)
+        
+        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]
+
+        real_frame_len = len(cond_images)
+        target_len = self.get_valid_len(real_frame_len, clip_len, overlap=refert_num)
+        logging.info('real frames: {} target frames: {}'.format(real_frame_len, target_len))
+        cond_images = self.inputs_padding(cond_images, target_len)
+        face_images = self.inputs_padding(face_images, target_len)
+        
+        if replace_flag:
+            src_bg_path = os.path.join(src_root_path, "src_bg.mp4")
+            src_mask_path = os.path.join(src_root_path, "src_mask.mp4")
+            bg_images, mask_images = self.prepare_source_for_replace(src_bg_path, src_mask_path)
+            bg_images = self.inputs_padding(bg_images, target_len)
+            mask_images = self.inputs_padding(mask_images, target_len)
+            self.noise_model.disable_adapters() 
+        else:
+            self.noise_model.disable_adapters() 
+
+
+        height, width = refer_images.shape[:2]
+        start = 0
+        end = clip_len
+        all_out_frames = []
+        while True:
+            if start + refert_num >= len(cond_images):
+                break
+
+            if start == 0:
+                mask_reft_len = 0
+            else:
+                mask_reft_len = refert_num
+
+            batch = {
+                        "conditioning_pixel_values": torch.zeros(1, 3, clip_len, height, width),
+                        "bg_pixel_values": torch.zeros(1, 3, clip_len, height, width),
+                        "mask_pixel_values": torch.zeros(1, 1, clip_len, height, width),
+                        "face_pixel_values": torch.zeros(1, 3, clip_len, 512, 512),
+                        "refer_pixel_values": torch.zeros(1, 3, height, width),
+                        "refer_t_pixel_values": torch.zeros(refert_num, 3, height, width)
+                    }   
+
+            batch["conditioning_pixel_values"] = rearrange(
+                torch.tensor(np.stack(cond_images[start:end]) / 127.5 - 1),
+                "t h w c -> 1 c t h w",
+            )
+            batch["face_pixel_values"] = rearrange(
+                torch.tensor(np.stack(face_images[start:end]) / 127.5 - 1),
+                "t h w c -> 1 c t h w",
+            )
+
+            batch["refer_pixel_values"] = rearrange(
+                torch.tensor(refer_images / 127.5 - 1), "h w c -> 1 c h w"
+            )
+
+            if start > 0:
+                batch["refer_t_pixel_values"] = rearrange(
+                    out_frames[0, :, -refert_num:].clone().detach(),
+                    "c t h w -> t c h w",
+                )
+
+            batch["refer_t_pixel_values"] = rearrange(batch["refer_t_pixel_values"],
+                                            "t c h w -> 1 c t h w",
+                                            )
+
+            if replace_flag:
+                batch["bg_pixel_values"] = rearrange(
+                    torch.tensor(np.stack(bg_images[start:end]) / 127.5 - 1),
+                    "t h w c -> 1 c t h w",
+                )
+
+                batch["mask_pixel_values"] = rearrange(
+                    torch.tensor(np.stack(mask_images[start:end])[:, :, :, None]),
+                    "t h w c -> 1 t c h w",
+                )
+                
+
+            for key, value in batch.items():
+                if isinstance(value, torch.Tensor):
+                    batch[key] = value.to(device=self.device, dtype=torch.bfloat16)
+
+            ref_pixel_values = batch["refer_pixel_values"]
+            refer_t_pixel_values = batch["refer_t_pixel_values"]
+            conditioning_pixel_values = batch["conditioning_pixel_values"]
+            face_pixel_values = batch["face_pixel_values"]
+
+            B, _, H, W = ref_pixel_values.shape
+            T = clip_len
+            lat_h = H // 8
+            lat_w = W // 8
+            lat_t = T // 4 + 1
+            target_shape = [lat_t + 1, lat_h, lat_w]
+            noise = [
+                torch.randn(
+                    16,
+                    target_shape[0],
+                    target_shape[1],
+                    target_shape[2],
+                    dtype=torch.float32,
+                    device=self.device,
+                    generator=seed_g,
+                )
+            ]
+        
+            max_seq_len = int(math.ceil(np.prod(target_shape) // 4 / self.sp_size)) * self.sp_size
+            if max_seq_len % self.sp_size != 0:
+                raise ValueError(f"max_seq_len {max_seq_len} is not divisible by sp_size {self.sp_size}")
+
+            with (
+                torch.autocast(device_type=str(self.device), dtype=torch.bfloat16, enabled=True),
+                torch.no_grad()
+            ):
+                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.")
+
+                latents = noise
+
+                pose_latents_no_ref =  self.vae.encode(conditioning_pixel_values.to(torch.bfloat16))
+                pose_latents_no_ref = torch.stack(pose_latents_no_ref)
+                pose_latents = torch.cat([pose_latents_no_ref], dim=2)
+
+                ref_pixel_values = rearrange(ref_pixel_values, "t c h w -> 1 c t h w")
+                ref_latents =  self.vae.encode(ref_pixel_values.to(torch.bfloat16))
+                ref_latents = torch.stack(ref_latents)
+
+                mask_ref = self.get_i2v_mask(1, lat_h, lat_w, 1, device=self.device)
+                y_ref = torch.concat([mask_ref, ref_latents[0]]).to(dtype=torch.bfloat16, device=self.device)
+
+                img = ref_pixel_values[0, :, 0]
+                clip_context = self.clip.visual([img[:, None, :, :]]).to(dtype=torch.bfloat16, device=self.device)
+
+                if mask_reft_len > 0:
+                    if replace_flag:
+                        bg_pixel_values = batch["bg_pixel_values"]
+                        y_reft = self.vae.encode(
+                            [
+                                torch.concat([refer_t_pixel_values[0, :, :mask_reft_len], bg_pixel_values[0, :, mask_reft_len:]], dim=1).to(self.device)
+                            ]
+                        )[0]
+                        mask_pixel_values = 1 - batch["mask_pixel_values"]
+                        mask_pixel_values = rearrange(mask_pixel_values, "b t c h w -> (b t) c h w")
+                        mask_pixel_values = F.interpolate(mask_pixel_values, size=(H//8, W//8), mode='nearest')
+                        mask_pixel_values = rearrange(mask_pixel_values, "(b t) c h w -> b t c h w", b=1)[:,:,0]
+                        msk_reft = self.get_i2v_mask(lat_t, lat_h, lat_w, mask_reft_len, mask_pixel_values=mask_pixel_values, device=self.device)
+                    else:
+                        y_reft = self.vae.encode(
+                            [
+                                torch.concat(
+                                    [
+                                        torch.nn.functional.interpolate(refer_t_pixel_values[0, :, :mask_reft_len].cpu(),
+                                                                        size=(H, W), mode="bicubic"),
+                                        torch.zeros(3, T - mask_reft_len, H, W),
+                                    ],
+                                    dim=1,
+                                ).to(self.device)
+                            ]
+                        )[0]
+                        msk_reft = self.get_i2v_mask(lat_t, lat_h, lat_w, mask_reft_len, device=self.device)
+                else:
+                    if replace_flag:
+                        bg_pixel_values = batch["bg_pixel_values"]
+                        mask_pixel_values = 1 - batch["mask_pixel_values"]
+                        mask_pixel_values = rearrange(mask_pixel_values, "b t c h w -> (b t) c h w")
+                        mask_pixel_values = F.interpolate(mask_pixel_values, size=(H//8, W//8), mode='nearest')
+                        mask_pixel_values = rearrange(mask_pixel_values, "(b t) c h w -> b t c h w", b=1)[:,:,0]
+                        y_reft = self.vae.encode(
+                            [
+                                torch.concat(
+                                    [
+                                        bg_pixel_values[0],
+                                    ],
+                                    dim=1,
+                                ).to(self.device)
+                            ]
+                        )[0]
+                        msk_reft = self.get_i2v_mask(lat_t, lat_h, lat_w, mask_reft_len, mask_pixel_values=mask_pixel_values, device=self.device)
+                    else:
+                        y_reft = self.vae.encode(
+                            [
+                                torch.concat(
+                                    [
+                                        torch.zeros(3, T - mask_reft_len, H, W),
+                                    ],
+                                    dim=1,
+                                ).to(self.device)
+                            ]
+                        )[0]
+                        msk_reft = self.get_i2v_mask(lat_t, lat_h, lat_w, mask_reft_len, device=self.device)
+
+                y_reft = torch.concat([msk_reft, y_reft]).to(dtype=torch.bfloat16, device=self.device)
+                y = torch.concat([y_ref, y_reft], dim=1)
+
+                arg_c = {
+                    "context": context, 
+                    "seq_len": max_seq_len,
+                    "clip_fea": clip_context.to(dtype=torch.bfloat16, device=self.device),
+                    "y": [y],
+                    "pose_latents": pose_latents,
+                    "face_pixel_values": face_pixel_values,
+                }
+
+                if guide_scale > 1:
+                    face_pixel_values_uncond = face_pixel_values * 0 - 1
+                    arg_null = {
+                        "context": context_null,
+                        "seq_len": max_seq_len,
+                        "clip_fea": clip_context.to(dtype=torch.bfloat16, device=self.device),
+                        "y": [y],
+                        "pose_latents": pose_latents,
+                        "face_pixel_values": face_pixel_values_uncond,
+                    }
+
+                for i, t in enumerate(tqdm(timesteps)):
+                    latent_model_input = latents
+                    timestep = [t]
+
+                    timestep = torch.stack(timestep)
+
+                    noise_pred_cond = TensorList(
+                         self.noise_model(TensorList(latent_model_input), t=timestep, **arg_c)
+                    )
+
+                    if guide_scale > 1:
+                        noise_pred_uncond = TensorList(
+                             self.noise_model(
+                                TensorList(latent_model_input), t=timestep, **arg_null
+                            )
+                        )
+                        noise_pred = noise_pred_uncond + guide_scale * (
+                            noise_pred_cond - noise_pred_uncond
+                        )
+                    else:
+                        noise_pred = noise_pred_cond
+
+                    temp_x0 = sample_scheduler.step(
+                        noise_pred[0].unsqueeze(0),
+                        t,
+                        latents[0].unsqueeze(0),
+                        return_dict=False,
+                        generator=seed_g,
+                    )[0]
+                    latents[0] = temp_x0.squeeze(0)
+
+                    x0 = latents
+
+                x0 = [x.to(dtype=torch.float32) for x in x0]
+                out_frames = torch.stack(self.vae.decode([x0[0][:, 1:]]))
+                
+                if start != 0:
+                    out_frames = out_frames[:, :, refert_num:]
+
+                all_out_frames.append(out_frames.cpu())
+
+                start += clip_len - refert_num
+                end += clip_len - refert_num
+
+        videos = torch.cat(all_out_frames, dim=2)[:, :, :real_frame_len]
+        return videos[0] if self.rank == 0 else None

wan/configs/__init__.py

@@ -0,0 +1,50 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import copy
+import os
+
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+from .wan_i2v_A14B import i2v_A14B
+from .wan_s2v_14B import s2v_14B
+from .wan_t2v_A14B import t2v_A14B
+from .wan_ti2v_5B import ti2v_5B
+from .wan_animate_14B import animate_14B
+
+WAN_CONFIGS = {
+    't2v-A14B': t2v_A14B,
+    'i2v-A14B': i2v_A14B,
+    'ti2v-5B': ti2v_5B,
+    'animate-14B': animate_14B,
+    's2v-14B': s2v_14B,
+}
+
+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 = {
+    't2v-A14B': ('720*1280', '1280*720', '480*832', '832*480'),
+    'i2v-A14B': ('720*1280', '1280*720', '480*832', '832*480'),
+    'ti2v-5B': ('704*1280', '1280*704'),
+    's2v-14B': ('720*1280', '1280*720', '480*832', '832*480', '1024*704',
+                '704*1024', '704*1280', '1280*704'),
+    'animate-14B': ('720*1280', '1280*720')
+}

wan/configs/shared_config.py

@@ -0,0 +1,20 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+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_animate_14B.py

@@ -0,0 +1,40 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from easydict import EasyDict
+
+from .shared_config import wan_shared_cfg
+
+#------------------------ Wan animate 14B ------------------------#
+animate_14B = EasyDict(__name__='Config: Wan animate 14B')
+animate_14B.update(wan_shared_cfg)
+
+animate_14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+animate_14B.t5_tokenizer = 'google/umt5-xxl'
+
+animate_14B.clip_checkpoint = 'models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth'
+animate_14B.clip_tokenizer = 'xlm-roberta-large'
+animate_14B.lora_checkpoint = 'relighting_lora.ckpt'
+# vae
+animate_14B.vae_checkpoint = 'Wan2.1_VAE.pth'
+animate_14B.vae_stride = (4, 8, 8)
+
+# transformer
+animate_14B.patch_size = (1, 2, 2)
+animate_14B.dim = 5120
+animate_14B.ffn_dim = 13824
+animate_14B.freq_dim = 256
+animate_14B.num_heads = 40
+animate_14B.num_layers = 40
+animate_14B.window_size = (-1, -1)
+animate_14B.qk_norm = True
+animate_14B.cross_attn_norm = True
+animate_14B.eps = 1e-6
+animate_14B.use_face_encoder = True
+animate_14B.motion_encoder_dim = 512
+
+# inference
+animate_14B.sample_shift = 5.0
+animate_14B.sample_steps = 20
+animate_14B.sample_guide_scale = 1.0
+animate_14B.frame_num = 77
+animate_14B.sample_fps = 30
+animate_14B.prompt = '视频中的人在做动作'

wan/configs/wan_i2v_A14B.py

@@ -0,0 +1,37 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+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 = 5.0
+i2v_A14B.sample_steps = 40
+i2v_A14B.boundary = 0.900
+i2v_A14B.sample_guide_scale = (3.5, 3.5)  # low noise, high noise

wan/configs/wan_s2v_14B.py

@@ -0,0 +1,59 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from easydict import EasyDict
+
+from .shared_config import wan_shared_cfg
+
+#------------------------ Wan S2V 14B ------------------------#
+
+s2v_14B = EasyDict(__name__='Config: Wan S2V 14B')
+s2v_14B.update(wan_shared_cfg)
+
+# t5
+s2v_14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+s2v_14B.t5_tokenizer = 'google/umt5-xxl'
+
+# vae
+s2v_14B.vae_checkpoint = 'Wan2.1_VAE.pth'
+s2v_14B.vae_stride = (4, 8, 8)
+
+# wav2vec
+s2v_14B.wav2vec = "wav2vec2-large-xlsr-53-english"
+
+s2v_14B.num_heads = 40
+# transformer
+s2v_14B.transformer = EasyDict(
+    __name__="Config: Transformer config for WanModel_S2V")
+s2v_14B.transformer.patch_size = (1, 2, 2)
+s2v_14B.transformer.dim = 5120
+s2v_14B.transformer.ffn_dim = 13824
+s2v_14B.transformer.freq_dim = 256
+s2v_14B.transformer.num_heads = 40
+s2v_14B.transformer.num_layers = 40
+s2v_14B.transformer.window_size = (-1, -1)
+s2v_14B.transformer.qk_norm = True
+s2v_14B.transformer.cross_attn_norm = True
+s2v_14B.transformer.eps = 1e-6
+s2v_14B.transformer.enable_adain = True
+s2v_14B.transformer.adain_mode = "attn_norm"
+s2v_14B.transformer.audio_inject_layers = [
+    0, 4, 8, 12, 16, 20, 24, 27, 30, 33, 36, 39
+]
+s2v_14B.transformer.zero_init = True
+s2v_14B.transformer.zero_timestep = True
+s2v_14B.transformer.enable_motioner = False
+s2v_14B.transformer.add_last_motion = True
+s2v_14B.transformer.trainable_token = False
+s2v_14B.transformer.enable_tsm = False
+s2v_14B.transformer.enable_framepack = True
+s2v_14B.transformer.framepack_drop_mode = 'padd'
+s2v_14B.transformer.audio_dim = 1024
+
+s2v_14B.transformer.motion_frames = 73
+s2v_14B.transformer.cond_dim = 16
+
+# inference
+s2v_14B.sample_neg_prompt = "画面模糊，最差质量，画面模糊，细节模糊不清，情绪激动剧烈，手快速抖动，字幕，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
+s2v_14B.drop_first_motion = True
+s2v_14B.sample_shift = 3
+s2v_14B.sample_steps = 40
+s2v_14B.sample_guide_scale = 4.5

wan/configs/wan_t2v_A14B.py

@@ -0,0 +1,37 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from easydict import EasyDict
+
+from .shared_config import wan_shared_cfg
+
+#------------------------ Wan T2V A14B ------------------------#
+
+t2v_A14B = EasyDict(__name__='Config: Wan T2V A14B')
+t2v_A14B.update(wan_shared_cfg)
+
+# t5
+t2v_A14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+t2v_A14B.t5_tokenizer = 'google/umt5-xxl'
+
+# vae
+t2v_A14B.vae_checkpoint = 'Wan2.1_VAE.pth'
+t2v_A14B.vae_stride = (4, 8, 8)
+
+# transformer
+t2v_A14B.patch_size = (1, 2, 2)
+t2v_A14B.dim = 5120
+t2v_A14B.ffn_dim = 13824
+t2v_A14B.freq_dim = 256
+t2v_A14B.num_heads = 40
+t2v_A14B.num_layers = 40
+t2v_A14B.window_size = (-1, -1)
+t2v_A14B.qk_norm = True
+t2v_A14B.cross_attn_norm = True
+t2v_A14B.eps = 1e-6
+t2v_A14B.low_noise_checkpoint = 'low_noise_model'
+t2v_A14B.high_noise_checkpoint = 'high_noise_model'
+
+# inference
+t2v_A14B.sample_shift = 12.0
+t2v_A14B.sample_steps = 40
+t2v_A14B.boundary = 0.875
+t2v_A14B.sample_guide_scale = (3.0, 4.0)  # low noise, high noise

wan/configs/wan_ti2v_5B.py

@@ -0,0 +1,36 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from easydict import EasyDict
+
+from .shared_config import wan_shared_cfg
+
+#------------------------ Wan TI2V 5B ------------------------#
+
+ti2v_5B = EasyDict(__name__='Config: Wan TI2V 5B')
+ti2v_5B.update(wan_shared_cfg)
+
+# t5
+ti2v_5B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+ti2v_5B.t5_tokenizer = 'google/umt5-xxl'
+
+# vae
+ti2v_5B.vae_checkpoint = 'Wan2.2_VAE.pth'
+ti2v_5B.vae_stride = (4, 16, 16)
+
+# transformer
+ti2v_5B.patch_size = (1, 2, 2)
+ti2v_5B.dim = 3072
+ti2v_5B.ffn_dim = 14336
+ti2v_5B.freq_dim = 256
+ti2v_5B.num_heads = 24
+ti2v_5B.num_layers = 30
+ti2v_5B.window_size = (-1, -1)
+ti2v_5B.qk_norm = True
+ti2v_5B.cross_attn_norm = True
+ti2v_5B.eps = 1e-6
+
+# inference
+ti2v_5B.sample_fps = 24
+ti2v_5B.sample_shift = 5.0
+ti2v_5B.sample_steps = 50
+ti2v_5B.sample_guide_scale = 5.0
+ti2v_5B.frame_num = 121

wan/distributed/__init__.py

@@ -0,0 +1 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.

wan/distributed/fsdp.py

@@ -0,0 +1,45 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+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,176 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+import torch.cuda.amp as amp
+
+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,
+):
+    """
+    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
+        ]))
+
+    # 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)
+
+    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,47 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+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,51 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+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,431 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+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
+
+
+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}")
+        self.low_noise_model = WanModel.from_pretrained(
+            checkpoint_dir, subfolder=config.low_noise_checkpoint)
+        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)
+        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,
+                 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)
+        """
+        # 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]
+
+        max_seq_len = ((F - 1) // self.vae_stride[0] + 1) * 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]
+
+        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],
+            }
+
+            arg_null = {
+                'context': context_null,
+                'seq_len': max_seq_len,
+                'y': [y],
+            }
+
+            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,19 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+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 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from .model_animate import WanAnimateModel
+from .clip import CLIPModel
+__all__ = ['WanAnimateModel', 'CLIPModel']

wan/modules/animate/animate_utils.py

@@ -0,0 +1,143 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+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''
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+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 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+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 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+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``
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+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

wan/modules/animate/preprocess/UserGuider.md

@@ -0,0 +1,70 @@
+# Wan-animate Preprocessing User Guider
+
+## 1. Introductions
+
+
+Wan-animate offers two generation modes: `animation` and `replacement`. While both modes extract the skeleton from the reference video, they each have a distinct preprocessing pipeline.
+
+### 1.1 Animation Mode
+
+In this mode, it is highly recommended to enable pose retargeting, especially if the body proportions of the reference and driving characters are dissimilar.
+
+ - A simplified version of pose retargeting pipeline is provided to help developers quickly implement this functionality.
+
+ - **NOTE:** Due to the potential complexity of input data, the results from this simplified retargeting version are NOT guaranteed to be perfect. It is strongly advised to verify the preprocessing results before proceeding.
+
+ - Community contributions to improve on this feature are welcome.
+
+### 1.2 Replacement Mode
+
+ - Pose retargeting is DISABLED by default in this mode. This is a deliberate choice to account for potential spatial interactions between the character and the environment.
+
+ - **WARNING**: If there is a significant mismatch in body proportions between the reference and driving characters, artifacts or deformations may appear in the final output.
+
+ - A simplified version for extracting the character's mask is also provided.
+ - **WARNING:** This mask extraction process is designed for **single-person videos ONLY** and may produce incorrect results or fail in multi-person videos (incorrect pose tracking). For multi-person video, users are required to either develop their own solution or integrate a suitable open-source tool.
+
+---
+
+## 2. Preprocessing Instructions and Recommendations
+
+### 2.1 Basic Usage
+
+- The preprocessing process requires some additional models, including pose detection (mandatory), and mask extraction and image editing models (optional, as needed). Place them according to the following directory structure:
+```
+    /path/to/your/ckpt_path/
+    ├── det/
+    │ └── yolov10m.onnx
+    ├── pose2d/
+    │ └── vitpose_h_wholebody.onnx
+    ├── sam2/
+    │ └── sam2_hiera_large.pt
+    └── FLUX.1-Kontext-dev/
+```
+- `video_path`, `refer_path`, and `save_path` correspond to the paths for the input driving video, the character image, and the preprocessed results.
+
+- When using `animation` mode, two videos, `src_face.mp4` and `src_pose.mp4`, will be generated in `save_path`. When using `replacement` mode, two additional videos, `src_bg.mp4` and `src_mask.mp4`, will also be generated.
+
+- The `resolution_area` parameter determines the resolution for both preprocessing and the generation model. Its size is determined by pixel area.
+
+- The `fps` parameter can specify the frame rate for video processing. A lower frame rate can improve generation efficiency, but may cause stuttering or choppiness.
+
+---
+
+### 2.2 Animation Mode
+
+- We support three forms: not using pose retargeting, using basic pose retargeting, and using enhanced pose retargeting based on the `FLUX.1-Kontext-dev` image editing model. These are specified via the `retarget_flag` and `use_flux` parameters.
+
+- Specifying `retarget_flag` to use basic pose retargeting requires ensuring that both the reference character and the character in the first frame of the driving video are in a front-facing, stretched pose.
+
+- Other than that, we recommend using enhanced pose retargeting by specifying both `retarget_flag` and `use_flux`. **NOTE:** Due to the limited capabilities of `FLUX.1-Kontext-dev`, it is NOT guaranteed to produce the expected results (e.g., consistency is not maintained, the pose is incorrect, etc.). It is recommended to check the intermediate results as well as the finally generated pose video; both are stored in `save_path`. Of course, users can also use a better image editing model, or explore the prompts for Flux on their own.
+
+---
+
+### 2.3 Replacement Mode
+
+- Specifying `replace_flag` to enable data preprocessing for this mode. The preprocessing will additionally process a mask for the character in the video, and its size and shape can be adjusted by specifying some parameters.
+- `iterations` and `k` can make the mask larger, covering more area.
+- `w_len` and `h_len` can adjust the mask's shape. Smaller values will make the outline coarser, while larger values will make it finer.
+
+- A smaller, finer-contoured mask can allow for more of the original background to be preserved, but may potentially limit the character's generation area (considering potential appearance differences, this can lead to some shape leakage). A larger, coarser mask can allow the character generation to be more flexible and consistent, but because it includes more of the background, it might affect the background's consistency. We recommend users to adjust the relevant parameters based on their specific input data.

wan/modules/animate/preprocess/__init__.py

@@ -0,0 +1,3 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from .process_pipepline import ProcessPipeline
+from .video_predictor import SAM2VideoPredictor

wan/modules/animate/preprocess/human_visualization.py

@@ -0,0 +1,1357 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import os
+import cv2
+import time
+import math
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+from typing import Dict, List
+import random
+from pose2d_utils import AAPoseMeta
+
+
+def draw_handpose(canvas, keypoints, hand_score_th=0.6):
+    """
+    Draw keypoints and connections representing hand pose on a given canvas.
+
+    Args:
+        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the hand pose.
+        keypoints (List[Keypoint]| None): A list of Keypoint objects representing the hand keypoints to be drawn
+                                          or None if no keypoints are present.
+
+    Returns:
+        np.ndarray: A 3D numpy array representing the modified canvas with the drawn hand pose.
+
+    Note:
+        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
+    """
+    eps = 0.01
+
+    H, W, C = canvas.shape
+    stickwidth = max(int(min(H, W) / 200), 1)
+
+    edges = [
+        [0, 1],
+        [1, 2],
+        [2, 3],
+        [3, 4],
+        [0, 5],
+        [5, 6],
+        [6, 7],
+        [7, 8],
+        [0, 9],
+        [9, 10],
+        [10, 11],
+        [11, 12],
+        [0, 13],
+        [13, 14],
+        [14, 15],
+        [15, 16],
+        [0, 17],
+        [17, 18],
+        [18, 19],
+        [19, 20],
+    ]
+
+    for ie, (e1, e2) in enumerate(edges):
+        k1 = keypoints[e1]
+        k2 = keypoints[e2]
+        if k1 is None or k2 is None:
+            continue
+        if k1[2] < hand_score_th or k2[2] < hand_score_th:
+            continue
+
+        x1 = int(k1[0])
+        y1 = int(k1[1])
+        x2 = int(k2[0])
+        y2 = int(k2[1])
+        if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
+            cv2.line(
+                canvas,
+                (x1, y1),
+                (x2, y2),
+                matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * 255,
+                thickness=stickwidth,
+            )
+
+    for keypoint in keypoints:
+
+        if keypoint is None:
+            continue
+        if keypoint[2] < hand_score_th:
+            continue
+
+        x, y = keypoint[0], keypoint[1]
+        x = int(x)
+        y = int(y)
+        if x > eps and y > eps:
+            cv2.circle(canvas, (x, y), stickwidth, (0, 0, 255), thickness=-1)
+    return canvas
+
+
+def draw_handpose_new(canvas, keypoints, stickwidth_type='v2', hand_score_th=0.6):
+    """
+    Draw keypoints and connections representing hand pose on a given canvas.
+
+    Args:
+        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the hand pose.
+        keypoints (List[Keypoint]| None): A list of Keypoint objects representing the hand keypoints to be drawn
+                                          or None if no keypoints are present.
+
+    Returns:
+        np.ndarray: A 3D numpy array representing the modified canvas with the drawn hand pose.
+
+    Note:
+        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
+    """
+    eps = 0.01
+
+    H, W, C = canvas.shape
+    if stickwidth_type == 'v1':
+        stickwidth = max(int(min(H, W) / 200), 1)
+    elif stickwidth_type == 'v2':
+        stickwidth = max(max(int(min(H, W) / 200) - 1, 1) // 2, 1)
+
+    edges = [
+        [0, 1],
+        [1, 2],
+        [2, 3],
+        [3, 4],
+        [0, 5],
+        [5, 6],
+        [6, 7],
+        [7, 8],
+        [0, 9],
+        [9, 10],
+        [10, 11],
+        [11, 12],
+        [0, 13],
+        [13, 14],
+        [14, 15],
+        [15, 16],
+        [0, 17],
+        [17, 18],
+        [18, 19],
+        [19, 20],
+    ]
+
+    for ie, (e1, e2) in enumerate(edges):
+        k1 = keypoints[e1]
+        k2 = keypoints[e2]
+        if k1 is None or k2 is None:
+            continue
+        if k1[2] < hand_score_th or k2[2] < hand_score_th:
+            continue
+
+        x1 = int(k1[0])
+        y1 = int(k1[1])
+        x2 = int(k2[0])
+        y2 = int(k2[1])
+        if x1 > eps and y1 > eps and x2 > eps and y2 > eps:
+            cv2.line(
+                canvas,
+                (x1, y1),
+                (x2, y2),
+                matplotlib.colors.hsv_to_rgb([ie / float(len(edges)), 1.0, 1.0]) * 255,
+                thickness=stickwidth,
+            )
+
+    for keypoint in keypoints:
+
+        if keypoint is None:
+            continue
+        if keypoint[2] < hand_score_th:
+            continue
+
+        x, y = keypoint[0], keypoint[1]
+        x = int(x)
+        y = int(y)
+        if x > eps and y > eps:
+            cv2.circle(canvas, (x, y), stickwidth, (0, 0, 255), thickness=-1)
+    return canvas
+
+
+def draw_ellipse_by_2kp(img, keypoint1, keypoint2, color, threshold=0.6):
+    H, W, C = img.shape
+    stickwidth = max(int(min(H, W) / 200), 1)
+
+    if keypoint1[-1] < threshold or keypoint2[-1] < threshold:
+        return img
+
+    Y = np.array([keypoint1[0], keypoint2[0]])
+    X = np.array([keypoint1[1], keypoint2[1]])
+    mX = np.mean(X)
+    mY = np.mean(Y)
+    length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
+    angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
+    polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
+    cv2.fillConvexPoly(img, polygon, [int(float(c) * 0.6) for c in color])
+    return img
+
+
+def split_pose2d_kps_to_aa(kp2ds: np.ndarray) -> List[np.ndarray]:
+    """Convert the 133 keypoints from pose2d to body and hands keypoints.
+
+    Args:
+        kp2ds (np.ndarray): [133, 2]
+
+    Returns:
+        List[np.ndarray]: _description_
+    """
+    kp2ds_body = (
+        kp2ds[[0, 6, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 17, 20]]
+        + kp2ds[[0, 5, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 18, 21]]
+    ) / 2
+    kp2ds_lhand = kp2ds[91:112]
+    kp2ds_rhand = kp2ds[112:133]
+    return kp2ds_body.copy(), kp2ds_lhand.copy(), kp2ds_rhand.copy()
+
+
+def draw_aapose_by_meta(img, meta: AAPoseMeta, threshold=0.5, stick_width_norm=200, draw_hand=True, draw_head=True):
+    kp2ds = np.concatenate([meta.kps_body, meta.kps_body_p[:, None]], axis=1)
+    kp2ds_lhand = np.concatenate([meta.kps_lhand, meta.kps_lhand_p[:, None]], axis=1)
+    kp2ds_rhand = np.concatenate([meta.kps_rhand, meta.kps_rhand_p[:, None]], axis=1)
+    pose_img = draw_aapose(img, kp2ds, threshold, kp2ds_lhand=kp2ds_lhand, kp2ds_rhand=kp2ds_rhand, stick_width_norm=stick_width_norm, draw_hand=draw_hand, draw_head=draw_head)
+    return pose_img
+
+def draw_aapose_by_meta_new(img, meta: AAPoseMeta, threshold=0.5, stickwidth_type='v2', draw_hand=True, draw_head=True):
+    kp2ds = np.concatenate([meta.kps_body, meta.kps_body_p[:, None]], axis=1)
+    kp2ds_lhand = np.concatenate([meta.kps_lhand, meta.kps_lhand_p[:, None]], axis=1)
+    kp2ds_rhand = np.concatenate([meta.kps_rhand, meta.kps_rhand_p[:, None]], axis=1)
+    pose_img = draw_aapose_new(img, kp2ds, threshold, kp2ds_lhand=kp2ds_lhand, kp2ds_rhand=kp2ds_rhand,
+                               stickwidth_type=stickwidth_type, draw_hand=draw_hand, draw_head=draw_head)
+    return pose_img
+
+def draw_hand_by_meta(img, meta: AAPoseMeta, threshold=0.5, stick_width_norm=200):
+    kp2ds = np.concatenate([meta.kps_body, meta.kps_body_p[:, None] * 0], axis=1)
+    kp2ds_lhand = np.concatenate([meta.kps_lhand, meta.kps_lhand_p[:, None]], axis=1)
+    kp2ds_rhand = np.concatenate([meta.kps_rhand, meta.kps_rhand_p[:, None]], axis=1)
+    pose_img = draw_aapose(img, kp2ds, threshold, kp2ds_lhand=kp2ds_lhand, kp2ds_rhand=kp2ds_rhand, stick_width_norm=stick_width_norm, draw_hand=True, draw_head=False)
+    return pose_img
+
+
+def draw_aaface_by_meta(img, meta: AAPoseMeta, threshold=0.5, stick_width_norm=200, draw_hand=False, draw_head=True):
+    kp2ds = np.concatenate([meta.kps_body, meta.kps_body_p[:, None]], axis=1)
+    # kp2ds_lhand = np.concatenate([meta.kps_lhand, meta.kps_lhand_p[:, None]], axis=1)
+    # kp2ds_rhand = np.concatenate([meta.kps_rhand, meta.kps_rhand_p[:, None]], axis=1)
+    pose_img = draw_M(img, kp2ds, threshold, kp2ds_lhand=None, kp2ds_rhand=None, stick_width_norm=stick_width_norm, draw_hand=draw_hand, draw_head=draw_head)
+    return pose_img
+
+
+def draw_aanose_by_meta(img, meta: AAPoseMeta, threshold=0.5, stick_width_norm=100, draw_hand=False):
+    kp2ds = np.concatenate([meta.kps_body, meta.kps_body_p[:, None]], axis=1)
+    # kp2ds_lhand = np.concatenate([meta.kps_lhand, meta.kps_lhand_p[:, None]], axis=1)
+    # kp2ds_rhand = np.concatenate([meta.kps_rhand, meta.kps_rhand_p[:, None]], axis=1)
+    pose_img = draw_nose(img, kp2ds, threshold, kp2ds_lhand=None, kp2ds_rhand=None, stick_width_norm=stick_width_norm, draw_hand=draw_hand)
+    return pose_img
+
+
+def gen_face_motion_seq(img, metas: List[AAPoseMeta], threshold=0.5, stick_width_norm=200):
+
+    return 
+
+
+def draw_M(
+    img,
+    kp2ds,
+    threshold=0.6,
+    data_to_json=None,
+    idx=-1,
+    kp2ds_lhand=None,
+    kp2ds_rhand=None,
+    draw_hand=False,
+    stick_width_norm=200,
+    draw_head=True
+):
+    """
+    Draw keypoints and connections representing hand pose on a given canvas.
+
+    Args:
+        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the hand pose.
+        keypoints (List[Keypoint]| None): A list of Keypoint objects representing the hand keypoints to be drawn
+                                          or None if no keypoints are present.
+
+    Returns:
+        np.ndarray: A 3D numpy array representing the modified canvas with the drawn hand pose.
+
+    Note:
+        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
+    """
+
+    new_kep_list = [
+        "Nose",
+        "Neck",
+        "RShoulder",
+        "RElbow",
+        "RWrist",  # No.4
+        "LShoulder",
+        "LElbow",
+        "LWrist",  # No.7
+        "RHip",
+        "RKnee",
+        "RAnkle",  # No.10
+        "LHip",
+        "LKnee",
+        "LAnkle",  # No.13
+        "REye",
+        "LEye",
+        "REar",
+        "LEar",
+        "LToe",
+        "RToe",
+    ]
+    # kp2ds_body = (kp2ds.copy()[[0, 6, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 17, 20]] + \
+    #              kp2ds.copy()[[0, 5, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 18, 21]]) / 2
+    kp2ds = kp2ds.copy()
+    # import ipdb; ipdb.set_trace()
+    kp2ds[[1,2,3,4,5,6,7,8,9,10,11,12,13,18,19], 2] = 0
+    if not draw_head:
+        kp2ds[[0,14,15,16,17], 2] = 0
+    kp2ds_body = kp2ds
+    # kp2ds_body = kp2ds_body[:18]
+
+    # kp2ds_lhand = kp2ds.copy()[91:112]
+    # kp2ds_rhand = kp2ds.copy()[112:133]
+
+    limbSeq = [
+        # [2, 3],
+        # [2, 6],  # shoulders
+        # [3, 4],
+        # [4, 5],  # left arm
+        # [6, 7],
+        # [7, 8],  # right arm
+        # [2, 9],
+        # [9, 10],
+        # [10, 11],  # right leg
+        # [2, 12],
+        # [12, 13],
+        # [13, 14],  # left leg
+        # [2, 1],
+        [1, 15],
+        [15, 17],
+        [1, 16],
+        [16, 18],  # face (nose, eyes, ears)
+        # [14, 19],
+        # [11, 20],  # foot
+    ]
+
+    colors = [
+        # [255, 0, 0],
+        # [255, 85, 0],
+        # [255, 170, 0],
+        # [255, 255, 0],
+        # [170, 255, 0],
+        # [85, 255, 0],
+        # [0, 255, 0],
+        # [0, 255, 85],
+        # [0, 255, 170],
+        # [0, 255, 255],
+        # [0, 170, 255],
+        # [0, 85, 255],
+        # [0, 0, 255],
+        # [85, 0, 255],
+        [170, 0, 255],
+        [255, 0, 255],
+        [255, 0, 170],
+        [255, 0, 85],
+        # foot
+        # [200, 200, 0],
+        # [100, 100, 0],
+    ]
+
+    H, W, C = img.shape
+    stickwidth = max(int(min(H, W) / stick_width_norm), 1)
+
+    for _idx, ((k1_index, k2_index), color) in enumerate(zip(limbSeq, colors)):
+        keypoint1 = kp2ds_body[k1_index - 1]
+        keypoint2 = kp2ds_body[k2_index - 1]
+
+        if keypoint1[-1] < threshold or keypoint2[-1] < threshold:
+            continue
+
+        Y = np.array([keypoint1[0], keypoint2[0]])
+        X = np.array([keypoint1[1], keypoint2[1]])
+        mX = np.mean(X)
+        mY = np.mean(Y)
+        length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
+        angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
+        polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
+        cv2.fillConvexPoly(img, polygon, [int(float(c) * 0.6) for c in color])
+
+    for _idx, (keypoint, color) in enumerate(zip(kp2ds_body, colors)):
+        if keypoint[-1] < threshold:
+            continue
+        x, y = keypoint[0], keypoint[1]
+        # cv2.circle(canvas, (int(x), int(y)), 4, color, thickness=-1)
+        cv2.circle(img, (int(x), int(y)), stickwidth, color, thickness=-1)
+
+    if draw_hand:
+        img = draw_handpose(img, kp2ds_lhand, hand_score_th=threshold)
+        img = draw_handpose(img, kp2ds_rhand, hand_score_th=threshold)
+
+    kp2ds_body[:, 0] /= W
+    kp2ds_body[:, 1] /= H
+
+    if data_to_json is not None:
+        if idx == -1:
+            data_to_json.append(
+                {
+                    "image_id": "frame_{:05d}.jpg".format(len(data_to_json) + 1),
+                    "height": H,
+                    "width": W,
+                    "category_id": 1,
+                    "keypoints_body": kp2ds_body.tolist(),
+                    "keypoints_left_hand": kp2ds_lhand.tolist(),
+                    "keypoints_right_hand": kp2ds_rhand.tolist(),
+                }
+            )
+        else:
+            data_to_json[idx] = {
+                "image_id": "frame_{:05d}.jpg".format(idx + 1),
+                "height": H,
+                "width": W,
+                "category_id": 1,
+                "keypoints_body": kp2ds_body.tolist(),
+                "keypoints_left_hand": kp2ds_lhand.tolist(),
+                "keypoints_right_hand": kp2ds_rhand.tolist(),
+            }
+    return img
+
+
+def draw_nose(
+    img,
+    kp2ds,
+    threshold=0.6,
+    data_to_json=None,
+    idx=-1,
+    kp2ds_lhand=None,
+    kp2ds_rhand=None,
+    draw_hand=False,
+    stick_width_norm=200,
+):
+    """
+    Draw keypoints and connections representing hand pose on a given canvas.
+
+    Args:
+        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the hand pose.
+        keypoints (List[Keypoint]| None): A list of Keypoint objects representing the hand keypoints to be drawn
+                                          or None if no keypoints are present.
+
+    Returns:
+        np.ndarray: A 3D numpy array representing the modified canvas with the drawn hand pose.
+
+    Note:
+        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
+    """
+
+    new_kep_list = [
+        "Nose",
+        "Neck",
+        "RShoulder",
+        "RElbow",
+        "RWrist",  # No.4
+        "LShoulder",
+        "LElbow",
+        "LWrist",  # No.7
+        "RHip",
+        "RKnee",
+        "RAnkle",  # No.10
+        "LHip",
+        "LKnee",
+        "LAnkle",  # No.13
+        "REye",
+        "LEye",
+        "REar",
+        "LEar",
+        "LToe",
+        "RToe",
+    ]
+    # kp2ds_body = (kp2ds.copy()[[0, 6, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 17, 20]] + \
+    #              kp2ds.copy()[[0, 5, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 18, 21]]) / 2
+    kp2ds = kp2ds.copy()
+    kp2ds[1:, 2] = 0
+    # kp2ds[0, 2] = 1
+    kp2ds_body = kp2ds
+    # kp2ds_body = kp2ds_body[:18]
+
+    # kp2ds_lhand = kp2ds.copy()[91:112]
+    # kp2ds_rhand = kp2ds.copy()[112:133]
+
+    limbSeq = [
+        # [2, 3],
+        # [2, 6],  # shoulders
+        # [3, 4],
+        # [4, 5],  # left arm
+        # [6, 7],
+        # [7, 8],  # right arm
+        # [2, 9],
+        # [9, 10],
+        # [10, 11],  # right leg
+        # [2, 12],
+        # [12, 13],
+        # [13, 14],  # left leg
+        # [2, 1],
+        [1, 15],
+        [15, 17],
+        [1, 16],
+        [16, 18],  # face (nose, eyes, ears)
+        # [14, 19],
+        # [11, 20],  # foot
+    ]
+
+    colors = [
+        # [255, 0, 0],
+        # [255, 85, 0],
+        # [255, 170, 0],
+        # [255, 255, 0],
+        # [170, 255, 0],
+        # [85, 255, 0],
+        # [0, 255, 0],
+        # [0, 255, 85],
+        # [0, 255, 170],
+        # [0, 255, 255],
+        # [0, 170, 255],
+        # [0, 85, 255],
+        # [0, 0, 255],
+        # [85, 0, 255],
+        [170, 0, 255],
+        # [255, 0, 255],
+        # [255, 0, 170],
+        # [255, 0, 85],
+        # foot
+        # [200, 200, 0],
+        # [100, 100, 0],
+    ]
+
+    H, W, C = img.shape
+    stickwidth = max(int(min(H, W) / stick_width_norm), 1)
+
+    # for _idx, ((k1_index, k2_index), color) in enumerate(zip(limbSeq, colors)):
+    #     keypoint1 = kp2ds_body[k1_index - 1]
+    #     keypoint2 = kp2ds_body[k2_index - 1]
+
+    #     if keypoint1[-1] < threshold or keypoint2[-1] < threshold:
+    #         continue
+
+    #     Y = np.array([keypoint1[0], keypoint2[0]])
+    #     X = np.array([keypoint1[1], keypoint2[1]])
+    #     mX = np.mean(X)
+    #     mY = np.mean(Y)
+    #     length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
+    #     angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
+    #     polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
+    #     cv2.fillConvexPoly(img, polygon, [int(float(c) * 0.6) for c in color])
+
+    for _idx, (keypoint, color) in enumerate(zip(kp2ds_body, colors)):
+        if keypoint[-1] < threshold:
+            continue
+        x, y = keypoint[0], keypoint[1]
+        # cv2.circle(canvas, (int(x), int(y)), 4, color, thickness=-1)
+        cv2.circle(img, (int(x), int(y)), stickwidth, color, thickness=-1)
+
+    if draw_hand:
+        img = draw_handpose(img, kp2ds_lhand, hand_score_th=threshold)
+        img = draw_handpose(img, kp2ds_rhand, hand_score_th=threshold)
+
+    kp2ds_body[:, 0] /= W
+    kp2ds_body[:, 1] /= H
+
+    if data_to_json is not None:
+        if idx == -1:
+            data_to_json.append(
+                {
+                    "image_id": "frame_{:05d}.jpg".format(len(data_to_json) + 1),
+                    "height": H,
+                    "width": W,
+                    "category_id": 1,
+                    "keypoints_body": kp2ds_body.tolist(),
+                    "keypoints_left_hand": kp2ds_lhand.tolist(),
+                    "keypoints_right_hand": kp2ds_rhand.tolist(),
+                }
+            )
+        else:
+            data_to_json[idx] = {
+                "image_id": "frame_{:05d}.jpg".format(idx + 1),
+                "height": H,
+                "width": W,
+                "category_id": 1,
+                "keypoints_body": kp2ds_body.tolist(),
+                "keypoints_left_hand": kp2ds_lhand.tolist(),
+                "keypoints_right_hand": kp2ds_rhand.tolist(),
+            }
+    return img
+
+
+def draw_aapose(
+    img,
+    kp2ds,
+    threshold=0.6,
+    data_to_json=None,
+    idx=-1,
+    kp2ds_lhand=None,
+    kp2ds_rhand=None,
+    draw_hand=False,
+    stick_width_norm=200,
+    draw_head=True
+):
+    """
+    Draw keypoints and connections representing hand pose on a given canvas.
+
+    Args:
+        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the hand pose.
+        keypoints (List[Keypoint]| None): A list of Keypoint objects representing the hand keypoints to be drawn
+                                          or None if no keypoints are present.
+
+    Returns:
+        np.ndarray: A 3D numpy array representing the modified canvas with the drawn hand pose.
+
+    Note:
+        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
+    """
+
+    new_kep_list = [
+        "Nose",
+        "Neck",
+        "RShoulder",
+        "RElbow",
+        "RWrist",  # No.4
+        "LShoulder",
+        "LElbow",
+        "LWrist",  # No.7
+        "RHip",
+        "RKnee",
+        "RAnkle",  # No.10
+        "LHip",
+        "LKnee",
+        "LAnkle",  # No.13
+        "REye",
+        "LEye",
+        "REar",
+        "LEar",
+        "LToe",
+        "RToe",
+    ]
+    # kp2ds_body = (kp2ds.copy()[[0, 6, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 17, 20]] + \
+    #              kp2ds.copy()[[0, 5, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 18, 21]]) / 2
+    kp2ds = kp2ds.copy()
+    if not draw_head:
+        kp2ds[[0,14,15,16,17], 2] = 0
+    kp2ds_body = kp2ds
+
+    # kp2ds_lhand = kp2ds.copy()[91:112]
+    # kp2ds_rhand = kp2ds.copy()[112:133]
+
+    limbSeq = [
+        [2, 3],
+        [2, 6],  # shoulders
+        [3, 4],
+        [4, 5],  # left arm
+        [6, 7],
+        [7, 8],  # right arm
+        [2, 9],
+        [9, 10],
+        [10, 11],  # right leg
+        [2, 12],
+        [12, 13],
+        [13, 14],  # left leg
+        [2, 1],
+        [1, 15],
+        [15, 17],
+        [1, 16],
+        [16, 18],  # face (nose, eyes, ears)
+        [14, 19],
+        [11, 20],  # foot
+    ]
+
+    colors = [
+        [255, 0, 0],
+        [255, 85, 0],
+        [255, 170, 0],
+        [255, 255, 0],
+        [170, 255, 0],
+        [85, 255, 0],
+        [0, 255, 0],
+        [0, 255, 85],
+        [0, 255, 170],
+        [0, 255, 255],
+        [0, 170, 255],
+        [0, 85, 255],
+        [0, 0, 255],
+        [85, 0, 255],
+        [170, 0, 255],
+        [255, 0, 255],
+        [255, 0, 170],
+        [255, 0, 85],
+        # foot
+        [200, 200, 0],
+        [100, 100, 0],
+    ]
+
+    H, W, C = img.shape
+    stickwidth = max(int(min(H, W) / stick_width_norm), 1)
+
+    for _idx, ((k1_index, k2_index), color) in enumerate(zip(limbSeq, colors)):
+        keypoint1 = kp2ds_body[k1_index - 1]
+        keypoint2 = kp2ds_body[k2_index - 1]
+
+        if keypoint1[-1] < threshold or keypoint2[-1] < threshold:
+            continue
+
+        Y = np.array([keypoint1[0], keypoint2[0]])
+        X = np.array([keypoint1[1], keypoint2[1]])
+        mX = np.mean(X)
+        mY = np.mean(Y)
+        length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
+        angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
+        polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
+        cv2.fillConvexPoly(img, polygon, [int(float(c) * 0.6) for c in color])
+
+    for _idx, (keypoint, color) in enumerate(zip(kp2ds_body, colors)):
+        if keypoint[-1] < threshold:
+            continue
+        x, y = keypoint[0], keypoint[1]
+        # cv2.circle(canvas, (int(x), int(y)), 4, color, thickness=-1)
+        cv2.circle(img, (int(x), int(y)), stickwidth, color, thickness=-1)
+
+    if draw_hand:
+        img = draw_handpose(img, kp2ds_lhand, hand_score_th=threshold)
+        img = draw_handpose(img, kp2ds_rhand, hand_score_th=threshold)
+
+    kp2ds_body[:, 0] /= W
+    kp2ds_body[:, 1] /= H
+
+    if data_to_json is not None:
+        if idx == -1:
+            data_to_json.append(
+                {
+                    "image_id": "frame_{:05d}.jpg".format(len(data_to_json) + 1),
+                    "height": H,
+                    "width": W,
+                    "category_id": 1,
+                    "keypoints_body": kp2ds_body.tolist(),
+                    "keypoints_left_hand": kp2ds_lhand.tolist(),
+                    "keypoints_right_hand": kp2ds_rhand.tolist(),
+                }
+            )
+        else:
+            data_to_json[idx] = {
+                "image_id": "frame_{:05d}.jpg".format(idx + 1),
+                "height": H,
+                "width": W,
+                "category_id": 1,
+                "keypoints_body": kp2ds_body.tolist(),
+                "keypoints_left_hand": kp2ds_lhand.tolist(),
+                "keypoints_right_hand": kp2ds_rhand.tolist(),
+            }
+    return img
+
+
+def draw_aapose_new(
+    img,
+    kp2ds,
+    threshold=0.6,
+    data_to_json=None,
+    idx=-1,
+    kp2ds_lhand=None,
+    kp2ds_rhand=None,
+    draw_hand=False,
+    stickwidth_type='v2',
+    draw_head=True
+):
+    """
+    Draw keypoints and connections representing hand pose on a given canvas.
+
+    Args:
+        canvas (np.ndarray): A 3D numpy array representing the canvas (image) on which to draw the hand pose.
+        keypoints (List[Keypoint]| None): A list of Keypoint objects representing the hand keypoints to be drawn
+                                          or None if no keypoints are present.
+
+    Returns:
+        np.ndarray: A 3D numpy array representing the modified canvas with the drawn hand pose.
+
+    Note:
+        The function expects the x and y coordinates of the keypoints to be normalized between 0 and 1.
+    """
+
+    new_kep_list = [
+        "Nose",
+        "Neck",
+        "RShoulder",
+        "RElbow",
+        "RWrist",  # No.4
+        "LShoulder",
+        "LElbow",
+        "LWrist",  # No.7
+        "RHip",
+        "RKnee",
+        "RAnkle",  # No.10
+        "LHip",
+        "LKnee",
+        "LAnkle",  # No.13
+        "REye",
+        "LEye",
+        "REar",
+        "LEar",
+        "LToe",
+        "RToe",
+    ]
+    # kp2ds_body = (kp2ds.copy()[[0, 6, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 17, 20]] + \
+    #              kp2ds.copy()[[0, 5, 6, 8, 10, 5, 7, 9, 12, 14, 16, 11, 13, 15, 2, 1, 4, 3, 18, 21]]) / 2
+    kp2ds = kp2ds.copy()
+    if not draw_head:
+        kp2ds[[0,14,15,16,17], 2] = 0
+    kp2ds_body = kp2ds
+
+    # kp2ds_lhand = kp2ds.copy()[91:112]
+    # kp2ds_rhand = kp2ds.copy()[112:133]
+
+    limbSeq = [
+        [2, 3],
+        [2, 6],  # shoulders
+        [3, 4],
+        [4, 5],  # left arm
+        [6, 7],
+        [7, 8],  # right arm
+        [2, 9],
+        [9, 10],
+        [10, 11],  # right leg
+        [2, 12],
+        [12, 13],
+        [13, 14],  # left leg
+        [2, 1],
+        [1, 15],
+        [15, 17],
+        [1, 16],
+        [16, 18],  # face (nose, eyes, ears)
+        [14, 19],
+        [11, 20],  # foot
+    ]
+
+    colors = [
+        [255, 0, 0],
+        [255, 85, 0],
+        [255, 170, 0],
+        [255, 255, 0],
+        [170, 255, 0],
+        [85, 255, 0],
+        [0, 255, 0],
+        [0, 255, 85],
+        [0, 255, 170],
+        [0, 255, 255],
+        [0, 170, 255],
+        [0, 85, 255],
+        [0, 0, 255],
+        [85, 0, 255],
+        [170, 0, 255],
+        [255, 0, 255],
+        [255, 0, 170],
+        [255, 0, 85],
+        # foot
+        [200, 200, 0],
+        [100, 100, 0],
+    ]
+
+    H, W, C = img.shape
+    H, W, C = img.shape
+
+    if stickwidth_type == 'v1':
+        stickwidth = max(int(min(H, W) / 200), 1)
+    elif stickwidth_type == 'v2':
+        stickwidth = max(int(min(H, W) / 200) - 1, 1)
+    else:
+        raise
+
+    for _idx, ((k1_index, k2_index), color) in enumerate(zip(limbSeq, colors)):
+        keypoint1 = kp2ds_body[k1_index - 1]
+        keypoint2 = kp2ds_body[k2_index - 1]
+
+        if keypoint1[-1] < threshold or keypoint2[-1] < threshold:
+            continue
+
+        Y = np.array([keypoint1[0], keypoint2[0]])
+        X = np.array([keypoint1[1], keypoint2[1]])
+        mX = np.mean(X)
+        mY = np.mean(Y)
+        length = ((X[0] - X[1]) ** 2 + (Y[0] - Y[1]) ** 2) ** 0.5
+        angle = math.degrees(math.atan2(X[0] - X[1], Y[0] - Y[1]))
+        polygon = cv2.ellipse2Poly((int(mY), int(mX)), (int(length / 2), stickwidth), int(angle), 0, 360, 1)
+        cv2.fillConvexPoly(img, polygon, [int(float(c) * 0.6) for c in color])
+
+    for _idx, (keypoint, color) in enumerate(zip(kp2ds_body, colors)):
+        if keypoint[-1] < threshold:
+            continue
+        x, y = keypoint[0], keypoint[1]
+        # cv2.circle(canvas, (int(x), int(y)), 4, color, thickness=-1)
+        cv2.circle(img, (int(x), int(y)), stickwidth, color, thickness=-1)
+
+    if draw_hand:
+        img = draw_handpose_new(img, kp2ds_lhand, stickwidth_type=stickwidth_type, hand_score_th=threshold)
+        img = draw_handpose_new(img, kp2ds_rhand, stickwidth_type=stickwidth_type, hand_score_th=threshold)
+
+    kp2ds_body[:, 0] /= W
+    kp2ds_body[:, 1] /= H
+
+    if data_to_json is not None:
+        if idx == -1:
+            data_to_json.append(
+                {
+                    "image_id": "frame_{:05d}.jpg".format(len(data_to_json) + 1),
+                    "height": H,
+                    "width": W,
+                    "category_id": 1,
+                    "keypoints_body": kp2ds_body.tolist(),
+                    "keypoints_left_hand": kp2ds_lhand.tolist(),
+                    "keypoints_right_hand": kp2ds_rhand.tolist(),
+                }
+            )
+        else:
+            data_to_json[idx] = {
+                "image_id": "frame_{:05d}.jpg".format(idx + 1),
+                "height": H,
+                "width": W,
+                "category_id": 1,
+                "keypoints_body": kp2ds_body.tolist(),
+                "keypoints_left_hand": kp2ds_lhand.tolist(),
+                "keypoints_right_hand": kp2ds_rhand.tolist(),
+            }
+    return img
+
+
+def draw_bbox(img, bbox, color=(255, 0, 0)):
+    img = load_image(img)
+    bbox = [int(bbox_tmp) for bbox_tmp in bbox]
+    cv2.rectangle(img, (bbox[0], bbox[1]), (bbox[2], bbox[3]), color, 2)
+    return img
+
+
+def draw_kp2ds(img, kp2ds, threshold=0, color=(255, 0, 0), skeleton=None, reverse=False):
+    img = load_image(img, reverse)
+
+    if skeleton is not None:
+        if skeleton == "coco17":
+            skeleton_list = [
+                [6, 8],
+                [8, 10],
+                [5, 7],
+                [7, 9],
+                [11, 13],
+                [13, 15],
+                [12, 14],
+                [14, 16],
+                [5, 6],
+                [6, 12],
+                [12, 11],
+                [11, 5],
+            ]
+            color_list = [
+                (255, 0, 0),
+                (0, 255, 0),
+                (0, 0, 255),
+                (255, 255, 0),
+                (255, 0, 255),
+                (0, 255, 255),
+            ]
+        elif skeleton == "cocowholebody":
+            skeleton_list = [
+                [6, 8],
+                [8, 10],
+                [5, 7],
+                [7, 9],
+                [11, 13],
+                [13, 15],
+                [12, 14],
+                [14, 16],
+                [5, 6],
+                [6, 12],
+                [12, 11],
+                [11, 5],
+                [15, 17],
+                [15, 18],
+                [15, 19],
+                [16, 20],
+                [16, 21],
+                [16, 22],
+                [91, 92, 93, 94, 95],
+                [91, 96, 97, 98, 99],
+                [91, 100, 101, 102, 103],
+                [91, 104, 105, 106, 107],
+                [91, 108, 109, 110, 111],
+                [112, 113, 114, 115, 116],
+                [112, 117, 118, 119, 120],
+                [112, 121, 122, 123, 124],
+                [112, 125, 126, 127, 128],
+                [112, 129, 130, 131, 132],
+            ]
+            color_list = [
+                (255, 0, 0),
+                (0, 255, 0),
+                (0, 0, 255),
+                (255, 255, 0),
+                (255, 0, 255),
+                (0, 255, 255),
+            ]
+        else:
+            color_list = [color]
+        for _idx, _skeleton in enumerate(skeleton_list):
+            for i in range(len(_skeleton) - 1):
+                cv2.line(
+                    img,
+                    (int(kp2ds[_skeleton[i], 0]), int(kp2ds[_skeleton[i], 1])),
+                    (int(kp2ds[_skeleton[i + 1], 0]), int(kp2ds[_skeleton[i + 1], 1])),
+                    color_list[_idx % len(color_list)],
+                    3,
+                )
+
+    for _idx, kp2d in enumerate(kp2ds):
+        if kp2d[2] > threshold:
+            cv2.circle(img, (int(kp2d[0]), int(kp2d[1])), 3, color, -1)
+            # cv2.putText(img,
+            #         str(_idx),
+            #         (int(kp2d[0, i, 0])*1,
+            #             int(kp2d[0, i, 1])*1),
+            #         cv2.FONT_HERSHEY_SIMPLEX,
+            #         0.75,
+            #         color,
+            #         2
+            #         )
+
+    return img
+
+
+def draw_mask(img, mask, background=0, return_rgba=False):
+    img = load_image(img)
+    h, w, _ = img.shape
+    if type(background) == int:
+        background = np.ones((h, w, 3)).astype(np.uint8) * 255 * background
+    backgournd = cv2.resize(background, (w, h))
+    img_rgba = np.concatenate([img, mask], -1)
+    return alphaMerge(img_rgba, background, 0, 0, return_rgba=True)
+
+
+def draw_pcd(pcd_list, save_path=None):
+    fig = plt.figure()
+    ax = fig.add_subplot(111, projection="3d")
+
+    color_list = ["r", "g", "b", "y", "p"]
+
+    for _idx, _pcd in enumerate(pcd_list):
+        ax.scatter(_pcd[:, 0], _pcd[:, 1], _pcd[:, 2], c=color_list[_idx], marker="o")
+
+    ax.set_xlabel("X")
+    ax.set_ylabel("Y")
+    ax.set_zlabel("Z")
+
+    if save_path is not None:
+        plt.savefig(save_path)
+    else:
+        plt.savefig("tmp.png")
+
+
+def load_image(img, reverse=False):
+    if type(img) == str:
+        img = cv2.imread(img)
+    if reverse:
+        img = img.astype(np.float32)
+        img = img[:, :, ::-1]
+        img = img.astype(np.uint8)
+    return img
+
+
+def draw_skeleten(meta):
+    kps = []
+    for i, kp in enumerate(meta["keypoints_body"]):
+        if kp is None:
+            # if kp is None:
+            kps.append([0, 0, 0])
+        else:
+            kps.append([*kp, 1])
+    kps = np.array(kps)
+
+    kps[:, 0] *= meta["width"]
+    kps[:, 1] *= meta["height"]
+    pose_img = np.zeros([meta["height"], meta["width"], 3], dtype=np.uint8)
+
+    pose_img = draw_aapose(
+        pose_img,
+        kps,
+        draw_hand=True,
+        kp2ds_lhand=meta["keypoints_left_hand"],
+        kp2ds_rhand=meta["keypoints_right_hand"],
+    )
+    return pose_img
+
+
+def draw_skeleten_with_pncc(pncc: np.ndarray, meta: Dict) -> np.ndarray:
+    """
+    Args:
+        pncc: [H,W,3]
+        meta: required keys: keypoints_body: [N, 3] keypoints_left_hand, keypoints_right_hand
+    Return:
+        np.ndarray [H, W, 3]
+    """
+    # preprocess keypoints
+    kps = []
+    for i, kp in enumerate(meta["keypoints_body"]):
+        if kp is None:
+            # if kp is None:
+            kps.append([0, 0, 0])
+        elif i in [14, 15, 16, 17]:
+            kps.append([0, 0, 0])
+        else:
+            kps.append([*kp])
+    kps = np.stack(kps)
+
+    kps[:, 0] *= pncc.shape[1]
+    kps[:, 1] *= pncc.shape[0]
+
+    # draw neck
+    canvas = np.zeros_like(pncc)
+    if kps[0][2] > 0.6 and kps[1][2] > 0.6:
+        canvas = draw_ellipse_by_2kp(canvas, kps[0], kps[1], [0, 0, 255])
+
+    # draw pncc
+    mask = (pncc > 0).max(axis=2)
+    canvas[mask] = pncc[mask]
+    pncc = canvas
+
+    # draw other skeleten
+    kps[0] = 0
+
+    meta["keypoints_left_hand"][:, 0] *= meta["width"]
+    meta["keypoints_left_hand"][:, 1] *= meta["height"]
+
+    meta["keypoints_right_hand"][:, 0] *= meta["width"]
+    meta["keypoints_right_hand"][:, 1] *= meta["height"]
+    pose_img = draw_aapose(
+        pncc,
+        kps,
+        draw_hand=True,
+        kp2ds_lhand=meta["keypoints_left_hand"],
+        kp2ds_rhand=meta["keypoints_right_hand"],
+    )
+    return pose_img
+
+
+FACE_CUSTOM_STYLE = {
+    "eyeball": {"indexs": [68, 69], "color": [255, 255, 255], "connect": False},
+    "left_eyebrow": {"indexs": [17, 18, 19, 20, 21], "color": [0, 255, 0]},
+    "right_eyebrow": {"indexs": [22, 23, 24, 25, 26], "color": [0, 0, 255]},
+    "left_eye": {"indexs": [36, 37, 38, 39, 40, 41], "color": [255, 255, 0], "close": True},
+    "right_eye": {"indexs": [42, 43, 44, 45, 46, 47], "color": [255, 0, 255], "close": True},
+    "mouth_outside": {"indexs": list(range(48, 60)), "color": [100, 255, 50], "close": True},
+    "mouth_inside": {"indexs": [60, 61, 62, 63, 64, 65, 66, 67], "color": [255, 100, 50], "close": True},
+}
+
+
+def draw_face_kp(img, kps, thickness=2, style=FACE_CUSTOM_STYLE):
+    """
+    Args:
+        img: [H, W, 3]
+        kps: [70, 2]
+    """
+    img = img.copy()
+    for key, item in style.items():
+        pts = np.array(kps[item["indexs"]]).astype(np.int32)
+        connect = item.get("connect", True)
+        color = item["color"]
+        close = item.get("close", False)
+        if connect:
+            cv2.polylines(img, [pts], close, color, thickness=thickness)
+        else:
+            for kp in pts:
+                kp = np.array(kp).astype(np.int32)
+                cv2.circle(img, kp, thickness * 2, color=color, thickness=-1)
+    return img
+
+
+def draw_traj(metas: List[AAPoseMeta], threshold=0.6):
+
+    colors = [[255, 0, 0], [255, 85, 0], [255, 170, 0], [255, 255, 0], [170, 255, 0], [85, 255, 0], [0, 255, 0], \
+                [0, 255, 85], [0, 255, 170], [0, 255, 255], [0, 170, 255], [0, 85, 255], [0, 0, 255], [85, 0, 255], \
+                [170, 0, 255], [255, 0, 255], [255, 0, 170], [255, 0, 85], [100, 255, 50], [255, 100, 50], 
+                # foot
+                [200, 200, 0],
+                [100, 100, 0]
+                ]
+    limbSeq = [
+                    [1, 2], [1, 5],     # shoulders
+                    [2, 3], [3, 4],     # left arm
+                    [5, 6], [6, 7],     # right arm
+                    [1, 8], [8, 9], [9, 10],    # right leg 
+                    [1, 11], [11, 12], [12, 13],  # left leg
+                     # face (nose, eyes, ears)
+                    [13, 18], [10, 19] # foot
+                ]
+    
+    face_seq = [[1, 0], [0, 14], [14, 16], [0, 15], [15, 17]] 
+    kp_body = np.array([meta.kps_body for meta in metas])
+    kp_body_p = np.array([meta.kps_body_p for meta in metas])
+    
+    
+    face_seq = random.sample(face_seq, 2)
+
+    kp_lh = np.array([meta.kps_lhand for meta in metas])
+    kp_rh = np.array([meta.kps_rhand for meta in metas])
+
+    kp_lh_p = np.array([meta.kps_lhand_p for meta in metas])
+    kp_rh_p = np.array([meta.kps_rhand_p for meta in metas])
+
+    # kp_lh = np.concatenate([kp_lh, kp_lh_p], axis=-1)
+    # kp_rh = np.concatenate([kp_rh, kp_rh_p], axis=-1)
+    
+    new_limbSeq = []
+    key_point_list = []
+    for _idx, ((k1_index, k2_index)) in enumerate(limbSeq):
+        
+        vis = (kp_body_p[:, k1_index] > threshold) * (kp_body_p[:, k2_index] > threshold) * 1
+        if vis.sum() * 1.0 / vis.shape[0] > 0.4:
+            new_limbSeq.append([k1_index, k2_index])
+
+    for _idx, ((k1_index, k2_index)) in enumerate(limbSeq):
+
+        keypoint1 = kp_body[:, k1_index - 1]
+        keypoint2 = kp_body[:, k2_index - 1]
+        interleave = random.randint(4, 7)
+        randind = random.randint(0, interleave - 1)
+        # randind = random.rand(range(interleave), sampling_num)
+
+        Y = np.array([keypoint1[:, 0], keypoint2[:, 0]]) 
+        X = np.array([keypoint1[:, 1], keypoint2[:, 1]])
+
+        vis = (keypoint1[:, -1] > threshold) * (keypoint2[:, -1] > threshold) * 1
+
+        # for randidx in randind:
+        t = randind / interleave
+        x = (1-t)*Y[0, :] + t*Y[1, :]
+        y = (1-t)*X[0, :] + t*X[1, :]
+
+        # np.array([1])
+        x = x.astype(int)
+        y = y.astype(int)
+
+        new_array = np.array([x, y, vis]).T
+        
+        key_point_list.append(new_array)
+    
+    indx_lh = random.randint(0, kp_lh.shape[1] - 1)
+    lh = kp_lh[:, indx_lh, :]
+    lh_p = kp_lh_p[:, indx_lh:indx_lh+1]
+    lh = np.concatenate([lh, lh_p], axis=-1)
+    
+    indx_rh = random.randint(0, kp_rh.shape[1] - 1)
+    rh = kp_rh[:, random.randint(0, kp_rh.shape[1] - 1), :]
+    rh_p = kp_rh_p[:, indx_rh:indx_rh+1]
+    rh = np.concatenate([rh, rh_p], axis=-1)
+
+
+
+    lh[-1, :] = (lh[-1, :] > threshold) * 1
+    rh[-1, :] = (rh[-1, :] > threshold) * 1
+
+    # print(rh.shape, new_array.shape)
+    # exit()
+    key_point_list.append(lh.astype(int))
+    key_point_list.append(rh.astype(int))
+
+    
+    key_points_list = np.stack(key_point_list)
+    num_points = len(key_points_list)
+    sample_colors = random.sample(colors, num_points)
+
+    stickwidth = max(int(min(metas[0].width, metas[0].height) / 150), 2)
+
+    image_list_ori = []
+    for i in range(key_points_list.shape[-2]):
+        _image_vis = np.zeros((metas[0].width, metas[0].height, 3))
+        points = key_points_list[:, i, :]
+        for idx, point in enumerate(points):
+            x, y, vis = point
+            if vis == 1:
+                cv2.circle(_image_vis, (x, y), stickwidth, sample_colors[idx], thickness=-1)
+        
+        image_list_ori.append(_image_vis)
+    
+    return image_list_ori
+
+    return [np.zeros([meta.width, meta.height, 3], dtype=np.uint8) for meta in metas]
+
+
+if __name__ == "__main__":
+    meta = {
+        "image_id": "00472.jpg",
+        "height": 540,
+        "width": 414,
+        "category_id": 1,
+        "keypoints_body": [
+            [0.5084776947463768, 0.11350188078703703],
+            [0.504467655495169, 0.20419560185185184],
+            [0.3982016153381642, 0.198046875],
+            [0.3841664779589372, 0.34869068287037036],
+            [0.3901815368357488, 0.4670536747685185],
+            [0.610733695652174, 0.2103443287037037],
+            [0.6167487545289855, 0.3517650462962963],
+            [0.6448190292874396, 0.4762767650462963],
+            [0.4523371452294686, 0.47320240162037036],
+            [0.4503321256038647, 0.6776475694444445],
+            [0.47639738073671495, 0.8544234664351852],
+            [0.5766483620169082, 0.47320240162037036],
+            [0.5666232638888888, 0.6761103877314815],
+            [0.534542949879227, 0.863646556712963],
+            [0.4864224788647343, 0.09505570023148148],
+            [0.5285278910024155, 0.09351851851851851],
+            [0.46236224335748793, 0.10581597222222222],
+            [0.5586031853864735, 0.10274160879629629],
+            [0.4994551064311594, 0.9405056423611111],
+            [0.4152442821557971, 0.9312825520833333],
+        ],
+        "keypoints_left_hand": [
+            [267.78515625, 263.830078125, 1.2840936183929443],
+            [265.294921875, 269.640625, 1.2546794414520264],
+            [263.634765625, 277.111328125, 1.2863062620162964],
+            [262.8046875, 285.412109375, 1.267038345336914],
+            [261.14453125, 292.8828125, 1.280144453048706],
+            [273.595703125, 281.26171875, 1.2592815160751343],
+            [271.10546875, 291.22265625, 1.3256099224090576],
+            [265.294921875, 294.54296875, 1.2368024587631226],
+            [261.14453125, 294.54296875, 0.9771889448165894],
+            [274.42578125, 282.091796875, 1.250044584274292],
+            [269.4453125, 291.22265625, 1.2571144104003906],
+            [264.46484375, 292.8828125, 1.177802324295044],
+            [260.314453125, 292.052734375, 0.9283463358879089],
+            [273.595703125, 282.091796875, 1.1834490299224854],
+            [269.4453125, 290.392578125, 1.188171625137329],
+            [265.294921875, 290.392578125, 1.192609429359436],
+            [261.974609375, 289.5625, 0.9366656541824341],
+            [271.935546875, 281.26171875, 1.0946396589279175],
+            [268.615234375, 287.072265625, 0.9906131029129028],
+            [265.294921875, 287.90234375, 1.0219476222991943],
+            [262.8046875, 287.072265625, 0.9240120053291321],
+        ],
+        "keypoints_right_hand": [
+            [161.53515625, 258.849609375, 1.2069408893585205],
+            [168.17578125, 263.0, 1.1846840381622314],
+            [173.986328125, 269.640625, 1.1435924768447876],
+            [173.986328125, 277.94140625, 1.1802611351013184],
+            [173.986328125, 286.2421875, 1.2599592208862305],
+            [165.685546875, 275.451171875, 1.0633569955825806],
+            [167.345703125, 286.2421875, 1.1693341732025146],
+            [169.8359375, 291.22265625, 1.2698509693145752],
+            [170.666015625, 294.54296875, 1.0619274377822876],
+            [160.705078125, 276.28125, 1.0995020866394043],
+            [163.1953125, 287.90234375, 1.2735884189605713],
+            [166.515625, 291.22265625, 1.339503526687622],
+            [169.005859375, 294.54296875, 1.0835273265838623],
+            [157.384765625, 277.111328125, 1.0866981744766235],
+            [161.53515625, 287.072265625, 1.2468621730804443],
+            [164.025390625, 289.5625, 1.2817761898040771],
+            [166.515625, 292.052734375, 1.099466323852539],
+            [155.724609375, 277.111328125, 1.1065717935562134],
+            [159.044921875, 285.412109375, 1.1924479007720947],
+            [160.705078125, 287.072265625, 1.1304771900177002],
+            [162.365234375, 287.90234375, 1.0040509700775146],
+        ],
+    }
+    demo_meta = AAPoseMeta(meta)
+    res = draw_traj([demo_meta]*5)
+    cv2.imwrite("traj.png", res[0][..., ::-1])

wan/modules/animate/preprocess/pose2d.py

@@ -0,0 +1,430 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import os
+import cv2
+from typing import Union, List
+
+import numpy as np
+import torch
+import onnxruntime
+
+from pose2d_utils import (
+    read_img,
+    box_convert_simple,
+    bbox_from_detector,
+    crop,
+    keypoints_from_heatmaps,
+    load_pose_metas_from_kp2ds_seq
+)
+
+
+class SimpleOnnxInference(object):
+    def __init__(self, checkpoint, device='cuda', reverse_input=False, **kwargs):
+        if isinstance(device, str):
+            device = torch.device(device)
+        if device.type == 'cuda':
+            device = '{}:{}'.format(device.type, device.index)
+            providers = [("CUDAExecutionProvider", {"device_id": device[-1:] if device[-1] in [str(_i) for _i in range(10)] else "0"}), "CPUExecutionProvider"]
+        else:
+            providers = ["CPUExecutionProvider"]
+        self.device = device
+        if not os.path.exists(checkpoint):
+            raise RuntimeError("{} is not existed!".format(checkpoint))
+        
+        if os.path.isdir(checkpoint):
+            checkpoint = os.path.join(checkpoint, 'end2end.onnx')
+
+        self.session = onnxruntime.InferenceSession(checkpoint,
+                                                    providers=providers
+                                                    )
+        self.input_name = self.session.get_inputs()[0].name
+        self.output_name = self.session.get_outputs()[0].name
+        self.input_resolution = self.session.get_inputs()[0].shape[2:] if not reverse_input else self.session.get_inputs()[0].shape[2:][::-1]
+        self.input_resolution = np.array(self.input_resolution)
+        
+
+    def __call__(self, *args, **kwargs):
+        return self.forward(*args, **kwargs)
+    
+
+    def get_output_names(self):
+        output_names = []
+        for node in self.session.get_outputs():
+            output_names.append(node.name)
+        return output_names
+
+
+    def set_device(self, device):
+        if isinstance(device, str):
+            device = torch.device(device)
+        if device.type == 'cuda':
+            device = '{}:{}'.format(device.type, device.index)
+            providers = [("CUDAExecutionProvider", {"device_id": device[-1:] if device[-1] in [str(_i) for _i in range(10)] else "0"}), "CPUExecutionProvider"]
+        else:
+            providers = ["CPUExecutionProvider"]
+        self.session.set_providers(["CUDAExecutionProvider"])
+        self.device = device
+
+
+class Yolo(SimpleOnnxInference):
+    def __init__(self, checkpoint, device='cuda', threshold_conf=0.05, threshold_multi_persons=0.1, input_resolution=(640, 640), threshold_iou=0.5, threshold_bbox_shape_ratio=0.4, cat_id=[1], select_type='max', strict=True, sorted_func=None, **kwargs):
+        super(Yolo, self).__init__(checkpoint, device=device, **kwargs)
+        self.session.set_providers(["CUDAExecutionProvider"])
+        model_inputs = self.session.get_inputs()
+        input_shape = model_inputs[0].shape
+
+        self.input_width = 640
+        self.input_height = 640
+        
+        self.threshold_multi_persons = threshold_multi_persons
+        self.threshold_conf = threshold_conf
+        self.threshold_iou = threshold_iou
+        self.threshold_bbox_shape_ratio = threshold_bbox_shape_ratio
+        self.input_resolution = input_resolution
+        self.cat_id = cat_id
+        self.select_type = select_type
+        self.strict = strict
+        self.sorted_func = sorted_func
+        
+        
+    def preprocess(self, input_image):
+        """
+        Preprocesses the input image before performing inference.
+
+        Returns:
+            image_data: Preprocessed image data ready for inference.
+        """
+        img = read_img(input_image)
+        # Get the height and width of the input image
+        img_height, img_width = img.shape[:2]
+        # Resize the image to match the input shape
+        img = cv2.resize(img, (self.input_resolution[1], self.input_resolution[0]))
+        # Normalize the image data by dividing it by 255.0
+        image_data = np.array(img) / 255.0
+        # Transpose the image to have the channel dimension as the first dimension
+        image_data = np.transpose(image_data, (2, 0, 1))  # Channel first
+        # Expand the dimensions of the image data to match the expected input shape
+        # image_data = np.expand_dims(image_data, axis=0).astype(np.float32)
+        image_data = image_data.astype(np.float32)
+        # Return the preprocessed image data
+        return image_data, np.array([img_height, img_width])
+
+    
+    def postprocess(self, output, shape_raw, cat_id=[1]):
+        """
+        Performs post-processing on the model's output to extract bounding boxes, scores, and class IDs.
+
+        Args:
+            input_image (numpy.ndarray): The input image.
+            output (numpy.ndarray): The output of the model.
+
+        Returns:
+            numpy.ndarray: The input image with detections drawn on it.
+        """
+        # Transpose and squeeze the output to match the expected shape
+
+        outputs = np.squeeze(output)
+        if len(outputs.shape) == 1:
+            outputs = outputs[None]
+        if output.shape[-1] != 6 and output.shape[1] == 84:
+            outputs = np.transpose(outputs)
+        
+        # Get the number of rows in the outputs array
+        rows = outputs.shape[0]
+
+        # Calculate the scaling factors for the bounding box coordinates
+        x_factor = shape_raw[1] / self.input_width
+        y_factor = shape_raw[0] / self.input_height
+
+        # Lists to store the bounding boxes, scores, and class IDs of the detections
+        boxes = []
+        scores = []
+        class_ids = []
+
+        if outputs.shape[-1] == 6:
+            max_scores = outputs[:, 4]
+            classid = outputs[:, -1]
+            
+            threshold_conf_masks = max_scores >= self.threshold_conf
+            classid_masks = classid[threshold_conf_masks] != 3.14159
+
+            max_scores = max_scores[threshold_conf_masks][classid_masks]
+            classid = classid[threshold_conf_masks][classid_masks]
+
+            boxes = outputs[:, :4][threshold_conf_masks][classid_masks]
+            boxes[:, [0, 2]] *= x_factor
+            boxes[:, [1, 3]] *= y_factor
+            boxes[:, 2] = boxes[:, 2] - boxes[:, 0]
+            boxes[:, 3] = boxes[:, 3] - boxes[:, 1]
+            boxes = boxes.astype(np.int32)
+
+        else:
+            classes_scores = outputs[:, 4:]
+            max_scores = np.amax(classes_scores, -1)
+            threshold_conf_masks = max_scores >= self.threshold_conf
+
+            classid = np.argmax(classes_scores[threshold_conf_masks], -1)
+
+            classid_masks = classid!=3.14159
+            
+            classes_scores = classes_scores[threshold_conf_masks][classid_masks]
+            max_scores = max_scores[threshold_conf_masks][classid_masks]
+            classid = classid[classid_masks]
+    
+            xywh = outputs[:, :4][threshold_conf_masks][classid_masks]
+
+            x = xywh[:, 0:1]
+            y = xywh[:, 1:2]
+            w = xywh[:, 2:3]
+            h = xywh[:, 3:4]
+    
+            left = ((x - w / 2) * x_factor)
+            top = ((y - h / 2) * y_factor)
+            width = (w * x_factor)
+            height = (h * y_factor)
+            boxes = np.concatenate([left, top, width, height], axis=-1).astype(np.int32)
+
+        boxes = boxes.tolist()
+        scores = max_scores.tolist()
+        class_ids = classid.tolist()
+
+        # Apply non-maximum suppression to filter out overlapping bounding boxes
+        indices = cv2.dnn.NMSBoxes(boxes, scores, self.threshold_conf, self.threshold_iou)
+        # Iterate over the selected indices after non-maximum suppression
+        
+        results = []
+        for i in indices:
+            # Get the box, score, and class ID corresponding to the index
+            box = box_convert_simple(boxes[i], 'xywh2xyxy')
+            score = scores[i]
+            class_id = class_ids[i]
+            results.append(box + [score] + [class_id])
+            # # Draw the detection on the input image
+
+        # Return the modified input image
+        return np.array(results)
+
+    
+    def process_results(self, results, shape_raw, cat_id=[1], single_person=True):
+        if isinstance(results, tuple):
+            det_results = results[0]
+        else:
+            det_results = results
+
+        person_results = []
+        person_count = 0
+        if len(results):
+            max_idx = -1
+            max_bbox_size = shape_raw[0] * shape_raw[1] * -10
+            max_bbox_shape = -1
+            
+            bboxes = []
+            idx_list = []
+            for i in range(results.shape[0]):
+                bbox = results[i]
+                if (bbox[-1] + 1 in cat_id) and (bbox[-2] > self.threshold_conf):
+                    idx_list.append(i)
+                    bbox_shape = max((bbox[2] - bbox[0]), ((bbox[3] - bbox[1])))
+                    if bbox_shape > max_bbox_shape:
+                        max_bbox_shape = bbox_shape
+            
+            results = results[idx_list]
+
+            for i in range(results.shape[0]):
+                bbox = results[i]
+                bboxes.append(bbox)
+                if self.select_type == 'max':
+                    bbox_size = (bbox[2] - bbox[0]) * ((bbox[3] - bbox[1]))
+                elif self.select_type == 'center':
+                    bbox_size = (abs((bbox[2] + bbox[0]) / 2 - shape_raw[1]/2)) * -1
+                bbox_shape = max((bbox[2] - bbox[0]), ((bbox[3] - bbox[1])))
+                if bbox_size > max_bbox_size:
+                    if (self.strict or max_idx != -1) and bbox_shape < max_bbox_shape * self.threshold_bbox_shape_ratio:
+                        continue
+                    max_bbox_size = bbox_size
+                    max_bbox_shape = bbox_shape
+                    max_idx = i
+
+            if self.sorted_func is not None and len(bboxes) > 0:
+                max_idx = self.sorted_func(bboxes, shape_raw)
+                bbox = bboxes[max_idx]
+                if self.select_type == 'max':
+                    max_bbox_size = (bbox[2] - bbox[0]) * ((bbox[3] - bbox[1]))
+                elif self.select_type == 'center':
+                    max_bbox_size = (abs((bbox[2] + bbox[0]) / 2 - shape_raw[1]/2)) * -1
+                
+            if max_idx != -1:
+                person_count = 1
+
+            if max_idx != -1:
+                person = {}
+                person['bbox'] = results[max_idx, :5]
+                person['track_id'] = int(0)
+                person_results.append(person)
+
+            for i in range(results.shape[0]):
+                bbox = results[i]
+                if (bbox[-1] + 1 in cat_id) and (bbox[-2] > self.threshold_conf):
+                    if self.select_type == 'max':
+                        bbox_size = (bbox[2] - bbox[0]) * ((bbox[3] - bbox[1]))
+                    elif self.select_type == 'center':
+                        bbox_size = (abs((bbox[2] + bbox[0]) / 2 - shape_raw[1]/2)) * -1
+                    if i != max_idx and bbox_size > max_bbox_size * self.threshold_multi_persons and bbox_size < max_bbox_size:
+                        person_count += 1
+                        if not single_person:
+                            person = {}
+                            person['bbox'] = results[i, :5]
+                            person['track_id'] = int(person_count - 1)
+                            person_results.append(person)                   
+            return person_results
+        else:
+            return None
+        
+
+    def postprocess_threading(self, outputs, shape_raw, person_results, i, single_person=True, **kwargs):
+        result = self.postprocess(outputs[i], shape_raw[i], cat_id=self.cat_id)
+        result = self.process_results(result, shape_raw[i], cat_id=self.cat_id, single_person=single_person)
+        if result is not None and len(result) != 0:
+            person_results[i] = result
+
+
+    def forward(self, img, shape_raw, **kwargs):
+        """
+        Performs inference using an ONNX model and returns the output image with drawn detections.
+
+        Returns:
+            output_img: The output image with drawn detections.
+        """
+        if isinstance(img, torch.Tensor):
+            img = img.cpu().numpy()
+            shape_raw = shape_raw.cpu().numpy()
+
+        outputs = self.session.run(None, {self.session.get_inputs()[0].name: img})[0]
+        person_results = [[{'bbox': np.array([0., 0., 1.*shape_raw[i][1], 1.*shape_raw[i][0], -1]), 'track_id': -1}] for i in range(len(outputs))]
+
+        for i in range(len(outputs)):
+            self.postprocess_threading(outputs, shape_raw, person_results, i, **kwargs)         
+        return person_results
+
+
+class ViTPose(SimpleOnnxInference):
+    def __init__(self, checkpoint, device='cuda', **kwargs):
+        super(ViTPose, self).__init__(checkpoint, device=device)
+        self.session.set_providers(["CUDAExecutionProvider"])
+
+    def forward(self, img, center, scale, **kwargs):
+        heatmaps = self.session.run([], {self.session.get_inputs()[0].name: img})[0]
+        points, prob = keypoints_from_heatmaps(heatmaps=heatmaps,
+                                            center=center,
+                                            scale=scale*200,
+                                            unbiased=True, 
+                                            use_udp=False)
+        return np.concatenate([points, prob], axis=2)
+
+
+    @staticmethod
+    def preprocess(img, bbox=None, input_resolution=(256, 192), rescale=1.25, mask=None, **kwargs):
+        if bbox is None or bbox[-1] <= 0 or (bbox[2] - bbox[0]) < 10 or (bbox[3] - bbox[1]) < 10:
+            bbox = np.array([0, 0, img.shape[1], img.shape[0]])
+        
+        bbox_xywh = bbox
+        if mask is not None:
+            img = np.where(mask>128, img, mask)
+
+        if isinstance(input_resolution, int):
+            center, scale = bbox_from_detector(bbox_xywh, (input_resolution, input_resolution), rescale=rescale)
+            img, new_shape, old_xy, new_xy = crop(img, center, scale, (input_resolution, input_resolution))
+        else:
+            center, scale = bbox_from_detector(bbox_xywh, input_resolution, rescale=rescale)
+            img, new_shape, old_xy, new_xy = crop(img, center, scale, (input_resolution[0], input_resolution[1]))
+
+        IMG_NORM_MEAN = np.array([0.485, 0.456, 0.406])
+        IMG_NORM_STD = np.array([0.229, 0.224, 0.225])
+        img_norm = (img / 255. - IMG_NORM_MEAN) / IMG_NORM_STD
+        img_norm = img_norm.transpose(2, 0, 1).astype(np.float32)
+        return img_norm, np.array(center), np.array(scale)
+
+
+class Pose2d:
+    def __init__(self, checkpoint, detector_checkpoint=None, device='cuda', **kwargs):
+
+        if detector_checkpoint is not None:
+            self.detector = Yolo(detector_checkpoint, device)
+        else:
+            self.detector = None
+
+        self.model = ViTPose(checkpoint, device)
+        self.device = device
+
+    def load_images(self, inputs):
+        """
+        Load images from various input types.
+        
+        Args:
+            inputs (Union[str, np.ndarray, List[np.ndarray]]): Input can be file path, 
+                     single image array, or list of image arrays
+            
+        Returns:
+            List[np.ndarray]: List of RGB image arrays
+            
+        Raises:
+            ValueError: If file format is unsupported or image cannot be read
+        """
+        if isinstance(inputs, str):
+            if inputs.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
+                cap = cv2.VideoCapture(inputs)
+                frames = []
+                while True:
+                    ret, frame = cap.read()
+                    if not ret:
+                        break
+                    frames.append(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+                cap.release()
+                images = frames
+            elif inputs.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp')):
+                img = cv2.cvtColor(cv2.imread(inputs), cv2.COLOR_BGR2RGB)
+                if img is None:
+                    raise ValueError(f"Cannot read image: {inputs}")
+                images = [img]
+            else:
+                raise ValueError(f"Unsupported file format: {inputs}")
+                
+        elif isinstance(inputs, np.ndarray):
+            images = [cv2.cvtColor(image, cv2.COLOR_BGR2RGB) for image in inputs]
+        elif isinstance(inputs, list):
+            images = [cv2.cvtColor(image, cv2.COLOR_BGR2RGB) for image in inputs]
+        return images
+
+    def __call__(
+        self, 
+        inputs: Union[str, np.ndarray, List[np.ndarray]],
+        return_image: bool = False,
+        **kwargs
+    ):
+        """
+        Process input and estimate 2D keypoints.
+        
+        Args:
+            inputs (Union[str, np.ndarray, List[np.ndarray]]): Input can be file path,
+                     single image array, or list of image arrays
+            **kwargs: Additional arguments for processing
+            
+        Returns:
+            np.ndarray: Array of detected 2D keypoints for all input images
+        """
+        images = self.load_images(inputs)
+        H, W = images[0].shape[:2]
+        if self.detector is not None:
+            bboxes = []
+            for _image in images:
+                img, shape = self.detector.preprocess(_image)
+                bboxes.append(self.detector(img[None], shape[None])[0][0]["bbox"])
+        else:
+            bboxes = [None] * len(images)
+
+        kp2ds = []
+        for _image, _bbox in zip(images, bboxes):
+            img, center, scale = self.model.preprocess(_image, _bbox)
+            kp2ds.append(self.model(img[None], center[None], scale[None]))
+        kp2ds = np.concatenate(kp2ds, 0)
+        metas = load_pose_metas_from_kp2ds_seq(kp2ds, width=W, height=H)
+        return metas