init

.gitattributes

@@ -30,6 +30,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
 *.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
+*.mp4 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

app.py

@@ -0,0 +1,386 @@
+# app.py
+import gradio as gr
+import subprocess
+import os
+import sys
+import datetime
+import shutil
+import time # Moved import time to the top for global access
+import argparse
+
+# --- Configuration ---
+# !!! IMPORTANT: Ensure this path is correct for your environment !!!
+CKPT_DIR = "./checkpoints/Wan2.1-T2V-1.3B"
+EDIT_SCRIPT_PATH = "edit.py"  # Assumes edit.py is in the same directory
+OUTPUT_DIR = "gradio_outputs"
+PYTHON_EXECUTABLE = sys.executable # Uses the same python that runs gradio
+VIDEO_EXAMPLES_DIR = "video_list" # Directory for example videos
+
+# Create output directory if it doesn't exist
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+os.makedirs(VIDEO_EXAMPLES_DIR, exist_ok=True) # Ensure video_list exists for clarity
+
+def _parse_args():
+    parser = argparse.ArgumentParser(
+        description="Generate a image or video from a text prompt or image using Wan"
+    )
+    parser.add_argument(
+        "--ckpt",
+        type=str,
+        default="./checkpoints/Wan2.1-T2V-1.3B",
+        help="The path to the checkpoint directory.")
+
+    return parser.parse_args()
+
+def generate_safe_filename_part(text, max_len=20):
+    """Generates a filesystem-safe string from text."""
+    if not text:
+        return "untitled"
+    safe_text = "".join(c if c.isalnum() or c in [' ', '_'] else '_' for c in text).strip()
+    safe_text = "_".join(safe_text.split()) # Replace spaces with underscores
+    return safe_text[:max_len]
+
+def run_video_edit(source_video_path, source_prompt, target_prompt, source_words, target_words, 
+                   omega_value, n_max_value, n_avg_value, progress=gr.Progress(track_tqdm=True)):
+    if not source_video_path:
+        raise gr.Error("Please upload a source video.")
+    if not source_prompt:
+        raise gr.Error("Please provide a source prompt.")
+    if not target_prompt:
+        raise gr.Error("Please provide a target prompt (the 'prompt' for edit.py).")
+    # Allow empty source_words for additive edits
+    if source_words is None: # Check for None, as empty string is valid
+         raise gr.Error("Please provide source words (can be empty string for additions).")
+    if not target_words:
+        raise gr.Error("Please provide target words.")
+
+    progress(0, desc="Preparing for video editing...")
+    print(f"Source video received at: {source_video_path}")
+    print(f"Omega value: {omega_value}")
+    print(f"N_max value: {n_max_value}")
+    print(f"N_avg value: {n_avg_value}")
+
+    worse_avg_value = n_avg_value // 2
+    print(f"Calculated Worse_avg value: {worse_avg_value}")
+
+    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+    src_words_fn = generate_safe_filename_part(source_words)
+    tar_words_fn = generate_safe_filename_part(target_words)
+    
+    output_filename_base = f"{timestamp}_{src_words_fn}_to_{tar_words_fn}_omega{omega_value}_nmax{n_max_value}_navg{n_avg_value}"
+    output_video_path = os.path.join(OUTPUT_DIR, f"{output_filename_base}.mp4")
+
+    cmd = [
+        PYTHON_EXECUTABLE, EDIT_SCRIPT_PATH,
+        "--task", "t2v-1.3B",
+        "--size", "832*480",
+        "--base_seed", "42",
+        "--ckpt_dir", CKPT_DIR,
+        "--sample_solver", "unipc",
+        "--source_video_path", source_video_path,
+        "--source_prompt", source_prompt,
+        "--source_words", source_words, # Pass as is, even if empty
+        "--prompt", target_prompt,
+        "--target_words", target_words,
+        "--sample_guide_scale", "3.5",
+        "--tar_guide_scale", "10.5",
+        "--sample_shift", "12",
+        "--sample_steps", "50",
+        "--n_max", str(n_max_value),
+        "--n_min", "0", 
+        "--n_avg", str(n_avg_value),
+        "--worse_avg", str(worse_avg_value), 
+        "--omega", str(omega_value),
+        "--window_size", "11",
+        "--decay_factor", "0.25",
+        "--frame_num", "41",
+        "--save_file", output_video_path
+    ]
+
+    print(f"Executing command: {' '.join(cmd)}")
+    progress(0.1, desc="Starting video editing process...")
+
+    try:
+        process = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, bufsize=1, universal_newlines=True)
+        
+        # Simulate progress
+        for i in range(10): 
+            if process.poll() is not None: 
+                break
+            progress(0.1 + i * 0.08, desc=f"Editing in progress... (simulated step {i+1}/10)") 
+            time.sleep(1) 
+
+        stdout, stderr = process.communicate() 
+
+        progress(0.9, desc="Finalizing video...")
+
+        if process.returncode != 0:
+            print(f"Error during video editing:\nStdout:\n{stdout}\nStderr:\n{stderr}")
+            raise gr.Error(f"Video editing failed. Stderr: {stderr[:500]}")
+        
+        print(f"Video editing successful. Output at: {output_video_path}")
+        if not os.path.exists(output_video_path):
+            print(f"Error: Output file {output_video_path} was not created.")
+            raise gr.Error(f"Output file not found, though script reported success. Stdout: {stdout}")
+        
+        progress(1, desc="Video ready!")
+        return output_video_path
+
+    except FileNotFoundError:
+        progress(1, desc="Error")
+        print(f"Error: The script '{EDIT_SCRIPT_PATH}' or python executable '{PYTHON_EXECUTABLE}' was not found.")
+        raise gr.Error(f"Execution error: Ensure '{EDIT_SCRIPT_PATH}' and Python are correctly pathed.")
+    except Exception as e:
+        progress(1, desc="Error")
+        print(f"An unexpected error occurred: {e}")
+        raise gr.Error(f"An unexpected error: {str(e)}")
+
+# --- Gradio UI Definition ---
+
+# Define all examples to be loaded
+examples_to_load_definitions = [
+    { # Original bear_g example (corresponds to bear_g_03 in YAML)
+        "video_base_name": "bear_g",
+        "src_prompt": "A large brown bear is walking slowly across a rocky terrain in a zoo enclosure, surrounded by stone walls and scattered greenery. The camera remains fixed, capturing the bear's deliberate movements.",
+        "tar_prompt": "A large dinosaur is walking slowly across a rocky terrain in a zoo enclosure, surrounded by stone walls and scattered greenery. The camera remains fixed, capturing the dinosaur's deliberate movements.",
+        "src_words": "large brown bear",
+        "tar_words": "large dinosaur",
+    },
+    { # blackswan_02
+        "video_base_name": "blackswan",
+        "src_prompt": "A black swan with a red beak swimming in a river near a wall and bushes.",
+        "tar_prompt": "A white duck with a red beak swimming in a river near a wall and bushes.",
+        "src_words": "black swan",
+        "tar_words": "white duck",
+    },
+    { # jeep_01
+        "video_base_name": "jeep",
+        "src_prompt": "A silver jeep driving down a curvy road in the countryside.",
+        "tar_prompt": "A Porsche car driving down a curvy road in the countryside.",
+        "src_words": "silver jeep",
+        "tar_words": "Porsche car",
+    },
+    { # woman_02 (additive edit)
+        "video_base_name": "woman",
+        "src_prompt": "A woman in a black dress is walking along a paved path in a lush green park, with trees and a wooden bench in the background. The camera remains fixed, capturing her steady movement.",
+        "tar_prompt": "A woman in a black dress and a red baseball cap is walking along a paved path in a lush green park, with trees and a wooden bench in the background. The camera remains fixed, capturing her steady movement.",
+        "src_words": "", # Empty source words for addition
+        "tar_words": "a red baseball cap",
+    }
+]
+
+examples_data = []
+# Default advanced parameters for all examples
+default_omega = 2.75
+default_n_max = 40
+default_n_avg = 4
+
+for ex_def in examples_to_load_definitions:
+    # Assuming .mp4 extension for all videos
+    video_file_name = f"{ex_def['video_base_name']}.mp4" 
+    example_video_path = os.path.join(VIDEO_EXAMPLES_DIR, video_file_name)
+    
+    if os.path.exists(example_video_path):
+        examples_data.append([
+            example_video_path,
+            ex_def["src_prompt"],
+            ex_def["tar_prompt"],
+            ex_def["src_words"],
+            ex_def["tar_words"],
+            default_omega,
+            default_n_max,
+            default_n_avg
+        ])
+    else:
+        print(f"Warning: Example video {example_video_path} not found. Example for '{ex_def['video_base_name']}' will be skipped.")
+
+if not examples_data:
+    print(f"Warning: No example videos found in '{VIDEO_EXAMPLES_DIR}'. Examples section will be empty or not show.")
+
+
+with gr.Blocks(theme=gr.themes.Soft(), css="""
+    /* Main container - maximize width and improve spacing */
+    .gradio-container {
+        max-width: 98% !important; 
+        width: 98% !important; 
+        margin: 0 auto !important; 
+        padding: 20px !important;
+        min-height: 100vh !important;
+    }
+    
+    /* All containers should use full width */
+    .contain, .container {
+        max-width: 100% !important;
+        width: 100% !important;
+        padding: 0 !important;
+    }
+    
+    /* Remove default padding from main wrapper */
+    .main, .wrap, .panel {
+        max-width: 100% !important;
+        width: 100% !important;
+        padding: 0 !important;
+    }
+    
+    /* Improve spacing for components */
+    .gap, .form {
+        gap: 15px !important;
+    }
+    
+    /* Make all components full width */
+    #component-0, .block {
+        max-width: 100% !important;
+        width: 100% !important;
+    }
+    
+    /* Better padding for groups */
+    .group {
+        padding: 20px !important;
+        margin-bottom: 15px !important;
+        border-radius: 8px !important;
+    }
+    
+    /* Make rows and columns use full space with better gaps */
+    .row {
+        gap: 30px !important;
+        margin-bottom: 20px !important;
+    }
+    
+    /* Improve column spacing */
+    .column {
+        padding: 0 10px !important;
+    }
+    
+    /* Better video component sizing */
+    .video-container {
+        width: 100% !important;
+    }
+    
+    /* Textbox improvements */
+    .textbox, .input-field {
+        width: 100% !important;
+    }
+    
+    /* Button styling */
+    .primary {
+        width: 100% !important;
+        padding: 12px !important;
+        font-size: 16px !important;
+        margin-top: 20px !important;
+    }
+    
+    /* Examples section spacing */
+    .examples {
+        margin-top: 30px !important;
+        padding: 20px !important;
+    }
+    
+    /* Accordion improvements */
+    .accordion {
+        margin: 15px 0 !important;
+    }
+    """) as demo:
+    gr.Markdown(
+        """
+        <h1 style="text-align: center; font-size: 2.5em;">🪄 FlowDirector Video Edit</h1>
+        <p style="text-align: center;">
+        Edit videos by providing a source video, descriptive prompts, and specifying words to change.<br>
+        Powered by FlowDirector.
+        </p>
+        """
+    )
+
+    with gr.Row():
+        with gr.Column(scale=5):  # Input column - increased scale for better space usage
+            with gr.Group():
+                gr.Markdown("### 🎬 Source Material")
+                source_video_input = gr.Video(label="Upload Source Video", height=540)
+                source_prompt_input = gr.Textbox(
+                    label="Source Prompt",
+                    placeholder="Describe the original video content accurately.",
+                    lines=3,
+                    show_label=True
+                )
+                target_prompt_input = gr.Textbox(
+                    label="Target Prompt (Desired Edit)",
+                    placeholder="Describe how you want the video to be after editing.",
+                    lines=3,
+                    show_label=True
+                )
+            
+            with gr.Group():
+                gr.Markdown("### ✍️ Editing Instructions")
+                source_words_input = gr.Textbox(
+                    label="Source Words (to be replaced, or empty for addition)",
+                    placeholder="e.g., large brown bear (leave empty to add target words globally)"
+                )
+                target_words_input = gr.Textbox(
+                    label="Target Words (replacement or addition)",
+                    placeholder="e.g., large dinosaur OR a red baseball cap"
+                )
+
+            with gr.Accordion("🔧 Advanced Parameters", open=False):
+                omega_slider = gr.Slider(
+                    minimum=0.0, maximum=5.0, step=0.05, value=default_omega, label="Omega (ω)",
+                    info="Controls the intensity/style of the edit. Higher values might lead to stronger edits."
+                )
+                n_max_slider = gr.Slider(
+                    minimum=0, maximum=50, step=1, value=default_n_max, label="N_max",
+                    info="Max value for an adaptive param. `n_min` is fixed at 0."
+                )
+                n_avg_slider = gr.Slider(
+                    minimum=0, maximum=5, step=1, value=default_n_avg, label="N_avg",
+                    info="Average value for an adaptive param. `worse_avg` will be N_avg // 2."
+                )
+
+            submit_button = gr.Button("✨ Generate Edited Video", variant="primary")
+
+        with gr.Column(scale=4):  # Output column - increased scale for better proportion
+            gr.Markdown("### 🖼️ Edited Video Output")
+            output_video = gr.Video(label="Result", height=540, show_label=False)
+
+
+    if examples_data: # Only show examples if some were successfully loaded
+        gr.Examples(
+            examples=examples_data,
+            inputs=[
+                source_video_input,
+                source_prompt_input,
+                target_prompt_input,
+                source_words_input,
+                target_words_input,
+                omega_slider,
+                n_max_slider, 
+                n_avg_slider  
+            ],
+            outputs=output_video,
+            fn=run_video_edit,
+            cache_examples=False # For long processes, False is better
+        )
+
+    all_process_inputs = [
+        source_video_input,
+        source_prompt_input,
+        target_prompt_input,
+        source_words_input,
+        target_words_input,
+        omega_slider,
+        n_max_slider, 
+        n_avg_slider  
+    ]
+
+    submit_button.click(
+        fn=run_video_edit,
+        inputs=all_process_inputs,
+        outputs=output_video
+    )
+
+if __name__ == "__main__":
+    # print(f"Make sure your checkpoint directory is correctly set to: {CKPT_DIR}")
+    # print(f"And that '{EDIT_SCRIPT_PATH}' is in the same directory as app.py or correctly pathed.")
+    # print(f"Outputs will be saved to: {os.path.abspath(OUTPUT_DIR)}")
+    # print(f"Place example videos (e.g., bear_g.mp4, blackswan.mp4, etc.) in: {os.path.abspath(VIDEO_EXAMPLES_DIR)}")
+    
+    args = _parse_args()
+    CKPT_DIR = args.ckpt
+    demo.launch()

edit.py

@@ -0,0 +1,486 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import argparse
+from datetime import datetime
+import logging
+import os
+import sys
+import warnings
+
+warnings.filterwarnings('ignore')
+
+import torch, random
+import torch.distributed as dist
+from PIL import Image
+
+import wan
+from wan.configs import WAN_CONFIGS, SIZE_CONFIGS, MAX_AREA_CONFIGS, SUPPORTED_SIZES
+from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander
+from wan.utils.utils import cache_video, cache_image, str2bool
+
+EXAMPLE_PROMPT = {
+    "t2v-1.3B": {
+        "prompt": "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
+    },
+    "t2v-14B": {
+        "prompt": "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
+    },
+    "t2i-14B": {
+        "prompt": "一个朴素端庄的美人",
+    },
+    "i2v-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",
+    },
+}
+
+
+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}"
+
+    # The default sampling steps are 40 for image-to-video tasks and 50 for text-to-video tasks.
+    if args.sample_steps is None:
+        args.sample_steps = 40 if "i2v" in args.task else 50
+
+    if args.sample_shift is None:
+        args.sample_shift = 5.0
+        if "i2v" in args.task and args.size in ["832*480", "480*832"]:
+            args.sample_shift = 3.0
+
+    # The default number of frames are 1 for text-to-image tasks and 81 for other tasks.
+    if args.frame_num is None:
+        args.frame_num = 1 if "t2i" in args.task else 81
+
+    # T2I frame_num check
+    if "t2i" in args.task:
+        assert args.frame_num == 1, f"Unsupport frame_num {args.frame_num} for task {args.task}"
+
+    args.base_seed = args.base_seed if args.base_seed >= 0 else random.randint(
+        0, sys.maxsize)
+    # Size check
+    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])}"
+
+
+def _parse_args():
+    parser = argparse.ArgumentParser(
+        description="Generate a image or video from a text prompt or image using Wan"
+    )
+    parser.add_argument(
+        "--task",
+        type=str,
+        default="t2v-14B",
+        choices=list(WAN_CONFIGS.keys()),
+        help="The task to run.")
+    parser.add_argument(
+        "--size",
+        type=str,
+        default="1280*720",
+        choices=list(SIZE_CONFIGS.keys()),
+        help="The area (width*height) of the generated video. For the I2V task, the aspect ratio of the output video will follow that of the input image."
+    )
+    parser.add_argument(
+        "--frame_num",
+        type=int,
+        default=None,
+        help="How many frames to sample from a image or video. The number should be 4n+1"
+    )
+    parser.add_argument(
+        "--ckpt_dir",
+        type=str,
+        default=None,
+        help="The path to the checkpoint directory.")
+    parser.add_argument(
+        "--offload_model",
+        type=str2bool,
+        default=None,
+        help="Whether to offload the model to CPU after each model forward, reducing GPU memory usage."
+    )
+    parser.add_argument(
+        "--ulysses_size",
+        type=int,
+        default=1,
+        help="The size of the ulysses parallelism in DiT.")
+    parser.add_argument(
+        "--ring_size",
+        type=int,
+        default=1,
+        help="The size of the ring attention parallelism in DiT.")
+    parser.add_argument(
+        "--t5_fsdp",
+        action="store_true",
+        default=False,
+        help="Whether to use FSDP for T5.")
+    parser.add_argument(
+        "--t5_cpu",
+        action="store_true",
+        default=False,
+        help="Whether to place T5 model on CPU.")
+    parser.add_argument(
+        "--dit_fsdp",
+        action="store_true",
+        default=False,
+        help="Whether to use FSDP for DiT.")
+    parser.add_argument(
+        "--save_file",
+        type=str,
+        default=None,
+        help="The file to save the generated image or video to.")
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default=None,
+        help="The prompt to generate the image or video from.")
+    parser.add_argument(
+        "--use_prompt_extend",
+        action="store_true",
+        default=False,
+        help="Whether to use prompt extend.")
+    parser.add_argument(
+        "--prompt_extend_method",
+        type=str,
+        default="local_qwen",
+        choices=["dashscope", "local_qwen"],
+        help="The prompt extend method to use.")
+    parser.add_argument(
+        "--prompt_extend_model",
+        type=str,
+        default=None,
+        help="The prompt extend model to use.")
+    parser.add_argument(
+        "--prompt_extend_target_lang",
+        type=str,
+        default="ch",
+        choices=["ch", "en"],
+        help="The target language of prompt extend.")
+    parser.add_argument(
+        "--base_seed",
+        type=int,
+        default=-1,
+        help="The seed to use for generating the image or video.")
+    parser.add_argument(
+        "--image",
+        type=str,
+        default=None,
+        help="The image to generate the video from.")
+    parser.add_argument(
+        "--sample_solver",
+        type=str,
+        default='unipc',
+        choices=['unipc', 'dpm++'],
+        help="The solver used to sample.")
+    parser.add_argument(
+        "--sample_steps", type=int, default=None, help="The sampling steps.")
+    parser.add_argument(
+        "--sample_shift",
+        type=float,
+        default=None,
+        help="Sampling shift factor for flow matching schedulers.")
+    parser.add_argument(
+        "--sample_guide_scale",
+        type=float,
+        default=5.0,
+        help="Classifier free guidance scale.")
+    parser.add_argument(
+        "--tar_guide_scale",
+        type=float,
+        default=10.0,
+        help="Classifier free guidance scale for target video.")
+    parser.add_argument(
+        "--source_video_path", 
+        type=str, 
+        default=None, 
+        help="Path to the source video for editing.")
+    parser.add_argument(
+        "--source_prompt", 
+        type=str, 
+        default=None, 
+        help="Text prompt describing the source video.")
+    parser.add_argument(
+        "--n_max",
+        type=int,
+        default=35,
+        help="Number of steps to start editing, controlling the editing intensity.")
+    parser.add_argument(
+        "--n_min",
+        type=int,
+        default=0,
+        help="Number of steps at the end of editing, using the vector from tar after completion to control the intensity of style transfer.")
+    parser.add_argument(
+        "--n_avg",
+        type=int,
+        default=5,
+        help="number of steps to average")
+    parser.add_argument(
+        "--worse_avg",
+        type=int,
+        default=3,
+        help="number of steps for worse average")
+    parser.add_argument(
+        "--omega",
+        type=float,
+        default=3,
+        help="omega")
+    parser.add_argument(
+        "--source_words",
+        type=str,
+        default=None,
+        help="Object edited in the source prompt.")
+    parser.add_argument(
+        "--target_words",
+        type=str,
+        default=None,
+        help="Object edited in the target prompt.")
+    parser.add_argument(
+        "--window_size",
+        type=int,
+        default=13,
+        help="window size")
+    parser.add_argument(
+        "--decay_factor",
+        type=float,
+        default=0.1,
+        help="Window decay factor")
+
+
+    args = parser.parse_args()
+
+    _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 edit(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)
+
+    if args.offload_model is None:
+        args.offload_model = False if world_size > 1 else True
+        logging.info(
+            f"offload_model is not specified, set to {args.offload_model}.")
+    if world_size > 1:
+        torch.cuda.set_device(local_rank)
+        dist.init_process_group(
+            backend="nccl",
+            init_method="env://",
+            rank=rank,
+            world_size=world_size)
+    else:
+        assert not (
+            args.t5_fsdp or args.dit_fsdp
+        ), f"t5_fsdp and dit_fsdp are not supported in non-distributed environments."
+        assert not (
+            args.ulysses_size > 1 or args.ring_size > 1
+        ), f"context parallel are not supported in non-distributed environments."
+
+    if args.ulysses_size > 1 or args.ring_size > 1:
+        assert args.ulysses_size * args.ring_size == world_size, f"The number of ulysses_size and ring_size should be equal to the world size."
+        from xfuser.core.distributed import (initialize_model_parallel,
+                                             init_distributed_environment)
+        init_distributed_environment(
+            rank=dist.get_rank(), world_size=dist.get_world_size())
+
+        initialize_model_parallel(
+            sequence_parallel_degree=dist.get_world_size(),
+            ring_degree=args.ring_size,
+            ulysses_degree=args.ulysses_size,
+        )
+
+    if args.use_prompt_extend:
+        if args.prompt_extend_method == "dashscope":
+            prompt_expander = DashScopePromptExpander(
+                model_name=args.prompt_extend_model, is_vl="i2v" in args.task)
+        elif args.prompt_extend_method == "local_qwen":
+            prompt_expander = QwenPromptExpander(
+                model_name=args.prompt_extend_model,
+                is_vl="i2v" in args.task,
+                device=rank)
+        else:
+            raise NotImplementedError(
+                f"Unsupport prompt_extend_method: {args.prompt_extend_method}")
+
+    cfg = WAN_CONFIGS[args.task]
+    if args.ulysses_size > 1:
+        assert cfg.num_heads % args.ulysses_size == 0, f"`num_heads` must be divisible by `ulysses_size`."
+
+    logging.info(f"Generation job args: {args}")
+    logging.info(f"Generation model config: {cfg}")
+
+    if dist.is_initialized():
+        base_seed = [args.base_seed] if rank == 0 else [None]
+        dist.broadcast_object_list(base_seed, src=0)
+        args.base_seed = base_seed[0]
+
+    if "t2v" in args.task or "t2i" in args.task:
+        if args.prompt is None:
+            args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
+        logging.info(f"Input prompt: {args.prompt}")
+        if args.use_prompt_extend:
+            logging.info("Extending prompt ...")
+            if rank == 0:
+                prompt_output = prompt_expander(
+                    args.prompt,
+                    tar_lang=args.prompt_extend_target_lang,
+                    seed=args.base_seed)
+                if prompt_output.status == False:
+                    logging.info(
+                        f"Extending prompt failed: {prompt_output.message}")
+                    logging.info("Falling back to original prompt.")
+                    input_prompt = args.prompt
+                else:
+                    input_prompt = prompt_output.prompt
+                input_prompt = [input_prompt]
+            else:
+                input_prompt = [None]
+            if dist.is_initialized():
+                dist.broadcast_object_list(input_prompt, src=0)
+            args.prompt = input_prompt[0]
+            logging.info(f"Extended prompt: {args.prompt}")
+
+        logging.info("Creating WanT2V pipeline.")
+        wan_t2v = wan.WanT2V(
+            config=cfg,
+            checkpoint_dir=args.ckpt_dir,
+            device_id=device,
+            rank=rank,
+            t5_fsdp=args.t5_fsdp,
+            dit_fsdp=args.dit_fsdp,
+            # use_usp=False,
+            use_usp=(args.ulysses_size > 1 or args.ring_size > 1),
+            t5_cpu=args.t5_cpu,
+        )
+
+        logging.info(
+            f"Generating {'image' if 't2i' in args.task else 'video'} ...")
+        video = wan_t2v.edit(
+            args.prompt,
+            size=SIZE_CONFIGS[args.size],
+            frame_num=args.frame_num,
+            shift=args.sample_shift,
+            sample_solver=args.sample_solver,
+            sampling_steps=args.sample_steps,
+            guide_scale=args.sample_guide_scale,
+            tar_guide_scale=args.tar_guide_scale,
+            seed=args.base_seed,
+            offload_model=args.offload_model,
+            source_video_path=args.source_video_path,
+            source_prompt=args.source_prompt,
+            nmax_step=args.n_max,
+            nmin_step=args.n_min,
+            n_avg=args.n_avg,
+            worse_avg=args.worse_avg,
+            omega=args.omega,
+            source_words=args.source_words,
+            target_words=args.target_words,
+            window_size=args.window_size,
+            decay_factor=args.decay_factor)
+
+    else:
+        if args.prompt is None:
+            args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
+        if args.image is None:
+            args.image = EXAMPLE_PROMPT[args.task]["image"]
+        logging.info(f"Input prompt: {args.prompt}")
+        logging.info(f"Input image: {args.image}")
+
+        img = Image.open(args.image).convert("RGB")
+        if args.use_prompt_extend:
+            logging.info("Extending prompt ...")
+            if rank == 0:
+                prompt_output = prompt_expander(
+                    args.prompt,
+                    tar_lang=args.prompt_extend_target_lang,
+                    image=img,
+                    seed=args.base_seed)
+                if prompt_output.status == False:
+                    logging.info(
+                        f"Extending prompt failed: {prompt_output.message}")
+                    logging.info("Falling back to original prompt.")
+                    input_prompt = args.prompt
+                else:
+                    input_prompt = prompt_output.prompt
+                input_prompt = [input_prompt]
+            else:
+                input_prompt = [None]
+            if dist.is_initialized():
+                dist.broadcast_object_list(input_prompt, src=0)
+            args.prompt = input_prompt[0]
+            logging.info(f"Extended prompt: {args.prompt}")
+
+        logging.info("Creating WanI2V pipeline.")
+        wan_i2v = wan.WanI2V(
+            config=cfg,
+            checkpoint_dir=args.ckpt_dir,
+            device_id=device,
+            rank=rank,
+            t5_fsdp=args.t5_fsdp,
+            dit_fsdp=args.dit_fsdp,
+            use_usp=(args.ulysses_size > 1 or args.ring_size > 1),
+            t5_cpu=args.t5_cpu,
+        )
+
+        logging.info("Generating video ...")
+        video = wan_i2v.edit(
+            args.prompt,
+            img,
+            max_area=MAX_AREA_CONFIGS[args.size],
+            frame_num=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:
+        if args.save_file is None:
+            formatted_time = datetime.now().strftime("%m%d_%H%M")
+            formatted_prompt = args.prompt.replace(" ", "_").replace("/",
+                                                                     "_")[:30]
+            suffix = '.png' if "t2i" in args.task else '.mp4'
+            args.save_file = f"videos/{formatted_time}_{args.source_words.replace(' ', '_').replace('/', '_')}_{args.target_words.replace(' ', '_').replace('/', '_')}_n{args.n_avg}_w{args.worse_avg}_omega{args.omega}_s{args.base_seed}" + suffix
+
+        if "t2i" in args.task:
+            logging.info(f"Saving generated image to {args.save_file}")
+            cache_image(
+                tensor=video.squeeze(1)[None],
+                save_file=args.save_file,
+                nrow=1,
+                normalize=True,
+                value_range=(-1, 1))
+        else:
+            logging.info(f"Saving generated video to {args.save_file}")
+            cache_video(
+                tensor=video[None],
+                save_file=args.save_file,
+                fps=cfg.sample_fps,
+                nrow=1,
+                normalize=True,
+                value_range=(-1, 1))
+    logging.info("Finished.")
+
+
+
+if __name__ == "__main__":
+    args = _parse_args()
+    edit(args)

generate.py

@@ -0,0 +1,412 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import argparse
+from datetime import datetime
+import logging
+import os
+import sys
+import warnings
+
+warnings.filterwarnings('ignore')
+
+import torch, random
+import torch.distributed as dist
+from PIL import Image
+
+import wan
+from wan.configs import WAN_CONFIGS, SIZE_CONFIGS, MAX_AREA_CONFIGS, SUPPORTED_SIZES
+from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander
+from wan.utils.utils import cache_video, cache_image, str2bool
+
+EXAMPLE_PROMPT = {
+    "t2v-1.3B": {
+        "prompt": "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
+    },
+    "t2v-14B": {
+        "prompt": "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
+    },
+    "t2i-14B": {
+        "prompt": "一个朴素端庄的美人",
+    },
+    "i2v-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",
+    },
+}
+
+
+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}"
+
+    # The default sampling steps are 40 for image-to-video tasks and 50 for text-to-video tasks.
+    if args.sample_steps is None:
+        args.sample_steps = 40 if "i2v" in args.task else 50
+
+    if args.sample_shift is None:
+        args.sample_shift = 5.0
+        if "i2v" in args.task and args.size in ["832*480", "480*832"]:
+            args.sample_shift = 3.0
+
+    # The default number of frames are 1 for text-to-image tasks and 81 for other tasks.
+    if args.frame_num is None:
+        args.frame_num = 1 if "t2i" in args.task else 81
+
+    # T2I frame_num check
+    if "t2i" in args.task:
+        assert args.frame_num == 1, f"Unsupport frame_num {args.frame_num} for task {args.task}"
+
+    args.base_seed = args.base_seed if args.base_seed >= 0 else random.randint(
+        0, sys.maxsize)
+    # Size check
+    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])}"
+
+
+def _parse_args():
+    parser = argparse.ArgumentParser(
+        description="Generate a image or video from a text prompt or image using Wan"
+    )
+    parser.add_argument(
+        "--task",
+        type=str,
+        default="t2v-14B",
+        choices=list(WAN_CONFIGS.keys()),
+        help="The task to run.")
+    parser.add_argument(
+        "--size",
+        type=str,
+        default="1280*720",
+        choices=list(SIZE_CONFIGS.keys()),
+        help="The area (width*height) of the generated video. For the I2V task, the aspect ratio of the output video will follow that of the input image."
+    )
+    parser.add_argument(
+        "--frame_num",
+        type=int,
+        default=None,
+        help="How many frames to sample from a image or video. The number should be 4n+1"
+    )
+    parser.add_argument(
+        "--ckpt_dir",
+        type=str,
+        default=None,
+        help="The path to the checkpoint directory.")
+    parser.add_argument(
+        "--offload_model",
+        type=str2bool,
+        default=None,
+        help="Whether to offload the model to CPU after each model forward, reducing GPU memory usage."
+    )
+    parser.add_argument(
+        "--ulysses_size",
+        type=int,
+        default=1,
+        help="The size of the ulysses parallelism in DiT.")
+    parser.add_argument(
+        "--ring_size",
+        type=int,
+        default=1,
+        help="The size of the ring attention parallelism in DiT.")
+    parser.add_argument(
+        "--t5_fsdp",
+        action="store_true",
+        default=False,
+        help="Whether to use FSDP for T5.")
+    parser.add_argument(
+        "--t5_cpu",
+        action="store_true",
+        default=False,
+        help="Whether to place T5 model on CPU.")
+    parser.add_argument(
+        "--dit_fsdp",
+        action="store_true",
+        default=False,
+        help="Whether to use FSDP for DiT.")
+    parser.add_argument(
+        "--save_file",
+        type=str,
+        default=None,
+        help="The file to save the generated image or video to.")
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default=None,
+        help="The prompt to generate the image or video from.")
+    parser.add_argument(
+        "--use_prompt_extend",
+        action="store_true",
+        default=False,
+        help="Whether to use prompt extend.")
+    parser.add_argument(
+        "--prompt_extend_method",
+        type=str,
+        default="local_qwen",
+        choices=["dashscope", "local_qwen"],
+        help="The prompt extend method to use.")
+    parser.add_argument(
+        "--prompt_extend_model",
+        type=str,
+        default=None,
+        help="The prompt extend model to use.")
+    parser.add_argument(
+        "--prompt_extend_target_lang",
+        type=str,
+        default="ch",
+        choices=["ch", "en"],
+        help="The target language of prompt extend.")
+    parser.add_argument(
+        "--base_seed",
+        type=int,
+        default=-1,
+        help="The seed to use for generating the image or video.")
+    parser.add_argument(
+        "--image",
+        type=str,
+        default=None,
+        help="The image to generate the video from.")
+    parser.add_argument(
+        "--sample_solver",
+        type=str,
+        default='unipc',
+        choices=['unipc', 'dpm++'],
+        help="The solver used to sample.")
+    parser.add_argument(
+        "--sample_steps", type=int, default=None, help="The sampling steps.")
+    parser.add_argument(
+        "--sample_shift",
+        type=float,
+        default=None,
+        help="Sampling shift factor for flow matching schedulers.")
+    parser.add_argument(
+        "--sample_guide_scale",
+        type=float,
+        default=5.0,
+        help="Classifier free guidance scale.")
+
+    args = parser.parse_args()
+
+    _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 generate(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)
+
+    if args.offload_model is None:
+        args.offload_model = False if world_size > 1 else True
+        logging.info(
+            f"offload_model is not specified, set to {args.offload_model}.")
+    if world_size > 1:
+        torch.cuda.set_device(local_rank)
+        dist.init_process_group(
+            backend="nccl",
+            init_method="env://",
+            rank=rank,
+            world_size=world_size)
+    else:
+        assert not (
+            args.t5_fsdp or args.dit_fsdp
+        ), f"t5_fsdp and dit_fsdp are not supported in non-distributed environments."
+        assert not (
+            args.ulysses_size > 1 or args.ring_size > 1
+        ), f"context parallel are not supported in non-distributed environments."
+
+    if args.ulysses_size > 1 or args.ring_size > 1:
+        assert args.ulysses_size * args.ring_size == world_size, f"The number of ulysses_size and ring_size should be equal to the world size."
+        from xfuser.core.distributed import (initialize_model_parallel,
+                                             init_distributed_environment)
+        init_distributed_environment(
+            rank=dist.get_rank(), world_size=dist.get_world_size())
+
+        initialize_model_parallel(
+            sequence_parallel_degree=dist.get_world_size(),
+            ring_degree=args.ring_size,
+            ulysses_degree=args.ulysses_size,
+        )
+
+    if args.use_prompt_extend:
+        if args.prompt_extend_method == "dashscope":
+            prompt_expander = DashScopePromptExpander(
+                model_name=args.prompt_extend_model, is_vl="i2v" in args.task)
+        elif args.prompt_extend_method == "local_qwen":
+            prompt_expander = QwenPromptExpander(
+                model_name=args.prompt_extend_model,
+                is_vl="i2v" in args.task,
+                device=rank)
+        else:
+            raise NotImplementedError(
+                f"Unsupport prompt_extend_method: {args.prompt_extend_method}")
+
+    cfg = WAN_CONFIGS[args.task]
+    if args.ulysses_size > 1:
+        assert cfg.num_heads % args.ulysses_size == 0, f"`num_heads` must be divisible by `ulysses_size`."
+
+    logging.info(f"Generation job args: {args}")
+    logging.info(f"Generation model config: {cfg}")
+
+    if dist.is_initialized():
+        base_seed = [args.base_seed] if rank == 0 else [None]
+        dist.broadcast_object_list(base_seed, src=0)
+        args.base_seed = base_seed[0]
+
+    if "t2v" in args.task or "t2i" in args.task:
+        if args.prompt is None:
+            args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
+        logging.info(f"Input prompt: {args.prompt}")
+        if args.use_prompt_extend:
+            logging.info("Extending prompt ...")
+            if rank == 0:
+                prompt_output = prompt_expander(
+                    args.prompt,
+                    tar_lang=args.prompt_extend_target_lang,
+                    seed=args.base_seed)
+                if prompt_output.status == False:
+                    logging.info(
+                        f"Extending prompt failed: {prompt_output.message}")
+                    logging.info("Falling back to original prompt.")
+                    input_prompt = args.prompt
+                else:
+                    input_prompt = prompt_output.prompt
+                input_prompt = [input_prompt]
+            else:
+                input_prompt = [None]
+            if dist.is_initialized():
+                dist.broadcast_object_list(input_prompt, src=0)
+            args.prompt = input_prompt[0]
+            logging.info(f"Extended prompt: {args.prompt}")
+
+        logging.info("Creating WanT2V pipeline.")
+        wan_t2v = wan.WanT2V(
+            config=cfg,
+            checkpoint_dir=args.ckpt_dir,
+            device_id=device,
+            rank=rank,
+            t5_fsdp=args.t5_fsdp,
+            dit_fsdp=args.dit_fsdp,
+            use_usp=(args.ulysses_size > 1 or args.ring_size > 1),
+            t5_cpu=args.t5_cpu,
+        )
+
+        logging.info(
+            f"Generating {'image' if 't2i' in args.task else 'video'} ...")
+        video = wan_t2v.generate(
+            args.prompt,
+            size=SIZE_CONFIGS[args.size],
+            frame_num=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)
+
+    else:
+        if args.prompt is None:
+            args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
+        if args.image is None:
+            args.image = EXAMPLE_PROMPT[args.task]["image"]
+        logging.info(f"Input prompt: {args.prompt}")
+        logging.info(f"Input image: {args.image}")
+
+        img = Image.open(args.image).convert("RGB")
+        if args.use_prompt_extend:
+            logging.info("Extending prompt ...")
+            if rank == 0:
+                prompt_output = prompt_expander(
+                    args.prompt,
+                    tar_lang=args.prompt_extend_target_lang,
+                    image=img,
+                    seed=args.base_seed)
+                if prompt_output.status == False:
+                    logging.info(
+                        f"Extending prompt failed: {prompt_output.message}")
+                    logging.info("Falling back to original prompt.")
+                    input_prompt = args.prompt
+                else:
+                    input_prompt = prompt_output.prompt
+                input_prompt = [input_prompt]
+            else:
+                input_prompt = [None]
+            if dist.is_initialized():
+                dist.broadcast_object_list(input_prompt, src=0)
+            args.prompt = input_prompt[0]
+            logging.info(f"Extended prompt: {args.prompt}")
+
+        logging.info("Creating WanI2V pipeline.")
+        wan_i2v = wan.WanI2V(
+            config=cfg,
+            checkpoint_dir=args.ckpt_dir,
+            device_id=device,
+            rank=rank,
+            t5_fsdp=args.t5_fsdp,
+            dit_fsdp=args.dit_fsdp,
+            use_usp=(args.ulysses_size > 1 or args.ring_size > 1),
+            t5_cpu=args.t5_cpu,
+        )
+
+        logging.info("Generating video ...")
+        video = wan_i2v.generate(
+            args.prompt,
+            img,
+            max_area=MAX_AREA_CONFIGS[args.size],
+            frame_num=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:
+        if args.save_file is None:
+            formatted_time = datetime.now().strftime("%Y%m%d_%H%M%S")
+            formatted_prompt = args.prompt.replace(" ", "_").replace("/",
+                                                                     "_")[:50]
+            suffix = '.png' if "t2i" in args.task else '.mp4'
+            args.save_file = f"videos/{args.task}_{args.size}_{args.ulysses_size}_{args.ring_size}_{formatted_prompt}_{formatted_time}" + suffix
+
+        if "t2i" in args.task:
+            logging.info(f"Saving generated image to {args.save_file}")
+            cache_image(
+                tensor=video.squeeze(1)[None],
+                save_file=args.save_file,
+                nrow=1,
+                normalize=True,
+                value_range=(-1, 1))
+        else:
+            logging.info(f"Saving generated video to {args.save_file}")
+            cache_video(
+                tensor=video[None],
+                save_file=args.save_file,
+                fps=cfg.sample_fps,
+                nrow=1,
+                normalize=True,
+                value_range=(-1, 1))
+    logging.info("Finished.")
+
+
+
+if __name__ == "__main__":
+    args = _parse_args()
+    generate(args)

requirements.txt

@@ -0,0 +1,19 @@
+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
+gradio>=5.0.0
+kornia
+scikit-image==0.25.2
+scipy==1.15.2
+xfuser==0.4.3.post3

video_list/bear_g.mp4

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

video_list/blackswan.mp4

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

video_list/cat_box.mp4

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

video_list/cockatiel.mp4

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

video_list/dog_flower_g.mp4

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

video_list/girl_and_dog.mp4

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

video_list/gym_woman.mp4

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

video_list/jeep.mp4

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

video_list/puppy.mp4

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

video_list/rabbit.mp4

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

video_list/sea_lion.mp4

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

video_list/sea_turtle.mp4

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

video_list/wolf.mp4

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

video_list/woman.mp4

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

wan/__init__.py

@@ -0,0 +1,3 @@
+from . import configs, distributed, modules
+from .image2video import WanI2V
+from .text2video import WanT2V

wan/configs/__init__.py

@@ -0,0 +1,42 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import copy
+import os
+
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+from .wan_i2v_14B import i2v_14B
+from .wan_t2v_1_3B import t2v_1_3B
+from .wan_t2v_14B import t2v_14B
+
+# the config of t2i_14B is the same as t2v_14B
+t2i_14B = copy.deepcopy(t2v_14B)
+t2i_14B.__name__ = 'Config: Wan T2I 14B'
+
+WAN_CONFIGS = {
+    't2v-14B': t2v_14B,
+    't2v-1.3B': t2v_1_3B,
+    'i2v-14B': i2v_14B,
+    't2i-14B': t2i_14B,
+}
+
+SIZE_CONFIGS = {
+    '720*1280': (720, 1280),
+    '1280*720': (1280, 720),
+    '480*832': (480, 832),
+    '832*480': (832, 480),
+    '1024*1024': (1024, 1024),
+}
+
+MAX_AREA_CONFIGS = {
+    '720*1280': 720 * 1280,
+    '1280*720': 1280 * 720,
+    '480*832': 480 * 832,
+    '832*480': 832 * 480,
+}
+
+SUPPORTED_SIZES = {
+    't2v-14B': ('720*1280', '1280*720', '480*832', '832*480'),
+    't2v-1.3B': ('480*832', '832*480'),
+    'i2v-14B': ('720*1280', '1280*720', '480*832', '832*480'),
+    't2i-14B': tuple(SIZE_CONFIGS.keys()),
+}

wan/configs/shared_config.py

@@ -0,0 +1,19 @@
+# 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/configs/wan_i2v_14B.py

@@ -0,0 +1,35 @@
+# 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 14B ------------------------#
+
+i2v_14B = EasyDict(__name__='Config: Wan I2V 14B')
+i2v_14B.update(wan_shared_cfg)
+
+i2v_14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+i2v_14B.t5_tokenizer = 'google/umt5-xxl'
+
+# clip
+i2v_14B.clip_model = 'clip_xlm_roberta_vit_h_14'
+i2v_14B.clip_dtype = torch.float16
+i2v_14B.clip_checkpoint = 'models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth'
+i2v_14B.clip_tokenizer = 'xlm-roberta-large'
+
+# vae
+i2v_14B.vae_checkpoint = 'Wan2.1_VAE.pth'
+i2v_14B.vae_stride = (4, 8, 8)
+
+# transformer
+i2v_14B.patch_size = (1, 2, 2)
+i2v_14B.dim = 5120
+i2v_14B.ffn_dim = 13824
+i2v_14B.freq_dim = 256
+i2v_14B.num_heads = 40
+i2v_14B.num_layers = 40
+i2v_14B.window_size = (-1, -1)
+i2v_14B.qk_norm = True
+i2v_14B.cross_attn_norm = True
+i2v_14B.eps = 1e-6

wan/configs/wan_t2v_14B.py

@@ -0,0 +1,29 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from easydict import EasyDict
+
+from .shared_config import wan_shared_cfg
+
+#------------------------ Wan T2V 14B ------------------------#
+
+t2v_14B = EasyDict(__name__='Config: Wan T2V 14B')
+t2v_14B.update(wan_shared_cfg)
+
+# t5
+t2v_14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+t2v_14B.t5_tokenizer = 'google/umt5-xxl'
+
+# vae
+t2v_14B.vae_checkpoint = 'Wan2.1_VAE.pth'
+t2v_14B.vae_stride = (4, 8, 8)
+
+# transformer
+t2v_14B.patch_size = (1, 2, 2)
+t2v_14B.dim = 5120
+t2v_14B.ffn_dim = 13824
+t2v_14B.freq_dim = 256
+t2v_14B.num_heads = 40
+t2v_14B.num_layers = 40
+t2v_14B.window_size = (-1, -1)
+t2v_14B.qk_norm = True
+t2v_14B.cross_attn_norm = True
+t2v_14B.eps = 1e-6

wan/configs/wan_t2v_1_3B.py

@@ -0,0 +1,29 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from easydict import EasyDict
+
+from .shared_config import wan_shared_cfg
+
+#------------------------ Wan T2V 1.3B ------------------------#
+
+t2v_1_3B = EasyDict(__name__='Config: Wan T2V 1.3B')
+t2v_1_3B.update(wan_shared_cfg)
+
+# t5
+t2v_1_3B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+t2v_1_3B.t5_tokenizer = 'google/umt5-xxl'
+
+# vae
+t2v_1_3B.vae_checkpoint = 'Wan2.1_VAE.pth'
+t2v_1_3B.vae_stride = (4, 8, 8)
+
+# transformer
+t2v_1_3B.patch_size = (1, 2, 2)
+t2v_1_3B.dim = 1536
+t2v_1_3B.ffn_dim = 8960
+t2v_1_3B.freq_dim = 256
+t2v_1_3B.num_heads = 12
+t2v_1_3B.num_layers = 30
+t2v_1_3B.window_size = (-1, -1)
+t2v_1_3B.qk_norm = True
+t2v_1_3B.cross_attn_norm = True
+t2v_1_3B.eps = 1e-6

wan/distributed/fsdp.py

@@ -0,0 +1,32 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+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
+
+
+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,
+):
+    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)
+    return model

wan/distributed/xdit_context_parallel.py

@@ -0,0 +1,420 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from time import time
+import torch
+import torch.cuda.amp as amp
+from xfuser.core.distributed import (get_sequence_parallel_rank,
+                                     get_sequence_parallel_world_size,
+                                     get_sp_group)
+from xfuser.core.long_ctx_attention import xFuserLongContextAttention
+
+from ..modules.model import sinusoidal_embedding_1d
+from typing import List, Union, Optional, Tuple
+import torch.nn.functional as F
+import torch
+
+
+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
+
+
+@amp.autocast(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_sequence_parallel_world_size()
+        sp_rank = get_sequence_parallel_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()
+
+
+@torch.no_grad() # Usually don't need gradients for mask generation
+def generate_attention_mask(
+    attention_map: torch.Tensor,
+    grid_sizes: torch.Tensor,
+    target_x_shape: Tuple[int, int, int, int], # Target shape: (C, T, H, W)
+    batch_index: int = 0,
+    target_word_indices: Union[List[int], slice] = None,
+    head_index: Optional[int] = None, # Process single head or average
+    word_aggregation_method: str = 'mean', # How to combine scores for multiple words
+    upsample_mode_spatial: str = 'nearest', # 'nearest', 'bilinear'
+    upsample_mode_temporal: str = 'nearest', # 'nearest', 'linear'
+    output_dtype: torch.dtype = torch.float32 # or torch.bool for soft mask before threshold
+) -> torch.Tensor:
+    """
+    Generates a binary mask from an attention map based on attention towards target words.
+
+    The mask identifies regions in the video (x) that attend strongly to the specified
+    context words, exceeding a given threshold. The mask has the same dimensions as x.
+
+    Args:
+        attention_map (torch.Tensor): Attention weights [B, Head_num, Lx, Lctx].
+                                      Lx = flattened video tokens (patches),
+                                      Lctx = context tokens (words).
+        target_word_indices (Union[List[int], slice]): Indices or slice for the target
+                                                     word(s) in the Lctx dimension.
+        grid_sizes (torch.Tensor): Patch grid dimensions [B, 3] -> (F, H_patch, W_patch)
+                                   for each batch item, corresponding to Lx.
+                                   F, H_patch, W_patch should be integers.
+        target_x_shape (Tuple[int, int, int, int]): The desired output shape [C, T, H, W],
+                                                    matching the original video tensor x.
+        threshold (float): Value between 0 and 1. Attention scores >= threshold become 1 (True),
+                           otherwise 0 (False).
+        batch_index (int, optional): Batch item to process. Defaults to 0.
+        head_index (Optional[int], optional): Specific head to use. If None, average
+                                              attention across all heads. Defaults to None.
+        word_aggregation_method (str, optional): How to aggregate scores if multiple
+                                                 target_word_indices are given ('mean',
+                                                 'sum', 'max'). Defaults to 'mean'.
+        upsample_mode_spatial (str, optional): PyTorch interpolate mode for H, W dimensions.
+                                           Defaults to 'nearest'.
+        upsample_mode_temporal (str, optional): PyTorch interpolate mode for T dimension.
+                                            Defaults to 'nearest'.
+        output_dtype (torch.dtype, optional): Data type of the output mask.
+                                              Defaults to torch.bool.
+
+    Returns:
+        torch.Tensor: A binary mask tensor of shape target_x_shape [C, T, H, W].
+
+    Raises:
+        TypeError: If inputs are not torch.Tensors.
+        ValueError: If tensor dimensions or indices are invalid, or if
+                    aggregation/upsample modes are unknown.
+        IndexError: If batch_index or head_index are out of bounds.
+    """
+    # --- Input Validation ---
+    if not isinstance(attention_map, torch.Tensor):
+        raise TypeError("attention_map must be a torch.Tensor")
+    if not isinstance(grid_sizes, torch.Tensor):
+        raise TypeError("grid_sizes must be a torch.Tensor")
+    if attention_map.dim() != 4:
+        raise ValueError(f"attention_map must be [B, H, Lx, Lctx], got {attention_map.dim()} dims")
+    if grid_sizes.dim() != 2 or grid_sizes.shape[1] != 3:
+        raise ValueError(f"grid_sizes must be [B, 3], got {grid_sizes.shape}")
+    if len(target_x_shape) != 4:
+         raise ValueError(f"target_x_shape must be [C, T, H, W], got length {len(target_x_shape)}")
+
+    B, H, Lx, Lctx = attention_map.shape
+    C_out, T_out, H_out, W_out = target_x_shape
+
+    if not 0 <= batch_index < B:
+        raise IndexError(f"batch_index {batch_index} out of range for batch size {B}")
+    if head_index is not None and not 0 <= head_index < H:
+         raise IndexError(f"head_index {head_index} out of range for head count {H}")
+    if word_aggregation_method not in ['mean', 'sum', 'max']:
+        raise ValueError(f"Unknown word_aggregation_method: {word_aggregation_method}")
+    if upsample_mode_spatial not in ['nearest', 'bilinear']:
+        raise ValueError(f"Unknown upsample_mode_spatial: {upsample_mode_spatial}")
+    if upsample_mode_temporal not in ['nearest', 'linear']:
+         raise ValueError(f"Unknown upsample_mode_temporal: {upsample_mode_temporal}")
+
+
+    # --- Select Head(s) ---
+    if head_index is None:
+        # Average across heads. Shape -> [Lx, Lctx]
+        attn_map_processed = attention_map[batch_index].mean(dim=0)
+    else:
+        # Select specific head. Shape -> [Lx, Lctx]
+        attn_map_processed = attention_map[batch_index, head_index]
+
+    # --- Select and Aggregate Word Attention ---
+    # Ensure target_word_indices are valid before slicing
+    if isinstance(target_word_indices, slice):
+        _slice_indices = range(*target_word_indices.indices(Lctx))
+        if not _slice_indices: # Empty slice
+             num_words = 0
+        elif _slice_indices.start >= Lctx or _slice_indices.stop < -Lctx : # Basic out of bounds check
+             num_words = len(_slice_indices) # Proceed cautiously or add stricter check
+        else:
+             num_words = len(_slice_indices)
+        word_indices_str = f"slice({_slice_indices.start}:{_slice_indices.stop}:{_slice_indices.step})"
+        word_attn_scores = attn_map_processed[:, target_word_indices] # Shape -> [Lx, num_words]
+    elif isinstance(target_word_indices, list):
+         # Check indices are within bounds
+         valid_indices = [idx for idx in target_word_indices if -Lctx <= idx < Lctx]
+         if not valid_indices:
+             num_words = 0
+             word_attn_scores = torch.empty((Lx, 0), device=attention_map.device, dtype=attention_map.dtype) # Handle empty case
+         else:
+            word_attn_scores = attn_map_processed[:, valid_indices] # Shape -> [Lx, num_words]
+            num_words = len(valid_indices)
+         word_indices_str = str(valid_indices) # Report used indices
+    else:
+        raise TypeError(f"target_word_indices must be list or slice, got {type(target_word_indices)}")
+
+    if num_words > 1:
+        if word_aggregation_method == 'mean':
+            aggregated_scores = word_attn_scores.mean(dim=-1)
+        elif word_aggregation_method == 'sum':
+            aggregated_scores = word_attn_scores.sum(dim=-1)
+        elif word_aggregation_method == 'max':
+            aggregated_scores = word_attn_scores.max(dim=-1).values
+        # aggregated_scores shape -> [Lx]
+    elif num_words == 1:
+         aggregated_scores = word_attn_scores.squeeze(-1) # Shape -> [Lx]
+    else: # No valid words selected
+         return torch.zeros(target_x_shape, dtype=output_dtype, device=attention_map.device)
+
+    # --- Reshape to Video Patch Grid ---
+    # Ensure grid sizes are integers
+    f_patch, h_patch, w_patch = map(int, grid_sizes[batch_index].tolist())
+    actual_num_tokens = f_patch * h_patch * w_patch
+
+    if actual_num_tokens == 0:
+         return torch.zeros(target_x_shape, dtype=output_dtype, device=attention_map.device)
+
+    # Handle mismatch between expected tokens (from grid) and actual attention length (Lx)
+    if actual_num_tokens > Lx:
+         # Pad aggregated_scores to actual_num_tokens size
+         padding_size = actual_num_tokens - aggregated_scores.numel()
+         scores_padded = F.pad(aggregated_scores, (0, padding_size), "constant", 0)
+         scores_unpadded = scores_padded # Use the padded version for reshaping
+         # This scenario is less common than Lx > actual_num_tokens
+    elif actual_num_tokens < Lx:
+        scores_unpadded = aggregated_scores[:actual_num_tokens]
+    else:
+        scores_unpadded = aggregated_scores # Shape [actual_num_tokens]
+
+    try:
+        # Reshape to [F_patch, H_patch, W_patch]
+        attention_patch_grid = scores_unpadded.reshape(f_patch, h_patch, w_patch)
+    except RuntimeError as e:
+        raise e
+
+    # --- Upsample to Original Video Resolution ---
+    # Add batch and channel dims for interpolation: [1, 1, F_patch, H_patch, W_patch]
+    # Note: Assuming attention is channel-agnostic here.
+    grid_for_upsample = attention_patch_grid.unsqueeze(0).unsqueeze(0).float() # Interpolate needs float
+
+ 
+     # --- SIMPLIFIED LOGIC: Always use 3D interpolation ---
+    target_size_3d = (T_out, H_out, W_out)
+
+    # Determine the 3D interpolation mode.
+    # Default to 'nearest' unless temporal dimension changes AND 'linear' is requested.
+    if upsample_mode_temporal == 'linear' and f_patch != T_out:
+        upsample_mode_3d = 'trilinear'
+        align_corners_3d = False # align_corners usually False for non-nearest modes
+    else:
+        # Use 'nearest' if T isn't changing, or if temporal mode is 'nearest'.
+        # 'nearest' is generally safer and handles spatial modes implicitly.
+        upsample_mode_3d = 'nearest'
+        align_corners_3d = None # align_corners=None for nearest
+
+    upsampled_scores_grid = F.interpolate(grid_for_upsample,
+                                       size=target_size_3d,
+                                       mode=upsample_mode_3d,
+                                       align_corners=align_corners_3d)
+    # Expected shape: [1, 1, T_out, H_out, W_out] == [1, 1, 21, 60, 104]
+
+    # --- END SIMPLIFIED LOGIC ---
+ 
+    # Remove batch and channel dims: [T_out, H_out, W_out]
+    upsampled_scores = upsampled_scores_grid.squeeze(0).squeeze(0)
+
+    # --- Thresholding ---
+    binary_mask_thw = (upsampled_scores / torch.max(upsampled_scores)) # Shape [T_out, H_out, W_out]
+
+    # --- Expand Channel Dimension ---
+    # Repeat the mask across the channel dimension C_out
+    # Input shape: [T_out, H_out, W_out]
+    # After unsqueeze(0): [1, T_out, H_out, W_out]
+    # Target shape:      [C_out, T_out, H_out, W_out]
+    # This expand operation is valid as explained above.
+    final_mask = binary_mask_thw.unsqueeze(0).expand(C_out, T_out, H_out, W_out)
+
+    return final_mask.to(dtype=output_dtype)
+
+
+def usp_dit_forward( 
+    self,
+    x,
+    t,
+    context,
+    seq_len,
+    clip_fea=None,
+    y=None,
+    words_indices=None,
+    block_id=-1,
+    type=None,
+    timestep=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 clip_fea is not None and 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
+    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 clip_fea is not None:
+        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,
+        collect_attn_map=False)
+
+    # Context Parallel
+    x = torch.chunk(
+        x, get_sequence_parallel_world_size(),
+        dim=1)[get_sequence_parallel_rank()]
+    
+    save_block_id = block_id 
+    attn_map = None
+    binary_mask = None
+    for i, block in enumerate(self.blocks):
+        kwargs["collect_attn_map"] = False
+        if i == save_block_id:
+            kwargs["collect_attn_map"] = True
+            x, attn_map = block(x, **kwargs)
+        else:
+            x = block(x, **kwargs)
+
+    # head
+    x = self.head(x, e)
+    # Context Parallel
+    x = get_sp_group().all_gather(x, dim=1)
+
+    # unpatchify
+    x = self.unpatchify(x, grid_sizes)
+
+    if save_block_id != -1 and words_indices is not None:
+        attention_map = get_sp_group().all_gather(attn_map, dim=2)
+        binary_mask = generate_attention_mask(
+                        attention_map=attention_map, # [1, 12, 32760, 512] batchsize, head_num, l_x, l_context
+                        target_word_indices=words_indices,
+                        grid_sizes=grid_sizes, # Make sure grid_sizes covers the full batch
+                        target_x_shape=x[0].shape, # channel, frames, h, W
+                        batch_index=0, # Process the first item in the batch
+                        head_index=None, # Average over heads
+                        word_aggregation_method='mean'
+                        )
+
+    return [u.float() for u in x], binary_mask
+
+
+
+def usp_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)
+
+    # TODO: We should use unpaded q,k,v for attention.
+    # k_lens = seq_lens // get_sequence_parallel_world_size()
+    # if k_lens is not None:
+    #     q = torch.cat([u[:l] for u, l in zip(q, k_lens)]).unsqueeze(0)
+    #     k = torch.cat([u[:l] for u, l in zip(k, k_lens)]).unsqueeze(0)
+    #     v = torch.cat([u[:l] for u, l in zip(v, k_lens)]).unsqueeze(0)
+
+    x = xFuserLongContextAttention()(
+        None,
+        query=half(q),
+        key=half(k),
+        value=half(v),
+        window_size=self.window_size)
+
+    # TODO: padding after attention.
+    # x = torch.cat([x, x.new_zeros(b, s - x.size(1), n, d)], dim=1)
+
+    # output
+    x = x.flatten(2)
+    x = self.o(x)
+    return x