Upload folder using huggingface_hub

.gitignore

@@ -0,0 +1,19 @@
+*output*
+*__pycache__*
+*samples*
+*runs*
+*checkpoints*
+master_ip
+*logs*
+*.DS_Store
+.idea
+.ipynb_checkpoints
+*ckpts*
+*kernel_meta*
+*egg*
+*wandb*
+output_root*
+fastvideo
+outputs*
+wandb*
+*test_data*

OSS/OSS.py

@@ -0,0 +1,357 @@
+import torch
+import numpy as np
+import logging
+
+
+from .utils import _broadcast_tensor
+
+def cal_medium(oss_steps_all):
+    ave_steps = []
+    for k in range(len(oss_steps_all[0])):
+        l = []
+        for i in range(len(oss_steps_all)):
+            l.append(oss_steps_all[i][k])
+        l.sort()
+        ave_steps.append((l[len(l)//2] + l[len(l)//2 - 1])//2)
+    return ave_steps
+
+
+
+
+@torch.no_grad()
+def infer_OSS(oss_steps, model, z, class_emb, device, renorm_flag=False, max_amp=None, min_amp=None, float32=True, model_kwargs=None):
+    # z [B,C.H,W]
+
+    N = len(oss_steps)
+    B = z.shape[0]
+    oss_steps = [0] + oss_steps
+    ori_z = z.clone()
+    
+    # First is zero
+    timesteps = model.fm_steps
+    
+    if model_kwargs is None:
+        model_kwargs = {}
+    
+    for i in reversed(range(1,N+1)):
+        logging.info(f"steps {i}")
+        t = torch.ones((B,), device=device, dtype=torch.long) * oss_steps[i]
+
+        if renorm_flag:
+            
+            max_s = torch.quantile(z.reshape((z.shape[0], -1)), 0.95, dim=1)
+            min_s = torch.quantile(z.reshape((z.shape[0], -1)), 0.05, dim=1)
+
+
+            max_amp_tmp = max_amp[oss_steps[i]-1]
+            min_amp_tmp = min_amp[oss_steps[i]-1]
+
+            ak = (max_amp_tmp - min_amp_tmp)/(max_s - min_s)
+            ab = min_amp_tmp - ak*min_s
+            
+            z = _broadcast_tensor(ak, z.shape) *z + _broadcast_tensor(ab, z.shape)
+
+        vt = model(z, t, class_emb, model_kwargs)        
+        if float32:
+            z = z.to(torch.float32)
+        z = z + (timesteps[oss_steps[i-1]] - timesteps[oss_steps[i]]) * vt
+        if float32:
+            z = z.to(ori_z.dtype)
+    
+    return z
+        
+@torch.no_grad()
+def search_OSS_video(model, z, batch_size, class_emb, device, teacher_steps=200, student_steps=5, norm=2, model_kwargs=None, frame_type="6", channel_type="4", random_channel=False, float32=True):
+    # z [B,C.H,W]
+    # model_kwargs doesn't contain class_embedding, which is another seperate input
+    # the class_embedding is the same for all the searching samples here.
+    
+    B = batch_size    
+    N = teacher_steps
+    STEP = student_steps
+    assert z.shape[0] == B
+    channel_size = z.shape[1]
+    frame_size = z.shape[2]
+
+    
+    # First is zero
+    timesteps = model.fm_steps
+    
+    
+    if model_kwargs is None:
+        model_kwargs = {}
+        
+    # Compute the teacher trajectory
+    traj_tea = torch.stack([torch.ones_like(z)]*(N+1), dim = 0)
+    traj_tea[N] = z
+    z_tea = z.clone()
+        
+    for i in reversed(range(1,N+1)):
+        print("teachering,%s"%i)
+
+        t = torch.ones((B,), device=device, dtype=torch.long) * i
+        vt = model(z_tea, t, class_emb, model_kwargs)
+        if float32:
+            z_tea = z_tea.to(torch.float32)
+        z_tea = z_tea + vt * (timesteps[i-1] - timesteps[i])
+        if float32:
+            z_tea = z_tea.to(z.dtype)
+        
+
+        traj_tea[i-1] = z_tea.clone() 
+    
+    
+    # solving dynamic programming
+    all_steps = []
+    
+    
+    for i_batch in range(B): # process each image separately
+        z_cur = z[i_batch].clone().unsqueeze(0)
+        if class_emb is not None:
+            class_emb_cur = class_emb[i_batch].clone().unsqueeze(0)
+        else:
+            class_emb_cur = None
+        
+        tracestep = [torch.ones(N+1, device=device, dtype=torch.long)*N for _ in range(STEP+1)]
+        dp = torch.ones((STEP+1,N+1), device=device)*torch.inf
+        z_prev = torch.cat([z_cur]*(N+1),dim=0) 
+        z_next = z_prev.clone() 
+
+        for k in range(STEP):
+            print("studenting,%s" % k)
+            logging.info(f"Doing k step solving {k}")
+
+            if random_channel and channel_type != "all":
+                channel_select = np.random.choice(range(channel_size), size=int(channel_type), replace=False)
+
+            for i in reversed(range(1,N+1)):
+                z_i = z_prev[i].unsqueeze(0)
+                t = torch.ones((1,), device=device, dtype=torch.long) * i
+                vt = model(z_i, t, class_emb_cur, model_kwargs)
+                
+                for j in reversed(range(0,i)):
+                    if float32:
+                        z_i = z_i.to(torch.float32)
+                    z_nxt = z_i + vt * (timesteps[j] - timesteps[i])
+                    if float32:
+                        z_nxt = z_nxt.to(z.dtype)
+                    
+                    
+                    if random_channel:
+                        pass
+                    elif channel_type == "all":
+                        channel_select = list(range(channel_size))
+                    else:
+                        channel_select = list(range(int(channel_type))) 
+
+                    if frame_type == "all":
+                        frame_select = list(range(frame_size))
+                    else:
+                        frame_select = torch.linspace(0,frame_size-1,int(frame_type),dtype=torch.long).tolist()
+
+                    
+                        
+                    channel_select_tensor = torch.tensor(channel_select, dtype=torch.long, device=device)
+                    frame_select_tensor = torch.tensor(frame_select, dtype=torch.long, device=device)
+                    
+                    z_nxt_select = z_nxt[0, channel_select_tensor,...]
+                    traj_tea_select = traj_tea[j, i_batch, channel_select_tensor, ...]
+                    
+                    z_nxt_select = z_nxt_select[:,frame_select_tensor,... ]
+                    traj_tea_select = traj_tea_select[:,frame_select_tensor,... ]
+                    
+
+                    cost = (torch.abs(z_nxt_select - traj_tea_select))**norm
+                    cost = cost.mean()
+                    
+                    if cost < dp[k][j]:
+                        dp[k][j] = cost
+                        tracestep[k][j] = i
+                        z_next[j] = z_nxt
+                        
+            dp[k+1] = dp[k].clone()
+            tracestep[k+1] = tracestep[k].clone()
+            z_prev = z_next.clone()
+
+
+            logging.info(f"finish {k} steps")
+            
+            cur_step = [0]
+            for kk in reversed(range(k+1)):
+                j = cur_step[-1]
+                cur_step.append(int(tracestep[kk][j].item()))
+
+            logging.info(cur_step)            
+            
+
+        # trace back
+        final_step = [0]
+        for k in reversed(range(STEP)):
+            j = final_step[-1]
+            final_step.append(int(tracestep[k][j].item()))
+        logging.info(final_step)
+        all_steps.append(final_step[1:])
+
+    return all_steps[0]
+
+
+
+@torch.no_grad()
+def search_OSS(model, z, batch_size, class_emb, device, teacher_steps=200, student_steps=5, model_kwargs=None):
+    # z [B,C.H,W]
+    # model_kwargs doesn't contain class_embedding, which is another seperate input
+    # the class_embedding is the same for all the searching samples here.
+    
+    B = batch_size    
+    N = teacher_steps
+    STEP = student_steps
+    assert z.shape[0] == B
+    
+    # First is zero
+    timesteps = model.fm_steps
+    
+    
+    if model_kwargs is None:
+        model_kwargs = {}
+        
+    # Compute the teacher trajectory
+    traj_tea = torch.stack([torch.ones_like(z)]*(N+1), dim = 0)
+    traj_tea[N] = z
+    z_tea = z.clone()
+        
+    for i in reversed(range(1,N+1)):
+
+        t = torch.ones((B,), device=device, dtype=torch.long) * i
+        vt = model(z_tea, t, class_emb, model_kwargs)
+        z_tea = z_tea + vt * (timesteps[i-1] - timesteps[i])
+        traj_tea[i-1] = z_tea.clone() 
+        
+    # solving dynamic programming
+    all_steps = []
+    
+    for i_batch in range(B): # process each image separately
+        z_cur = z[i_batch].clone().unsqueeze(0)
+        if class_emb is not None:
+            class_emb_cur = class_emb[i_batch].clone().unsqueeze(0)
+        else:
+            class_emb_cur = None
+        tracestep = [torch.ones(N+1, device=device, dtype=torch.long)*N for _ in range(STEP+1)]
+        dp = torch.ones((STEP+1,N+1), device=device)*torch.inf
+        z_prev = torch.cat([z_cur]*(N+1),dim=0) 
+        z_next = z_prev.clone() 
+
+
+        for k in range(STEP):
+            logging.info(f"Doing k step solving {k}")
+            for i in reversed(range(1,N+1)):
+                z_i = z_prev[i].unsqueeze(0)
+                t = torch.ones((1,), device=device, dtype=torch.long) * i
+                vt = model(z_i, t, class_emb_cur, model_kwargs)
+                
+                for j in reversed(range(0,i)):
+                    z_j = z_i + vt * (timesteps[j] - timesteps[i])
+                    cost = (z_j - traj_tea[j, i_batch])**2
+                    cost = cost.mean()
+                    
+                    if cost < dp[k][j]:
+                        dp[k][j] = cost
+                        tracestep[k][j] = i
+                        z_next[j] = z_j
+                        
+            dp[k+1] = dp[k].clone()
+            tracestep[k+1] = tracestep[k].clone()
+            z_prev = z_next.clone()
+
+        # trace back
+        final_step = [0]
+        for k in reversed(range(STEP)):
+            j = final_step[-1]
+            final_step.append(int(tracestep[k][j].item()))
+        logging.info(final_step)
+        all_steps.append(final_step[1:])
+
+    return all_steps
+
+
+
+
+@torch.no_grad()
+def search_OSS_batch(model, z, batch_size, class_emb, device, teacher_steps=200, student_steps=5, model_kwargs=None):
+    # z [B,C.H,W]
+    # model_kwargs doesn't contain class_embedding, which is another seperate input
+    # the class_embedding is the same for all the searching samples here.
+    
+    B = batch_size    
+    N = teacher_steps
+    STEP = student_steps
+    
+    
+    # First is zero
+    timesteps = model.fm_steps
+    
+    if model_kwargs is None:
+        model_kwargs = {}
+        
+    # Compute the teacher trajectory
+    traj_tea = torch.stack([torch.ones_like(z)]*(N+1), dim = 0)
+    traj_tea[N] = z
+    z_tea = z.clone()
+
+    for i in reversed(range(1,N+1)):
+        t = torch.ones((B,), device=device, dtype=torch.long) * i
+        vt = model(z_tea, t, class_emb, model_kwargs)
+        z_tea = z_tea + (timesteps[i-1] - timesteps[i]) * vt
+        traj_tea[i-1] = z_tea.clone() 
+
+
+
+    times = torch.linspace(0,N,N+1, device=device).long()
+    t_in = torch.linspace(1, N, N, device=device).long()
+    # solving dynamic programming
+    all_steps = []
+    
+    for i_batch in range(B): # process each image separately
+        z_cur = z[i_batch].clone().unsqueeze(0)
+        if class_emb is not None:
+            class_emb_cur = class_emb[i_batch].clone().unsqueeze(0).expand(N)
+        else:
+            class_emb_cur = None
+        tracestep = [torch.ones(N+1, device=device, dtype=torch.long)*N for _ in range(STEP+1)]
+        dp = torch.ones((STEP+1,N+1), device=device)*torch.inf
+        z_prev = torch.cat([z_cur]*(N+1),dim=0) 
+        z_next = z_prev.clone() 
+
+
+        for k in range(STEP):
+            logging.info(f"Doing k step solving {k}")
+            vt = model(z_prev[1:], t_in, class_emb_cur, model_kwargs) 
+            
+            for i in reversed(range(1,N+1)):
+                t = torch.ones((i,), device=device, dtype=torch.long) * i
+                z_i_batch = torch.stack([z_prev[i]]*i ,dim=0)
+                dt = timesteps[times[:i]] - timesteps[t]
+                z_j_batch = z_i_batch[:i] + _broadcast_tensor(dt, z_i_batch.shape) * vt[i-1].unsqueeze(0)
+
+                cost = (z_j_batch - traj_tea[:i,i_batch,...])**2
+                cost = cost.mean(dim = tuple(range(1, len(cost.shape))))
+                mask = cost < dp[k, :i]
+                
+                dp[k, :i] = torch.where(mask, cost, dp[k, :i])
+                tracestep[k][:i] = torch.where(mask, i, tracestep[k][:i])
+                expanded_mask = _broadcast_tensor(mask,z_i_batch.shape) 
+                z_next[:i] = torch.where(expanded_mask, z_j_batch, z_next[:i])
+                    
+                        
+            dp[k+1] = dp[k].clone()
+            tracestep[k+1] = tracestep[k].clone()
+            z_prev = z_next.clone()
+
+        # trace back
+        final_step = [0]
+        for k in reversed(range(STEP)):
+            j = final_step[-1]
+            final_step.append(int(tracestep[k][j].item()))
+        logging.info(final_step)
+        all_steps.append(final_step[1:])
+    
+    return all_steps

OSS/model_wrap.py

@@ -0,0 +1,124 @@
+import pdb
+
+import torch
+import numpy as np
+from .utils import _broadcast_tensor, _extract_into_tensor
+
+
+
+class _WrappedModel_DiT:
+    def __init__(self, model, diffusion, device=None, class_emb_null=None):
+        self.model = model
+        self.diffusion = diffusion
+        self._predict_xstart_from_eps = diffusion._predict_xstart_from_eps
+        
+        self.diffusion_t_map = list(diffusion.use_timesteps)
+        self.diffusion_t_map.sort()
+        
+        self.diffusion_t = [self.diffusion_t_map[i] for i in range(diffusion.num_timesteps)] # list(range(diffusion.num_timesteps))
+        self.diffusion_t = np.array(self.diffusion_t)
+        
+        self.diffusion_sqrt_alpha_cumprod = np.array([diffusion.sqrt_alphas_cumprod[i] for i in range(diffusion.num_timesteps)])
+        self.fm_steps = [(1 - self.diffusion_sqrt_alpha_cumprod[i]**2)**0.5/(self.diffusion_sqrt_alpha_cumprod[i] + (1 - self.diffusion_sqrt_alpha_cumprod[i]**2)**0.5) for i in range(len(self.diffusion_t))]
+        
+        self.fm_steps = torch.tensor([0] + self.fm_steps, device=device)
+        self.y_null = class_emb_null
+
+
+
+
+        
+    def __call__(self, x, t, y, kwargs):
+
+        N = len(self.diffusion_t)
+        B,C,H,W = x.shape
+        diffusion_x = torch.zeros_like(x)
+        diffusion_t = _extract_into_tensor(self.diffusion_t, t-1, t.shape).long()
+
+
+        t_fm = self.fm_steps[t]
+        diffusion_x_tmp = _extract_into_tensor(self.diffusion.sqrt_alphas_cumprod, t-1, x.shape) * x / ( 1 + 1e-4 - _broadcast_tensor(t_fm,x.shape)) 
+        diffusion_x_tmp = diffusion_x_tmp.to(torch.float)
+        diffusion_x = torch.where(_broadcast_tensor(t,x.shape) == N, x, diffusion_x_tmp)
+
+
+        y_null_batch = torch.cat([self.y_null[0].unsqueeze(0)]*B, dim=0)
+        y_new = torch.cat([y, y_null_batch], 0)
+        
+
+        model_output = self.model(torch.cat([diffusion_x,diffusion_x],dim=0), torch.cat([diffusion_t,diffusion_t],dim=0), y_new, **kwargs)
+        model_output = model_output[:B]
+        model_output, _ = torch.split(model_output, C, dim=1)
+        x0_diffusion = self._predict_xstart_from_eps(x_t=diffusion_x, t=t-1, eps=model_output)
+        vt = (x - x0_diffusion) / (_broadcast_tensor(t_fm,x.shape))
+        vt = vt.to(diffusion_x.dtype)
+        return vt
+
+
+
+class _WrappedModel_Sora:
+    def __init__(self, model, guidance_scale, y_null, timesteps, num_timesteps, mask_t):
+        self.model = model
+        self.guidance_scale = guidance_scale    
+        self.y_null = y_null
+        
+        self.timesteps = [torch.tensor([0], device=model.device)] + timesteps[::-1]
+        self.timesteps = torch.cat(self.timesteps, dim=0)
+        self.fm_steps = [x/num_timesteps for x in self.timesteps]
+        self.mask_t = mask_t
+        
+    def __call__(self, x, t, y, kwargs):
+        y = torch.cat([y, self.y_null], dim=0)
+        
+        t_in = self.timesteps[t]        
+        
+        x_in = torch.cat([x,x], dim=0)
+        # breakpoint()
+        mask_t_upper = self.mask_t >= t_in.unsqueeze(1)
+        kwargs["x_mask"] = mask_t_upper.repeat(2, 1)
+
+        t_in = torch.cat([t_in,t_in], dim=0)
+        with torch.no_grad():
+            pred = self.model(x_in, t_in, y,  **kwargs).chunk(2, dim=1)[0]
+        # breakpoint()
+        pred_cond, pred_uncond = pred.chunk(2, dim=0)
+        v_pred = pred_uncond + self.guidance_scale * (pred_cond - pred_uncond)
+        
+        return -v_pred
+
+
+
+class _WrappedModel_Wan:
+    def __init__(self, model, timesteps, num_timesteps, context_null, guide_scale):
+        self.model = model
+        self.context_null = context_null
+        self.guide_scale = guide_scale
+        fm_steps = torch.cat([timesteps,torch.zeros_like(timesteps[0]).view(1)])
+        self.time_steps = torch.flip(fm_steps, dims=[0])
+        self.fm_steps = self.time_steps/num_timesteps
+        
+        
+    def __call__(self, x, t, y, kwargs):
+        self.time_steps = self.time_steps.to(t.device)
+        t = self.time_steps[t]
+        noise_pred_cond = self.model(x, t=t, context=y, **kwargs)[0]
+        noise_pred_uncond = self.model(x, t=t, context=self.context_null, **kwargs)[0]
+        noise_pred = noise_pred_uncond + self.guide_scale * (noise_pred_cond - noise_pred_uncond)
+        return noise_pred
+
+
+
+
+class _WrappedModel_FLUX:
+    def __init__(self, model, timesteps, num_timesteps):
+        self.model = model
+        fm_steps = torch.cat([timesteps,torch.zeros_like(timesteps[0]).view(1)])
+        self.time_steps = torch.flip(fm_steps, dims=[0])
+        self.fm_steps = self.time_steps/num_timesteps
+            
+        
+    def __call__(self, x, t, y, kwargs):
+        t = self.time_steps[t]
+        t = t.expand(x.shape[0]).to(x.dtype) / 1000
+        pred = self.model(hidden_states=x, timestep=t, **kwargs)[0]
+        return pred

OSS/utils.py

@@ -0,0 +1,22 @@
+
+import torch
+
+
+def _broadcast_tensor(a, broadcast_shape):
+    while len(a.shape) < len(broadcast_shape):
+        a = a[..., None]
+    return a.expand(broadcast_shape)
+
+def _extract_into_tensor(arr, timesteps, broadcast_shape):
+    """
+    Extract values from a 1-D numpy array for a batch of indices.
+    :param arr: the 1-D numpy array.
+    :param timesteps: a tensor of indices into the array to extract.
+    :param broadcast_shape: a larger shape of K dimensions with the batch
+                            dimension equal to the length of timesteps.
+    :return: a tensor of shape [batch_size, 1, ...] where the shape has K dims.
+    """
+    res = torch.from_numpy(arr).to(device=timesteps.device)[timesteps].float()
+    while len(res.shape) < len(broadcast_shape):
+        res = res[..., None]
+    return res + torch.zeros(broadcast_shape, device=timesteps.device)

app.py

@@ -0,0 +1,306 @@
+# -*- coding: utf-8 -*-
+'''
+可以自选帧数、
+    49  3s
+    57  3.5s
+    65  4s
+    73  4.5s
+    81  5s
+加速模式
+    40/tea40/10(暂时不支持)
+
+image2video_teacache1是支持是否传tea参并且能传秒数的版本
+'''
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved
+import logging,os
+os.makedirs("/root/weights",exist_ok=True)
+cmd="huggingface-cli download IndexTeam/AnisoraV3 --include=\"14B/*\" --local-dir=/root/weights --token %s"%os.environ['token']
+os.system(cmd)
+
+# os.environ["HF_ENDPOINT"]          = "https://hf-mirror.com"
+from time import time as ttime
+import argparse
+from datetime import datetime
+import logging
+import sys
+import warnings
+from fastapi import FastAPI
+import uvicorn
+import gradio as gr
+warnings.filterwarnings('ignore')
+
+import torch, random
+import torch.distributed as dist
+from PIL import Image
+
+import wan
+from wan.image2video_if_oss import WanI2V
+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
+
+value2speed={
+    "原版":0,
+    "加速版":1,
+}
+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}.")
+
+    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]
+
+    logging.info("Creating WanI2V pipeline.")
+    # wan_i2v = wan.WanI2V(
+    wan_i2v = 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,
+    )
+
+    def generate_i2v(prompt,img,seed,nf,speed):
+        logging.info("Generating video ...")
+        save_file="output/%s-%s-%s-%s.mp4"%(seed,nf,speed,int(ttime()))
+        video = wan_i2v.generate(
+            prompt,
+            img,
+            max_area=MAX_AREA_CONFIGS[args.size],
+            frame_num=int(nf)*16+1,#args.frame_num
+            shift=args.sample_shift,
+            sample_solver=args.sample_solver,
+            sampling_steps=args.sample_steps,
+            guide_scale=args.sample_guide_scale,
+            seed=seed,#args.base_seed,
+            offload_model=args.offload_model,
+            speed=value2speed[speed]
+        )
+        if rank==0:
+            video_update = gr.update(visible=True, value=save_file)
+            seed_update = gr.update(visible=True, value=seed)
+            cache_video(
+                tensor=video[None],
+                save_file=save_file,
+                fps=cfg.sample_fps,
+                nrow=1,
+                normalize=True,
+                value_range=(-1, 1))
+            return save_file, video_update, seed_update
+    if rank == 0:
+        from app_os import DEMO
+        demo=DEMO(generate_i2v).demo
+        demo.launch()
+
+
+if __name__ == "__main__":
+    args = _parse_args()
+    generate(args)

app_os.py

@@ -0,0 +1,90 @@
+# -*- coding: utf-8 -*-
+"""
+THis is the main file for the gradio web demo. It uses the CogVideoX-5B model to generate videos gradio web demo.
+set environment variable OPENAI_API_KEY to use the OpenAI API to enhance the prompt.
+
+Usage:
+    OpenAI_API_KEY=your_openai_api_key OPENAI_BASE_URL=https://api.openai.com/v1 python inference/gradio_web_demo.py
+"""
+
+import logging
+import math
+import os
+import sys
+from fastapi.responses import PlainTextResponse
+from PIL import Image
+from huggingface_hub.utils.tqdm import progress_bar_states
+from numpy import ndarray
+
+current_dir = os.path.abspath(os.path.dirname(__file__))
+sys.path.append(os.path.join(current_dir, '../'))
+import random
+import threading
+import time
+
+import cv2
+import tempfile
+import imageio_ffmpeg
+import gradio as gr
+
+from datetime import datetime, timedelta
+
+os.makedirs("./output", exist_ok=True)
+os.makedirs("./input", exist_ok=True)
+os.makedirs("./gradio_tmp", exist_ok=True)
+
+class DEMO:
+    def __init__(self,generate):
+        with gr.Blocks() as self.demo:
+            gr.Markdown("""
+                   <div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
+                       AniSora-Bilibili动画视频生成模型
+                   </div>
+                   """)
+            with gr.Row():
+                with gr.Column():
+                    with gr.Accordion("I2V: Image Input (cannot be used simultaneously with video input)", open=True):
+                        image_input = gr.Image(label="Input Image")
+                    prompt = gr.Textbox(label="Prompt (Less than 200 Words)", placeholder="Enter your prompt here", lines=5)
+                    nf = gr.Slider(label="秒数", minimum=3,maximum=5, step=0.5, value=5)
+                    speed = gr.Radio(label="加速模式",value='加速版',choices=['原版','加速版'])
+                    with gr.Group():
+                        with gr.Column():
+                            with gr.Row():
+                                seed_param = gr.Number(
+                                    label="Inference Seed (Enter a positive number, -1 for random)", value=233
+                                )
+
+                    generate_button = gr.Button("🎬 Generate Video")
+
+                with gr.Column():
+                    video_output = gr.Video(label="Generated Video")
+                    with gr.Row():
+                        download_video_button = gr.File(label="📥 Download Video", visible=False)
+                        seed_text = gr.Number(label="Seed Used for Video Generation", visible=False)
+
+            generate_button.click(
+                generate,
+                inputs=[prompt, image_input,seed_param,nf,speed],
+                outputs=[video_output, download_video_button, seed_text],
+            )
+
+
+if __name__ == "__main__":
+    from fastapi import FastAPI
+    import uvicorn
+
+    app = FastAPI()
+
+
+    @app.get('/v2/health/ready')
+    def health():
+        return ""
+
+
+    demoo=DEMO()
+    demo=demoo.demo
+    demo.queue(max_size=15)
+    app = gr.mount_gradio_app(app,demo, path="/api/adhoc/ttv/demo")
+    uvicorn.run(app,host="0.0.0.0",port=26780)#
+    # demo.launch(server_name="0.0.0.0",server_port=16780)

generate-pi-i2v-myinfer-oss-stu.py

@@ -0,0 +1,452 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import os
+os.environ["HF_ENDPOINT"]          = "https://hf-mirror.com"
+import argparse
+from datetime import datetime
+import logging
+import sys
+import warnings
+
+warnings.filterwarnings('ignore')
+
+import torch, random
+import torch.distributed as dist
+from PIL import Image
+
+import wan
+from wan.image2video_mdinfer_oss_stu import WanI2V
+from wan.text2video import WanT2V
+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])}"
+
+    args.sample_steps=96#64########todo
+
+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(
+        "--student_steps",
+        type=int,
+        default=20,
+        help="The student steps during searching!")
+    parser.add_argument(
+        "--norm",
+        type=float,
+        default=2.0,
+        help="Norm of the cost function.")
+    parser.add_argument(
+        "--frame_type",
+        type=str,
+        default='all',
+        help="The cost frames of video.")
+    parser.add_argument(
+        "--channel_type",
+        type=str,
+        default="all",
+        choices=['2', '4', '8','12',"all"],
+        help="The cost channel of video.")
+    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:
+        opt_dir=args.image
+        with open(args.prompt,"r")as f:
+            lines=f.read().strip("\n").split("\n")
+        # if args.prompt is None:
+        #     args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
+
+        logging.info("Creating WanT2V pipeline.")
+        # wan_t2v = wan.WanT2V(
+        wan_t2v = 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,
+        )
+        for idx,line in enumerate(lines):
+            args.save_file="%s/%s.mp4"%(opt_dir,idx)
+            prompt,image=line.split("@@")
+            args.image=image
+            args.prompt=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(
+                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)
+            if rank==0:
+                cache_video(
+                    tensor=video[None],
+                    save_file=args.save_file,
+                    fps=cfg.sample_fps,
+                    nrow=1,
+                    normalize=True,
+                    value_range=(-1, 1))
+    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}")
+
+        opt_dir=args.image
+        with open(args.prompt,"r",encoding="gbk")as f:
+            lines=f.read().strip("\n").split("\n")
+        logging.info("Creating WanI2V pipeline.")
+        # wan_i2v = wan.WanI2V(
+        wan_i2v = 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,
+        )
+
+        for idx,line in enumerate(lines):
+            args.save_file="%s/%s.mp4"%(opt_dir,idx)
+            prompt,image=line.split("@@")
+            args.image=image
+            args.prompt=prompt
+            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("Generating video ...")
+            if os.path.exists(args.save_file)==False:
+                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,
+
+                    student_steps=16,
+                    norm=2,
+                    frame_type="4",
+                    channel_type="all",
+
+                )
+                if rank==0:
+                    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)

generate-pi-i2v-myinfer-oss-tea.py

@@ -0,0 +1,453 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import os
+os.environ["HF_ENDPOINT"]          = "https://hf-mirror.com"
+import argparse
+from datetime import datetime
+import logging
+import sys
+import warnings
+
+warnings.filterwarnings('ignore')
+
+import torch, random
+import torch.distributed as dist
+from PIL import Image
+
+import wan
+from wan.image2video_mdinfer_oss_tea import WanI2V
+from wan.text2video import WanT2V
+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(
+        "--student_steps",
+        type=int,
+        default=20,
+        help="The student steps during searching!")
+    parser.add_argument(
+        "--norm",
+        type=float,
+        default=2.0,
+        help="Norm of the cost function.")
+    parser.add_argument(
+        "--frame_type",
+        type=str,
+        default='all',
+        help="The cost frames of video.")
+    parser.add_argument(
+        "--channel_type",
+        type=str,
+        default="all",
+        choices=['2', '4', '8','12',"all"],
+        help="The cost channel of video.")
+    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:
+        opt_dir=args.image
+        with open(args.prompt,"r")as f:
+            lines=f.read().strip("\n").split("\n")
+        # if args.prompt is None:
+        #     args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
+
+        logging.info("Creating WanT2V pipeline.")
+        # wan_t2v = wan.WanT2V(
+        wan_t2v = 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,
+        )
+        for idx,line in enumerate(lines):
+            args.save_file="%s/%s.mp4"%(opt_dir,idx)
+            prompt,image=line.split("@@")
+            args.image=image
+            args.prompt=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(
+                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)
+            if rank==0:
+                cache_video(
+                    tensor=video[None],
+                    save_file=args.save_file,
+                    fps=cfg.sample_fps,
+                    nrow=1,
+                    normalize=True,
+                    value_range=(-1, 1))
+    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}")
+
+        opt_dir=args.image
+        with open(args.prompt,"r",encoding="gbk")as f:
+            lines=f.read().strip("\n").split("\n")
+        logging.info("Creating WanI2V pipeline.")
+        # wan_i2v = wan.WanI2V(
+        wan_i2v = 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,
+        )
+
+        for idx,line in enumerate(lines):
+            # args.save_file="%s/%s.mp4"%(opt_dir,idx)
+            args.save_file="%s-%s.mp4"%(opt_dir,idx)
+            prompt,image=line.split("@@")
+            args.image=image
+            args.prompt=prompt
+            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("Generating video ...")
+            if os.path.exists(args.save_file)==False:
+                video = wan_i2v.generate(
+                    args,
+                    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,
+
+                    student_steps=args.student_steps,#12,
+                    norm=2,
+                    frame_type="4",
+                    channel_type="all",
+
+                )
+                if rank==0:
+                    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)

generate-pi-i2v.py

@@ -0,0 +1,418 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import os
+os.environ["HF_ENDPOINT"]          = "https://hf-mirror.com"
+import argparse
+from datetime import datetime
+import logging
+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:
+        opt_dir=args.image
+        with open(args.prompt,"r")as f:
+            lines=f.read().strip("\n").split("\n")
+        # if args.prompt is None:
+        #     args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
+        for idx,line in enumerate(lines):
+            args.save_file="%s/%s.mp4"%(opt_dir,idx)
+            prompt,image=line.split("@@")
+            args.image=image
+            args.prompt=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)
+            if rank==0:
+                cache_video(
+                    tensor=video[None],
+                    save_file=args.save_file,
+                    fps=cfg.sample_fps,
+                    nrow=1,
+                    normalize=True,
+                    value_range=(-1, 1))
+    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}")
+
+        opt_dir=args.image
+        with open(args.prompt,"r",encoding="gbk")as f:
+            lines=f.read().strip("\n").split("\n")
+        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,
+        )
+
+        for idx,line in enumerate(lines):
+            args.save_file="%s/%s.mp4"%(opt_dir,idx)
+            prompt,image=line.split("@@")
+            args.image=image
+            args.prompt=prompt
+            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("Generating video ...")
+            if os.path.exists(args.save_file)==False:
+                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:
+                    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)

get-med.py

@@ -0,0 +1,25 @@
+datas=[
+    [3, 12, 38, 55, 68, 72, 76, 79, 82, 84, 87, 89, 91, 92, 94, 96],
+    [2, 5, 9, 16, 29, 42, 58, 68, 75, 79, 83, 88, 92, 94, 95, 96],
+    [1, 3, 7, 13, 24, 36, 44, 54, 63, 71, 78, 85, 89, 93, 95, 96],
+    [3, 10, 39, 60, 73, 77, 81, 84, 86, 88, 90, 92, 93, 94, 95, 96],
+    [2, 5, 11, 21, 36, 51, 65, 73, 80, 85, 88, 91, 93, 94, 95, 96],
+    [1, 3, 6, 11, 20, 36, 49, 65, 71, 77, 83, 87, 91, 94, 95, 96],
+    [2, 6, 15, 26, 39, 49, 57, 64, 70, 74, 80, 86, 90, 93, 95, 96],
+    [1, 4, 9, 17, 30, 49, 63, 71, 81, 86, 89, 91, 93, 94, 95, 96],
+    [4, 13, 29, 45, 58, 67, 73, 76, 80, 84, 88, 90, 92, 94, 95, 96],
+]
+def calculate_median(arr):
+    sorted_arr = sorted(arr)
+    n = len(sorted_arr)
+    middle_index = n // 2
+    return sorted_arr[middle_index]
+
+opt=[]
+leng=len(datas[0])
+len_datas=len(datas)
+for i in range(leng):
+    tmp=[datas[j][i]for j in range(len_datas)]
+    opt.append(calculate_median(tmp))
+print(opt)
+#544P#96-12#[5, 20, 54, 63, 71, 79, 82, 87, 91, 94, 95, 96]

note-webui.txt

@@ -0,0 +1,10 @@
+python generate-pi-i2v-app-os.py --task i2v-14B --size 960*544 --ckpt_dir /data/docker/eryu51mmd
+
+
+
+
+python generate-pi-i2v-app-os.py --task i2v-14B --size 960*544 --ckpt_dir /data/docker/eryu-mysp2v2-500/lns
+python generate-pi-i2v-app-os.py --task i2v-14B --size 960*544 --ckpt_dir /data/docker/eryu0417
+
+todo:
+    add motion

preprocess/extract-clip.py

@@ -0,0 +1,57 @@
+'''
+nohup python preprocess/extract-clip.py 0 2 >> clip.log 2>&1 &
+nohup python preprocess/extract-clip.py 1 2 >> clip.log 2>&1 &
+
+0/1: What part does this process do
+2:   It consists of two parts in total.
+'''
+root_mp4s="test_data/mp4root"
+h,w=480,832
+opt_root="output_root/clip"
+checkpoint_dir="/DATA/bvac/personal/wan21/Wan2.1-I2V-14B-720P"
+
+
+import os,sys,traceback
+import pdb
+all=int(sys.argv[2])
+i_part=int(sys.argv[1])
+# os.environ["CUDA_VISIBLE_DEVICES"]=str(int(sys.argv[1])%4)
+os.environ["CUDA_VISIBLE_DEVICES"]=sys.argv[1]
+import pdb,torch
+from wan.modules.clip import CLIPModel
+device="cuda"
+clip = CLIPModel(
+    dtype=torch.float16,
+    device=device,
+    checkpoint_path=os.path.join(checkpoint_dir,'models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth'),
+    tokenizer_path=os.path.join(checkpoint_dir, 'xlm-roberta-large'))
+
+clip.model = clip.model.to(device)
+from decord import VideoReader
+import torchvision.transforms.functional as TF
+def read_img(path):
+    vr = VideoReader(uri=path, height=-1, width=-1)
+    temp_frms = vr.get_batch([2])
+    return (TF.to_tensor(temp_frms.asnumpy().astype("float32")[0])/255).sub_(0.5).div_(0.5).to(device)
+
+os.makedirs(opt_root,exist_ok=True)
+def go(todos):
+    for path in todos:
+        try:
+            name=os.path.basename(path).replace(".mp4",".pt")
+            if os.path.exists("%s/%s"%(opt_root,name)):continue
+            img = read_img(path)
+            clip_context = clip.visual([img[:, None, :, :]])
+            save_path="%s/%s"%(opt_root,name)
+            torch.save(clip_context, save_path)
+        except:
+            print(path,traceback.format_exc())
+
+todo=[]
+for name in os.listdir(root_mp4s):
+    todo.append("%s/%s"%(root_mp4s,name))
+todo=sorted(todo)
+todo=todo[i_part::all]
+go(todo)
+
+

preprocess/extract-t5.py

@@ -0,0 +1,46 @@
+'''
+nohup python preprocess/extract-t5.py 0 2 >> t5.log 2>&1 &
+nohup python preprocess/extract-t5.py 1 2 >> t5.log 2>&1 &
+
+0/1: What part does this process do
+2:   It consists of two parts in total.
+'''
+txt_path="test_data/train_data_prompts.txt"
+opt_root="output_root/t5"
+checkpoint_dir="/DATA/bvac/personal/wan21/Wan2.1-I2V-14B-720P"
+
+import os,sys
+all=int(sys.argv[2])
+i_part=int(sys.argv[1])
+# os.environ["CUDA_VISIBLE_DEVICES"]=str(int(sys.argv[1])%4)
+os.environ["CUDA_VISIBLE_DEVICES"]=sys.argv[1]
+import pdb,torch
+from wan.modules.t5 import T5EncoderModel
+device="cuda"
+text_encoder = T5EncoderModel(
+    text_len=512,
+    dtype=torch.bfloat16,
+    device=torch.device('cpu'),
+    checkpoint_path=os.path.join(checkpoint_dir, 'models_t5_umt5-xxl-enc-bf16.pth'),
+    tokenizer_path=os.path.join(checkpoint_dir, 'google/umt5-xxl'),
+    shard_fn=None,
+)
+text_encoder.model=text_encoder.model.to(device)
+os.makedirs(opt_root,exist_ok=True)
+def go(todos):
+    for name,text in todos:
+        try:
+            if os.path.exists("%s/%s"%(opt_root,name)):continue
+            context = text_encoder([text], device)[0].cpu()#torch.Size([138, 4096])#"In this scene, a man with a beard is seen tending to a woman who lies in bed, her face illuminated by the soft glow of a nearby light source. The man, dressed in a blue robe adorned with intricate designs, holds a bowl, possibly containing a healing potion or a magical elixir. The woman, clad in a pink garment, appears to be resting or possibly unwell, as she lies on her side with her eyes closed. The setting suggests a historical or medieval context, with the dimly lit room and the man's attire evoking a sense of timelessness and mystery. "
+            save_path="%s/%s"%(opt_root,name)
+            torch.save(context,save_path)
+        except:
+            print(text,traceback.format_exc())
+
+
+todo=[]
+with open(txt_path,"r")as f:lines=f.read().strip("\n").split("\n")
+for line in lines:
+    todo.append(line.split("|"))
+todo=sorted(todo)[i_part::all]
+go(todo)

preprocess/extract-vae1.py

@@ -0,0 +1,62 @@
+'''
+nohup python preprocess/extract-vae1.py 0 2 >> vae1.log 2>&1 &
+nohup python preprocess/extract-vae1.py 1 2 >> vae1.log 2>&1 &
+
+0/1: What part does this process do
+2:   It consists of two parts in total.
+'''
+root_mp4s="test_data/mp4root"
+h,w=480,832
+num_frames = 49
+opt_root="output_root/vae1"
+checkpoint_dir="/DATA/bvac/personal/wan21/Wan2.1-I2V-14B-480P"
+
+
+
+
+import os,sys,traceback
+import pdb
+os.environ["CUDA_VISIBLE_DEVICES"]=sys.argv[1]
+all=int(sys.argv[2])
+i_part=int(os.environ["CUDA_VISIBLE_DEVICES"])
+import pdb,torch
+from wan.modules.vae import WanVAE
+device="cuda"
+vae = WanVAE(vae_pth=os.path.join(checkpoint_dir, 'Wan2.1_VAE.pth'),device=device)
+
+from decord import VideoReader
+import torchvision.transforms.functional as TF
+def read_img(path):
+    vr = VideoReader(uri=path, height=-1, width=-1)
+    temp_frms = vr.get_batch([2])
+    return (TF.to_tensor(temp_frms.asnumpy().astype("float32")[0])/255).sub_(0.5).div_(0.5).to(device)
+
+os.makedirs(opt_root,exist_ok=True)
+def go(todos):
+    for path in todos:
+        try:
+            name=os.path.basename(path).replace(".mp4",".pt")
+            if os.path.exists("%s/%s"%(opt_root,name)):continue
+            img = read_img(path)
+            tensorr = vae.encode([
+                torch.concat([
+                    torch.nn.functional.interpolate(
+                        img[None].cpu(), size=(h, w), mode='bicubic').transpose(
+                        0, 1),
+                    torch.zeros(3, num_frames-1, h, w)
+                ], dim=1).to(device)
+            ])[0].cpu()  # torch.Size([16, 21, 90, 160])#21->13
+            save_path="%s/%s"%(opt_root,name)
+            torch.save(tensorr, save_path)
+        except:
+            print(path,traceback.format_exc())
+
+todo=[]
+for name in os.listdir(root_mp4s):
+    todo.append("%s/%s"%(root_mp4s,name))
+todo=sorted(todo)
+todo=todo[i_part::all]
+go(todo)
+
+
+

preprocess/extract-vae_all.py

@@ -0,0 +1,64 @@
+'''
+nohup python preprocess/extract-vae_all.py 0 2 >> t5.log 2>&1 &
+nohup python preprocess/extract-vae_all.py 1 2 >> t5.log 2>&1 &
+
+0/1: What part does this process do
+2:   It consists of two parts in total.
+'''
+root_mp4s="test_data/mp4root"
+h,w=480,832
+num_frames = 49
+opt_root="output_root/vae_all"
+checkpoint_dir="/DATA/bvac/personal/wan21/Wan2.1-I2V-14B-720P"
+
+
+import os,sys,traceback,numpy as np
+import pdb
+os.environ["CUDA_VISIBLE_DEVICES"]=sys.argv[1]
+all=int(sys.argv[2])
+i_part=int(os.environ["CUDA_VISIBLE_DEVICES"])
+import pdb,torch
+from wan.modules.vae import WanVAE
+device="cuda"
+vae = WanVAE(vae_pth=os.path.join(checkpoint_dir, 'Wan2.1_VAE.pth'),device=device)
+from decord import VideoReader
+import torchvision.transforms.functional as TF
+def read_img(path):
+    vr = VideoReader(uri=path, height=-1, width=-1)
+    actual_fps=vr.get_avg_fps()
+    start=2
+    wanted_fps=16
+    end = int(start + num_frames / wanted_fps * actual_fps)
+    indices = np.arange(start, end, (end - start) / num_frames).astype(int)
+    # print(100000,start,end,indices,len(indices),actual_fps,len(vr))
+    temp_frms = vr.get_batch(indices)
+    temp_frms = torch.from_numpy(temp_frms.asnumpy()).to(device).float()/255
+    temp_frms-=0.5
+    temp_frms/=0.5#torch.Size([49, 1080, 1920, 3])
+    temp_frms=temp_frms.permute(3,0,1,2)#torch.Size([3, 49, 1080, 1920])
+    return temp_frms
+
+
+os.makedirs(opt_root,exist_ok=True)
+def go(todos):
+    for path in todos:
+        try:
+            name=os.path.basename(path).rsplit('.', maxsplit=1)[0]+'.pt'
+            if os.path.exists("%s/%s"%(opt_root,name)):continue
+            img = read_img(path)
+            aa=torch.nn.functional.interpolate(img, size=(h, w), mode='bicubic')
+            tensorr = vae.encode(aa.unsqueeze(0))[0].cpu()  # torch.Size([16, 21, 90, 160])#21->13
+            save_path="%s/%s"%(opt_root,name)
+            torch.save(tensorr, save_path)
+        except:
+            print(path,traceback.format_exc())
+
+todo=[]
+for name in os.listdir(root_mp4s):
+    todo.append("%s/%s"%(root_mp4s,name))
+todo=sorted(todo)
+todo=todo[i_part::all]
+go(todo)
+
+
+

pyproject.toml

@@ -0,0 +1,41 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "fastvideo"
+version = "1.2.0"
+description = "FastVideo"
+readme = "README1.md"
+requires-python = ">=3.8"
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: Apache Software License",
+]
+dependencies = [
+#"transformers==4.46.1", "accelerate==1.0.1", "tokenizers==0.20.1", "albumentations==1.4.20", "av==13.1.0",
+#"decord==0.6.0", "einops==0.8.0", "fastapi==0.115.3", "gdown==5.2.0", "h5py==3.12.1", "idna==3.6", "imageio==2.36.0",
+#"matplotlib==3.9.2", "numpy==1.26.3", "omegaconf==2.3.0", "opencv-python==4.10.0.84", "opencv-python-headless==4.10.0.84",
+#"pandas==2.2.3", "pillow==10.2.0", "pydub==0.25.1", "pytorch-lightning==2.4.0", "pytorchvideo==0.1.5", "PyYAML==6.0.1",
+#"regex==2024.9.11", "requests==2.31.0", "scikit-learn==1.5.2", "scipy==1.14.1", "six==1.16.0", "test-tube==0.7.5",
+#"timm==1.0.11", "torchdiffeq==0.2.4", "torchmetrics==1.5.1", "tqdm==4.66.5", "urllib3==2.2.0", "uvicorn==0.32.0",
+#"scikit-video==1.1.11", "imageio-ffmpeg==0.5.1", "sentencepiece==0.2.0", "beautifulsoup4==4.12.3", "ftfy==6.3.0",
+#"moviepy==1.0.3", "wandb==0.18.5", "tensorboard==2.18.0", "pydantic==2.9.2", "gradio==5.3.0", "huggingface_hub==0.26.1", "protobuf==5.28.3",
+#"watch", "gpustat", "peft==0.13.2", "liger_kernel==0.4.1", "einops==0.8.0", "wheel==0.44.0", "loguru", "diffusers==0.32.0", "bitsandbytes"]
+#"watch", "gpustat", "bitsandbytes"
+]
+
+
+[tool.setuptools.packages.find]
+exclude = ["assets*", "docker*", "docs", "scripts*"]
+
+[tool.wheel]
+exclude = ["assets*", "docker*", "docs", "scripts*"]
+
+[tool.mypy]
+warn_return_any = true
+warn_unused_configs = true
+ignore_missing_imports = true
+disallow_untyped_calls = true
+check_untyped_defs = true
+no_implicit_optional = true

req-fastvideo.txt

@@ -0,0 +1,59 @@
+transformers>=4.46.1
+accelerate>=1.0.1
+tokenizers>=0.20.1
+albumentations>=1.4.20
+av>=13.1.0
+decord>=0.6.0
+einops>=0.8.0
+fastapi>=0.115.3
+gdown>=5.2.0
+h5py>=3.12.1
+idna>=3.6
+imageio>=2.36.0
+matplotlib>=3.9.2
+numpy>=1.26.3
+omegaconf>=2.3.0
+opencv-python>=4.10.0.84
+opencv-python-headless>=4.10.0.84
+pandas>=2.2.3
+pillow>=10.2.0
+pydub>=0.25.1
+pytorch-lightning>=2.4.0
+pytorchvideo>=0.1.5
+PyYAML>=6.0.1
+regex>=2024.9.11
+requests>=2.31.0
+scikit-learn>=1.5.2
+scipy>=1.14.1
+six>=1.16.0
+test-tube>=0.7.5
+timm>=1.0.11
+torchdiffeq>=0.2.4
+torchmetrics>=1.5.1
+tqdm>=4.66.5
+urllib3>=2.2.0
+uvicorn>=0.32.0
+scikit-video>=1.1.11
+imageio-ffmpeg>=0.5.1
+sentencepiece>=0.2.0
+beautifulsoup4>=4.12.3
+ftfy>=6.3.0
+moviepy>=1.0.3
+wandb>=0.18.5
+tensorboard>=2.18.0
+pydantic>=2.9.2
+gradio>=5.3.0
+huggingface_hub>=0.26.1
+protobuf>=5.28.3
+watch
+gpustat
+peft>=0.13.2
+liger_kernel>=0.4.1
+einops>=0.8.0
+wheel>=0.44.0
+loguru
+torch
+torchvision
+ninja
+safetensors
+packaging

requirements.txt

@@ -0,0 +1,15 @@
+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
+gradio>=5.0.0
+numpy>=1.23.5,<2

scripts/distill/distill_cog.sh

@@ -0,0 +1,40 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export HF_ENDPOINT="https://hf-mirror.com"
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+#based on fm1500
+torchrun --nnodes 1 --nproc_per_node 8  \
+    fastvideo/distill_i2v_cog.py  \
+    --seed 42 \
+    --pretrained_model_name_or_path /DATA/bvac/personal/opensora/zhipu/pretrained/danzhen-i2v-97f-720P-16fps/danzhen-i2v-97f-720P-16fps  \
+    --dit_model_name_or_path /DATA/bvac/personal/opensora/zhipu/pretrained/danzhen-i2v-97f-720P-16fps/danzhen-i2v-97f-720P-16fps  \
+    --model_type "cog"  \
+    --cache_dir .cache  \
+    --data_json_path /DATA/bvac/personal/fastvideo/FastVideo-main/data/Image-Vid-Finetune-cog480P-49f12fps-fixvae/videos2caption.json  \
+    --validation_prompt_dir /DATA/bvac/personal/fastvideo/FastVideo-main/data/Image-Vid-Finetune-cog480P-49f12fps-fixvae/validation  \
+    --gradient_checkpointing  \
+    --train_batch_size=1  \
+    --num_latent_t 24  \
+    --sp_size 1  \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=20000 \
+    --learning_rate=3e-6 \
+    --mixed_precision="bf16" \
+    --checkpointing_steps=400 \
+    --validation_steps 15000000 \
+    --validation_sampling_steps "2,4,8"  \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --log_validation \
+    --output_dir=outputs-test-cog_i2v_5B-wukong49-distill-3e-6-fixvae2 \
+    --tracker_project_name video3w480Ptest_cog49i2v_fix_distill-3e-6-fixvae2  \
+    --num_frames  49  \
+    --shift 17  \
+    --validation_guidance_scale "1.0"  \
+    --num_euler_timesteps 50  \
+    --multi_phased_distill_schedule "4000-1"  \
+    --not_apply_cfg_solver 

scripts/distill/distill_cog720-49.sh

@@ -0,0 +1,40 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export HF_ENDPOINT="https://hf-mirror.com"
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+#based on fm1500
+torchrun --nnodes 3 --nproc_per_node 8  --node_rank=${NODE_RANK} --master_addr=10.156.32.11 --master_port=14431 \
+    fastvideo/distill_i2v_cog720-49.py  \
+    --seed 42 \
+    --pretrained_model_name_or_path /DATA/bvac/personal/opensora/zhipu/pretrained/danzhen-i2v-97f-720P-16fps/danzhen-i2v-97f-720P-16fps-1600steps_bs24-fm6000  \
+    --dit_model_name_or_path /DATA/bvac/personal/opensora/zhipu/pretrained/danzhen-i2v-97f-720P-16fps/danzhen-i2v-97f-720P-16fps-1600steps_bs24-fm6000  \
+    --model_type "cog"  \
+    --cache_dir /data/docker/data/video14w/.cache \
+    --data_json_path /data/docker/data/video14w/outputs/videos2caption.json \
+    --validation_prompt_dir /data/docker/data/video14w/outputs/validation \
+    --gradient_checkpointing  \
+    --train_batch_size=1  \
+    --num_latent_t 24  \
+    --sp_size 1  \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=20000 \
+    --learning_rate=3e-6 \
+    --mixed_precision="bf16" \
+    --checkpointing_steps=250 \
+    --validation_steps 15000000 \
+    --validation_sampling_steps "2,4,8"  \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --log_validation \
+    --output_dir=outputs-cog_i2v_5B_down_vae_eryu49_3gpus_fix_text_rot_fvfm6k_distill \
+    --tracker_project_name cog_i2v_5B_down_vae_eryu49_3gpus_fix_text_rot_fvfm6k_distill  \
+    --num_frames  49  \
+    --shift 17  \
+    --validation_guidance_scale "1.0"  \
+    --num_euler_timesteps 50  \
+    --multi_phased_distill_schedule "4000-1"  \
+    --not_apply_cfg_solver 

scripts/distill/distill_cog720-49mix246adv.sh

@@ -0,0 +1,40 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export HF_ENDPOINT="https://hf-mirror.com"
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+#based on fm1500
+torchrun --nnodes 3 --nproc_per_node 8  --node_rank=${NODE_RANK} --master_addr=10.156.32.11 --master_port=14431 \
+    fastvideo/distill_adv-cog720-49.py  \
+    --seed 42 \
+    --pretrained_model_name_or_path /DATA/bvac/personal/opensora/zhipu/pretrained/danzhen-i2v-97f-720P-16fps/danzhen-i2v-97f-720P-16fps-fm246mix  \
+    --dit_model_name_or_path /DATA/bvac/personal/opensora/zhipu/pretrained/danzhen-i2v-97f-720P-16fps/danzhen-i2v-97f-720P-16fps-fm246mix  \
+    --model_type "cog"  \
+    --cache_dir /data/docker/data/video14w/.cache \
+    --data_json_path /data/docker/data/video14w/outputs/videos2caption.json \
+    --validation_prompt_dir /data/docker/data/video14w/outputs/validation \
+    --gradient_checkpointing  \
+    --train_batch_size=1  \
+    --num_latent_t 24  \
+    --sp_size 1  \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=20000 \
+    --learning_rate=1e-6 \
+    --mixed_precision="bf16" \
+    --checkpointing_steps=64 \
+    --validation_steps 15000000 \
+    --validation_sampling_steps "2,4,8"  \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --log_validation \
+    --output_dir=outputs-cog_i2v_5B_down_vae_eryu49_3gpus_fix_text_rot_fvfm_mix246_distill1e6adv \
+    --tracker_project_name cog_i2v_5B_down_vae_eryu49_3gpus_fix_text_rot_fvfm_mix246_distill1e6adv  \
+    --num_frames  49  \
+    --shift 17  \
+    --validation_guidance_scale "1.0"  \
+    --num_euler_timesteps 50  \
+    --multi_phased_distill_schedule "4000-1"  \
+    --not_apply_cfg_solver 

scripts/distill/distill_cog720-49mix26.sh

@@ -0,0 +1,42 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export HF_ENDPOINT="https://hf-mirror.com"
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+
+
+torchrun --nnodes 5 --nproc_per_node 8  --node_rank=${NODE_RANK} --master_addr=10.156.32.11 --master_port=14431 \
+    fastvideo/distill_i2v_cog720-4133-nonegrid.py  \
+    --seed 42 \
+    --pretrained_model_name_or_path /DATA/bvac/personal/opensora/zhipu/pretrained/mix26  \
+    --dit_model_name_or_path /DATA/bvac/personal/opensora/zhipu/pretrained/mix26  \
+    --model_type "cog"  \
+    --cache_dir /data/docker/data/video14w/.cache \
+    --data_json_path2 /DATA/bvac/personal/fastvideo/make-data/all33-41-49/videos2caption41.json \
+    --data_json_path3 /DATA/bvac/personal/fastvideo/make-data/all33-41-49/videos2caption33.json \
+    --validation_prompt_dir /data/docker/data/video14w/outputs/validation \
+    --gradient_checkpointing  \
+    --train_batch_size=1  \
+    --num_latent_t 24  \
+    --sp_size 1  \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=20000 \
+    --learning_rate=1.5e-6 \
+    --mixed_precision="bf16" \
+    --checkpointing_steps=64 \
+    --validation_steps 15000000 \
+    --validation_sampling_steps "2,4,8"  \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --log_validation \
+    --output_dir=outputs-cog_i2sv_5B_down_vae_eryu4133_5gpus_fix_text_rot_fvfm_mix26_distill1d5e6 \
+    --tracker_project_name cog_i2v_5B_down_vae_eryu4133_5gpus_fix_text_rot_fvfm_mix26_distill1d5e6  \
+    --num_frames  41  \
+    --shift 17  \
+    --validation_guidance_scale "1.0"  \
+    --num_euler_timesteps 50  \
+    --multi_phased_distill_schedule "4000-1"  \
+    --not_apply_cfg_solver 

scripts/distill/distill_cog720-49mix26b.sh

@@ -0,0 +1,42 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+export HF_ENDPOINT="https://hf-mirror.com"
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
+
+
+torchrun --nnodes 4 --nproc_per_node 8  --node_rank=${NODE_RANK} --master_addr=10.156.32.11 --master_port=14431 \
+    fastvideo/distill_i2v_cog720-4133-nonegrid.py  \
+    --seed 42 \
+    --pretrained_model_name_or_path /DATA/bvac/personal/opensora/zhipu/pretrained/cog_i2sv_5B_down_vae_eryu4133_5gpus_fix_text_rot_fvfm_mix26_distill1d5e6/mix5221  \
+    --dit_model_name_or_path /DATA/bvac/personal/opensora/zhipu/pretrained/cog_i2sv_5B_down_vae_eryu4133_5gpus_fix_text_rot_fvfm_mix26_distill1d5e6/mix5221  \
+    --model_type "cog"  \
+    --cache_dir /data/docker/data/video14w/.cache \
+    --data_json_path2 /DATA/bvac/personal/fastvideo/make-data/all33-41-49/videos2caption41.json \
+    --data_json_path3 /DATA/bvac/personal/fastvideo/make-data/all33-41-49/videos2caption33.json \
+    --validation_prompt_dir /data/docker/data/video14w/outputs/validation \
+    --gradient_checkpointing  \
+    --train_batch_size=1  \
+    --num_latent_t 24  \
+    --sp_size 1  \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=20000 \
+    --learning_rate=1e-6 \
+    --mixed_precision="bf16" \
+    --checkpointing_steps=64 \
+    --validation_steps 15000000 \
+    --validation_sampling_steps "2,4,8"  \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --log_validation \
+    --output_dir=outputs-cog_i2sv_5B_down_vae_eryu4133_4gpus_fix_text_rot_fvfm_mix26_distill1e6b \
+    --tracker_project_name cog_i2v_5B_down_vae_eryu4133_4gpus_fix_text_rot_fvfm_mix26_distill1e6b  \
+    --num_frames  41  \
+    --shift 17  \
+    --validation_guidance_scale "1.0"  \
+    --num_euler_timesteps 50  \
+    --multi_phased_distill_schedule "4000-1"  \
+    --not_apply_cfg_solver 

scripts/distill/distill_hunyuan.sh

@@ -0,0 +1,40 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+
+DATA_DIR=./data
+
+torchrun --nnodes 1 --nproc_per_node 8\
+    fastvideo/distill.py\
+    --seed 42\
+    --pretrained_model_name_or_path $DATA_DIR/hunyuan\
+    --dit_model_name_or_path $DATA_DIR/hunyuan/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt\
+    --model_type "hunyuan" \
+    --cache_dir "$DATA_DIR/.cache"\
+    --data_json_path "$DATA_DIR/Hunyuan-30K-Distill-Data/videos2caption.json"\
+    --validation_prompt_dir "$DATA_DIR/Hunyuan-Distill-Data/validation"\
+    --gradient_checkpointing\
+    --train_batch_size=1\
+    --num_latent_t 24\
+    --sp_size 1\
+    --train_sp_batch_size 1\
+    --dataloader_num_workers 4\
+    --gradient_accumulation_steps=1\
+    --max_train_steps=2000\
+    --learning_rate=1e-6\
+    --mixed_precision="bf16"\
+    --checkpointing_steps=64\
+    --validation_steps 64\
+    --validation_sampling_steps "2,4,8" \
+    --checkpoints_total_limit 3\
+    --allow_tf32\
+    --ema_start_step 0\
+    --cfg 0.0\
+    --log_validation\
+    --output_dir="$DATA_DIR/outputs/hy_phase1_shift17_bs_32"\
+    --tracker_project_name Hunyuan_Distill \
+    --num_frames  93 \
+    --shift 17 \
+    --validation_guidance_scale "1.0" \
+    --num_euler_timesteps 50 \
+    --multi_phased_distill_schedule "4000-1" \
+    --not_apply_cfg_solver 

scripts/distill/distill_mochi.sh

@@ -0,0 +1,38 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+export WANDB_MODE=online
+
+torchrun --nnodes 1 --nproc_per_node 4 \
+    fastvideo/distill.py \
+    --seed 42 \
+    --pretrained_model_name_or_path data/mochi \
+    --model_type "mochi" \
+    --cache_dir data/.cache \
+    --data_json_path data/Merge-30k-Data/video2caption.json \
+    --validation_prompt_dir data/Image-Vid-Finetune-Mochi/validation \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --num_latent_t 28 \
+    --sp_size 4 \
+    --train_sp_batch_size 2 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=4000 \
+    --learning_rate=1e-6 \
+    --mixed_precision=bf16 \
+    --checkpointing_steps=64 \
+    --validation_steps 1 \
+    --validation_sampling_steps 8 \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --log_validation \
+    --output_dir="data/outputs/lq_euler_50_thres0.1_lrg_0.75_phase1_lr1e-6_repro" \
+    --tracker_project_name PCM \
+    --num_frames 163 \
+    --scheduler_type pcm_linear_quadratic \
+    --validation_guidance_scale 0.5,1.5,2.5 \
+    --num_euler_timesteps 50 \
+    --linear_quadratic_threshold 0.1 \
+    --linear_range 0.75 \
+    --multi_phased_distill_schedule 4000-1

scripts/finetune/finetune_wan.sh

@@ -0,0 +1,39 @@
+export WANDB_BASE_URL="https://api.wandb.ai"
+#export WANDB_MODE=online
+export WANDB_MODE=offline
+export WANDB_API_KEY=xxx
+export HF_ENDPOINT="https://hf-mirror.com"
+export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True #
+export TORCH_NCCL_TRACE_BUFFER_SIZE=1048576  # 1MB
+export TORCH_NCCL_ASYNC_ERROR_HANDLING=1
+export CUBLAS_WORKSPACE_CONFIG=:4096:8
+
+torchrun --nnodes 1 --nproc_per_node 8  --node_rank=0 --master_addr=127.0.0.1 --master_port=24431 \
+    fastvideo/train.py \
+    --seed 142 \
+    --pretrained_model_name_or_path /DATA/bvac/personal/wan21/Wan2.1-I2V-14B-480P \
+    --model_type "wan" \
+    --data_json_path test_data/data.json \
+    --gradient_checkpointing \
+    --train_batch_size=1 \
+    --num_latent_t 240 \
+    --sp_size 8 \
+    --train_sp_batch_size 1 \
+    --dataloader_num_workers 4 \
+    --gradient_accumulation_steps=1 \
+    --max_train_steps=20000 \
+    --learning_rate=1e-5 \
+    --mixed_precision=bf16 \
+    --checkpointing_steps=400 \
+    --validation_steps 20000 \
+    --validation_sampling_steps 64 \
+    --checkpoints_total_limit 3 \
+    --allow_tf32 \
+    --ema_start_step 0 \
+    --cfg 0.0 \
+    --ema_decay 0.999 \
+    --log_validation \
+    --output_dir=outputs-sp8 \
+    --tracker_project_name sp8 \
+    --validation_guidance_scale "1.0" \
+    --group_frame

scripts/huggingface/download_hf.py

@@ -0,0 +1,39 @@
+from huggingface_hub import snapshot_download, hf_hub_download
+import argparse
+
+# set args for repo_id, local_dir, repo_type,
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Download a dataset or model from the Hugging Face Hub"
+    )
+    parser.add_argument(
+        "--repo_id", type=str, help="The ID of the repository to download"
+    )
+    parser.add_argument(
+        "--local_dir",
+        type=str,
+        help="The local directory to download the repository to",
+    )
+    parser.add_argument(
+        "--repo_type",
+        type=str,
+        help="The type of repository to download (dataset or model)",
+    )
+    parser.add_argument("--file_name", type=str, help="The file name to download")
+    args = parser.parse_args()
+    if args.file_name:
+        hf_hub_download(
+            repo_id=args.repo_id,
+            filename=args.file_name,
+            repo_type=args.repo_type,
+            local_dir=args.local_dir,
+        )
+    else:
+        snapshot_download(
+            repo_id=args.repo_id,
+            local_dir=args.local_dir,
+            repo_type=args.repo_type,
+            local_dir_use_symlinks=False,
+            resume_download=True,
+        )

scripts/huggingface/upload_hf.py

@@ -0,0 +1,9 @@
+from huggingface_hub import HfApi
+
+api = HfApi()
+
+api.upload_folder(
+    folder_path="data/Black-Myth-Taylor-Src",
+    repo_id="FastVideo/Image-Vid-Finetune-Src",
+    repo_type="dataset",
+)

scripts/inference/inference_diffusers_hunyuan.sh

@@ -0,0 +1,20 @@
+#!/bin/bash
+
+num_gpus=1
+export MODEL_BASE="data/FastHunyuan-diffusers"
+torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 12345 \
+    fastvideo/sample/sample_t2v_diffusers_hunyuan.py \
+    --height 720 \
+    --width 1280 \
+    --num_frames 45 \
+    --num_inference_steps 6 \
+    --guidance_scale 1 \
+    --embedded_cfg_scale 6 \
+    --flow_shift 17 \
+    --flow-reverse \
+    --prompt ./assets/prompt.txt \
+    --seed 1024 \
+    --output_path outputs_video/hunyuan_quant/nf4/ \
+    --model_path $MODEL_BASE \
+    --quantization "nf4" \
+    --cpu_offload

scripts/inference/inference_hunyuan.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+num_gpus=4
+export MODEL_BASE=data/FastHunyuan
+torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
+    fastvideo/sample/sample_t2v_hunyuan.py \
+    --height 720 \
+    --width 1280 \
+    --num_frames 125 \
+    --num_inference_steps 6 \
+    --guidance_scale 1 \
+    --embedded_cfg_scale 6 \
+    --flow_shift 17 \
+    --flow-reverse \
+    --prompt ./assets/prompt.txt \
+    --seed 1024 \
+    --output_path outputs_video/hunyuan/cfg6/ \
+    --model_path $MODEL_BASE \
+    --dit-weight ${MODEL_BASE}/hunyuan-video-t2v-720p/transformers/mp_rank_00_model_states.pt

scripts/inference/inference_mochi_sp.sh

@@ -0,0 +1,19 @@
+#!/bin/bash
+
+num_gpus=4
+
+torchrun --nnodes=1 --nproc_per_node=$num_gpus --master_port 29503 \
+    fastvideo/sample/sample_t2v_mochi.py \
+    --model_path data/FastMochi-diffusers \
+    --prompt_path "assets/prompt.txt" \
+    --num_frames 163 \
+    --height 480 \
+    --width 848 \
+    --num_inference_steps 8 \
+    --guidance_scale 1.5 \
+    --output_path outputs_video/mochi_sp/ \
+    --seed 1024 \
+    --scheduler_type "pcm_linear_quadratic" \
+    --linear_threshold 0.1 \
+    --linear_range 0.75
+

scripts/preprocess/preprocess_cog_data.sh

@@ -0,0 +1,35 @@
+# export WANDB_MODE="offline"#
+######test-t2v-5B-wukong
+export HF_ENDPOINT="https://hf-mirror.com"
+GPU_NUM=4 # 2,4,8
+MODEL_PATH="/DATA/bvac/personal/opensora/zhipu/pretrained/hf2b/cache/models--THUDM--CogVideoX-5b/snapshots/8d6ea3f817438460b25595a120f109b88d5fdfad"
+MODEL_TYPE="cog"
+DATA_MERGE_PATH="/DATA/bvac/personal/fastvideo/FastVideo-main/data/Image-Vid-Finetune-Src/merge.txt"
+OUTPUT_DIR="/DATA/bvac/personal/fastvideo/FastVideo-main/data/Image-Vid-Finetune-cog480P-49f8fps"
+VALIDATION_PATH="/DATA/bvac/personal/fastvideo/FastVideo-main/assets/prompt1.txt"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/data_preprocess/preprocess_vae_latents.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --train_batch_size=1 \
+    --max_height=480 \
+    --max_width=720 \
+    --num_frames=49 \
+    --dataloader_num_workers 1 \
+    --output_dir=$OUTPUT_DIR \
+    --model_type $MODEL_TYPE \
+    --train_fps 12
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/data_preprocess/preprocess_text_embeddings.py \
+    --model_type $MODEL_TYPE \
+    --model_path $MODEL_PATH \
+    --output_dir=$OUTPUT_DIR 
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/data_preprocess/preprocess_validation_text_embeddings.py \
+    --model_type $MODEL_TYPE \
+    --model_path $MODEL_PATH \
+    --output_dir=$OUTPUT_DIR \
+    --validation_prompt_txt $VALIDATION_PATH

scripts/preprocess/preprocess_hunyuan_data.sh

@@ -0,0 +1,33 @@
+# export WANDB_MODE="offline"
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="data/hunyuan"
+MODEL_TYPE="hunyuan"
+DATA_MERGE_PATH="data/Image-Vid-Finetune-Src/merge.txt"
+OUTPUT_DIR="data/Image-Vid-Finetune-HunYuan"
+VALIDATION_PATH="assets/prompt.txt"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/data_preprocess/preprocess_vae_latents.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --train_batch_size=1 \
+    --max_height=480 \
+    --max_width=848 \
+    --num_frames=93 \
+    --dataloader_num_workers 1 \
+    --output_dir=$OUTPUT_DIR \
+    --model_type $MODEL_TYPE \
+    --train_fps 24 
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/data_preprocess/preprocess_text_embeddings.py \
+    --model_type $MODEL_TYPE \
+    --model_path $MODEL_PATH \
+    --output_dir=$OUTPUT_DIR 
+
+torchrun --nproc_per_node=1 \
+    fastvideo/data_preprocess/preprocess_validation_text_embeddings.py \
+    --model_type $MODEL_TYPE \
+    --model_path $MODEL_PATH \
+    --output_dir=$OUTPUT_DIR \
+    --validation_prompt_txt $VALIDATION_PATH

scripts/preprocess/preprocess_mochi_data.sh

@@ -0,0 +1,33 @@
+# export WANDB_MODE="offline"
+GPU_NUM=1 # 2,4,8
+MODEL_PATH="data/FastMochi-diffusers"
+MODEL_TYPE="mochi"
+DATA_MERGE_PATH="data/Image-Vid-Finetune-Src/merge.txt"
+OUTPUT_DIR="data/Image-Vid-Finetune-Mochi"
+VALIDATION_PATH="assets/prompt.txt"
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/data_preprocess/preprocess_vae_latents.py \
+    --model_path $MODEL_PATH \
+    --data_merge_path $DATA_MERGE_PATH \
+    --train_batch_size=1 \
+    --max_height=480 \
+    --max_width=848 \
+    --num_frames=93 \
+    --dataloader_num_workers 1 \
+    --output_dir=$OUTPUT_DIR \
+    --model_type $MODEL_TYPE \
+    --train_fps 24 
+
+torchrun --nproc_per_node=$GPU_NUM \
+    fastvideo/data_preprocess/preprocess_text_embeddings.py \
+    --model_type $MODEL_TYPE \
+    --model_path $MODEL_PATH \
+    --output_dir=$OUTPUT_DIR 
+
+torchrun --nproc_per_node=1 \
+    fastvideo/data_preprocess/preprocess_validation_text_embeddings.py \
+    --model_type $MODEL_TYPE \
+    --model_path $MODEL_PATH \
+    --output_dir=$OUTPUT_DIR \
+    --validation_prompt_txt $VALIDATION_PATH

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,44 @@
+# 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),
+    '960*544': (960, 544),
+}
+
+MAX_AREA_CONFIGS = {
+    '720*1280': 720 * 1280,
+    '1280*720': 1280 * 720,
+    '480*832': 480 * 832,
+    '832*480': 832 * 480,
+    '960*544': 960 * 544,
+}
+
+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', '960*544'),
+    '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,192 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+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
+
+
+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()
+
+
+def usp_dit_forward(
+    self,
+    x,
+    t,
+    context,
+    seq_len,
+    clip_fea=None,
+    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 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)
+
+    # Context Parallel
+    x = torch.chunk(
+        x, get_sequence_parallel_world_size(),
+        dim=1)[get_sequence_parallel_rank()]
+
+    for block in self.blocks:
+        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)
+    return [u.float() for u in x]
+
+
+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

wan/image2video.py

@@ -0,0 +1,345 @@
+# 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 .modules.clip import CLIPModel
+from .modules.model_infer import WanModel
+from .modules.t5 import T5EncoderModel
+from .modules.vae import WanVAE
+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_usp=False,
+        t5_cpu=False,
+        init_on_cpu=True,
+    ):
+        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_usp (`bool`, *optional*, defaults to False):
+                Enable distribution strategy of USP.
+            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.
+        """
+        self.device = torch.device(f"cuda:{device_id}")
+        self.config = config
+        self.rank = rank
+        self.use_usp = use_usp
+        self.t5_cpu = t5_cpu
+
+        self.num_train_timesteps = config.num_train_timesteps
+        self.param_dtype = config.param_dtype
+
+        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 = WanVAE(
+            vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
+            device=self.device)
+
+        self.clip = CLIPModel(
+            dtype=config.clip_dtype,
+            device=self.device,
+            checkpoint_path=os.path.join(checkpoint_dir,
+                                         config.clip_checkpoint),
+            tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer))
+
+        logging.info(f"Creating WanModel from {checkpoint_dir}")
+        self.model = WanModel.from_pretrained(checkpoint_dir)
+        self.model.eval().requires_grad_(False)
+
+        if t5_fsdp or dit_fsdp or use_usp:
+            init_on_cpu = False
+
+        if use_usp:
+            from xfuser.core.distributed import \
+                get_sequence_parallel_world_size
+
+            from .distributed.xdit_context_parallel import (usp_attn_forward,
+                                                            usp_dit_forward)
+            for block in self.model.blocks:
+                block.self_attn.forward = types.MethodType(
+                    usp_attn_forward, block.self_attn)
+            self.model.forward = types.MethodType(usp_dit_forward, self.model)
+            self.sp_size = get_sequence_parallel_world_size()
+        else:
+            self.sp_size = 1
+
+        if dist.is_initialized():
+            dist.barrier()
+        if dit_fsdp:
+            self.model = shard_fn(self.model)
+        else:
+            if not init_on_cpu:
+                self.model=self.model.to(self.device)
+
+        self.sample_neg_prompt = config.sample_neg_prompt
+
+    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`, *optional*, defaults 5.0):
+                Classifier-free guidance scale. Controls prompt adherence vs. creativity
+            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)
+        """
+        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//4+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=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=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]
+
+        self.clip.model=self.clip.model.to(self.device)
+        clip_context = self.clip.visual([img[:, None, :, :]])
+        if offload_model:
+            self.clip.model=self.clip.model.cpu()
+        torch.cuda.empty_cache()
+        y = self.vae.encode([
+            torch.concat([
+                torch.nn.functional.interpolate(
+                    img[None].cpu(), size=(h, w), mode='bicubic').transpose(
+                        0, 1),
+                # torch.zeros(3, 80, h, w)
+                torch.zeros(3, 48, h, w)
+            ],dim=1).to(self.device)
+        ])[0]
+        y = torch.concat([msk, y])
+
+        @contextmanager
+        def noop_no_sync():
+            yield
+
+        no_sync = getattr(self.model, 'no_sync', noop_no_sync)
+
+        # evaluation mode
+        with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync():
+
+            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]],
+                'clip_fea': clip_context,
+                'seq_len': max_seq_len,
+                'y': [y],
+            }
+
+            arg_null = {
+                'context': context_null,
+                'clip_fea': clip_context,
+                'seq_len': max_seq_len,
+                'y': [y],
+            }
+
+            if offload_model:
+                torch.cuda.empty_cache()
+
+            self.model=self.model.to(self.device)
+            for _, t in enumerate(tqdm(timesteps)):
+                latent_model_input = [latent.to(self.device)]
+                timestep = [t]
+
+                timestep = torch.stack(timestep).to(self.device)
+                # print(timestep)
+                noise_pred_cond = self.model(latent_model_input, t=timestep, **arg_c)[0].to(torch.device('cpu') if offload_model else self.device)
+                # noise_pred_cond = latent_model_input[0]
+                if offload_model:
+                    torch.cuda.empty_cache()
+                noise_pred_uncond = self.model(latent_model_input, t=timestep, **arg_null)[0].to(torch.device('cpu') if offload_model else self.device)
+
+                # noise_pred_uncond = latent_model_input[0]
+                if offload_model:
+                    torch.cuda.empty_cache()
+                noise_pred = noise_pred_uncond + guide_scale * (
+                    noise_pred_cond - noise_pred_uncond)
+
+                latent = latent.to(
+                    torch.device('cpu') if offload_model else self.device)
+
+                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.to(self.device)]
+                del latent_model_input, timestep
+
+            if offload_model:
+                self.model=self.model.cpu()
+                torch.cuda.empty_cache()
+
+            if self.rank == 0:
+                videos = self.vae.decode(x0)
+
+        del noise, latent
+        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/image2video_if_oss.py

@@ -0,0 +1,380 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import sys
+sys.path.append('../OSS')
+from OSS.OSS import search_OSS_video, infer_OSS
+from OSS.model_wrap import _WrappedModel_Wan
+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 .modules.clip import CLIPModel
+from .modules.model_infer import WanModel
+from .modules.t5 import T5EncoderModel
+from .modules.vae import WanVAE
+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_usp=False,
+        t5_cpu=False,
+        init_on_cpu=True,
+    ):
+        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_usp (`bool`, *optional*, defaults to False):
+                Enable distribution strategy of USP.
+            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.
+        """
+        self.device = torch.device(f"cuda:{device_id}")
+        self.config = config
+        self.rank = rank
+        self.use_usp = use_usp
+        self.t5_cpu = t5_cpu
+
+        self.num_train_timesteps = config.num_train_timesteps
+        self.param_dtype = config.param_dtype
+
+        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 = WanVAE(
+            vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
+            device=self.device)
+
+        self.clip = CLIPModel(
+            dtype=config.clip_dtype,
+            device=self.device,
+            checkpoint_path=os.path.join(checkpoint_dir,
+                                         config.clip_checkpoint),
+            tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer))
+
+        logging.info(f"Creating WanModel from {checkpoint_dir}")
+        self.model = WanModel.from_pretrained(checkpoint_dir)
+        self.model.eval().requires_grad_(False)
+
+        if t5_fsdp or dit_fsdp or use_usp:
+            init_on_cpu = False
+
+        if use_usp:
+            from xfuser.core.distributed import \
+                get_sequence_parallel_world_size
+
+            from .distributed.xdit_context_parallel import (usp_attn_forward,
+                                                            usp_dit_forward)
+            for block in self.model.blocks:
+                block.self_attn.forward = types.MethodType(
+                    usp_attn_forward, block.self_attn)
+            self.model.forward = types.MethodType(usp_dit_forward, self.model)
+            self.sp_size = get_sequence_parallel_world_size()
+        else:
+            self.sp_size = 1
+
+        if dist.is_initialized():
+            dist.barrier()
+        if dit_fsdp:
+            self.model = shard_fn(self.model)
+        else:
+            if not init_on_cpu:
+                self.model=self.model.to(self.device)
+
+        self.sample_neg_prompt = config.sample_neg_prompt
+
+    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,speed=0):
+        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`, *optional*, defaults 5.0):
+                Classifier-free guidance scale. Controls prompt adherence vs. creativity
+            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)
+        """
+        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//4+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=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=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]
+
+        self.clip.model=self.clip.model.to(self.device)
+        clip_context = self.clip.visual([img[:, None, :, :]])
+        if offload_model:
+            self.clip.model=self.clip.model.cpu()
+        torch.cuda.empty_cache()
+        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 = getattr(self.model, 'no_sync', noop_no_sync)
+
+        # evaluation mode
+        with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync():
+            # sample videos
+            latents =latent = noise
+
+            arg_c = {
+                'context': [context[0]],
+                'clip_fea': clip_context,
+                'seq_len': max_seq_len,
+                'y': [y],
+            }
+
+            arg_null = {
+                'context': context_null,
+                'clip_fea': clip_context,
+                'seq_len': max_seq_len,
+                'y': [y],
+            }
+
+            if offload_model:
+                torch.cuda.empty_cache()
+
+            self.model=self.model.to(self.device)
+            if speed==0:
+                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.")
+                for _, t in enumerate(tqdm(timesteps)):
+                    latent_model_input = [latent.to(self.device)]
+                    timestep = [t]
+
+                    timestep = torch.stack(timestep).to(self.device)
+                    # print(timestep)
+                    noise_pred_cond = self.model(latent_model_input, t=timestep, **arg_c)[0].to(torch.device('cpu') if offload_model else self.device)
+                    # noise_pred_cond = latent_model_input[0]
+                    if offload_model:
+                        torch.cuda.empty_cache()
+                    noise_pred_uncond = self.model(latent_model_input, t=timestep, **arg_null)[0].to(torch.device('cpu') if offload_model else self.device)
+
+                    # noise_pred_uncond = latent_model_input[0]
+                    if offload_model:
+                        torch.cuda.empty_cache()
+                    noise_pred = noise_pred_uncond + guide_scale * (
+                        noise_pred_cond - noise_pred_uncond)
+
+                    latent = latent.to(
+                        torch.device('cpu') if offload_model else self.device)
+
+                    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.to(self.device)]
+                    del latent_model_input, timestep
+            else:
+                n_ts=96
+                if sample_solver == 'unipc':
+                    sample_scheduler = FlowUniPCMultistepScheduler(
+                        num_train_timesteps=self.num_train_timesteps,
+                        shift=1,
+                        use_dynamic_shifting=False)
+                    sample_scheduler.set_timesteps(
+                        n_ts, 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(n_ts, shift)
+                    timesteps, _ = retrieve_timesteps(
+                        sample_scheduler,
+                        device=self.device,
+                        sigmas=sampling_sigmas)
+                else:
+                    raise NotImplementedError("Unsupported solver.")
+                # pre-process
+                model = _WrappedModel_Wan(self.model, timesteps, self.num_train_timesteps, context_null, guide_scale)
+                model_kwargs = {
+                    'seq_len': max_seq_len,
+                    'y': [y],
+                    'clip_fea': clip_context,
+                }
+                latents = latents.unsqueeze(0)
+
+                oss_steps = [2, 6, 14, 28, 44, 56, 66, 74, 79, 84, 87, 90, 93, 94, 95, 96]  ####oss544Pmed96-16
+                x0 = infer_OSS(oss_steps, model, latents, context, self.device, model_kwargs=model_kwargs)
+
+            if offload_model:
+                self.model=self.model.cpu()
+                torch.cuda.empty_cache()
+
+            if self.rank == 0:
+                videos = self.vae.decode(x0)
+
+        del noise, latent, latents
+        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/image2video_mdinfer_oss_stu.py

@@ -0,0 +1,454 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import sys
+sys.path.append('../OSS')
+from OSS.OSS import search_OSS_video, infer_OSS
+from OSS.model_wrap import _WrappedModel_Wan
+import gc
+import logging
+import math
+import os
+import pdb
+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 .modules.clip import CLIPModel
+from .modules.model_infer import WanModel
+from .modules.t5 import T5EncoderModel
+from .modules.vae import WanVAE
+# from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler,get_sampling_sigmas, retrieve_timesteps)
+from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler)
+from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+
+from diffusers import FlowMatchEulerDiscreteScheduler
+
+import inspect
+import math
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+import numpy as np
+import random
+def set_seed(seed):
+    if seed == -1:
+        seed = random.randint(0, 1000000)
+    seed = int(seed)
+    random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+class FlowMatchScheduler():
+
+    def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.set_timesteps(num_inference_steps)
+
+    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False, shift=None):
+        if shift is not None:
+            self.shift = shift
+        sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength
+        if self.extra_one_step:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]
+        else:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps)
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        if training:
+            x = self.timesteps
+            y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2)
+            y_shifted = y - y.min()
+            bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum())
+            self.linear_timesteps_weights = bsmntw_weighing
+
+    def step(self, model_output, timestep, sample, to_final=False):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        if to_final or timestep_id + 1 >= len(self.timesteps):
+            sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1]
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        return prev_sample
+
+    def return_to_timestep(self, timestep, sample, sample_stablized):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        model_output = (sample - sample_stablized) / sigma
+        return model_output
+
+    def add_noise(self, original_samples, noise, timestep):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample
+
+    def training_target(self, sample, noise, timestep):
+        target = noise - sample
+        return target
+
+    def training_weight(self, timestep):
+        timestep_id = torch.argmin((self.timesteps - timestep.to(self.timesteps.device)).abs())
+        weights = self.linear_timesteps_weights[timestep_id]
+        return weights
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+class WanI2V:
+
+    def __init__(
+        self,
+        config,
+        checkpoint_dir,
+        device_id=0,
+        rank=0,
+        t5_fsdp=False,
+        dit_fsdp=False,
+        use_usp=False,
+        t5_cpu=False,
+        init_on_cpu=True,
+    ):
+        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_usp (`bool`, *optional*, defaults to False):
+                Enable distribution strategy of USP.
+            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.
+        """
+        self.device = torch.device(f"cuda:{device_id}")
+        self.config = config
+        self.rank = rank
+        self.use_usp = use_usp
+        self.t5_cpu = t5_cpu
+        self.scheduler =FlowMatchScheduler(shift=5, sigma_min=0.0, extra_one_step=True)
+        # self.scheduler =FlowMatchScheduler(shift=17, sigma_min=0.0, extra_one_step=True)
+        self.num_train_timesteps = config.num_train_timesteps
+        self.param_dtype = config.param_dtype
+
+        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 = WanVAE(
+            vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
+            device=self.device)
+
+        self.clip = CLIPModel(
+            dtype=config.clip_dtype,
+            device=self.device,
+            checkpoint_path=os.path.join(checkpoint_dir,config.clip_checkpoint),
+            tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer))
+
+        logging.info(f"Creating WanModel from {checkpoint_dir}")
+        self.model = WanModel.from_pretrained(checkpoint_dir)
+        self.model.eval().requires_grad_(False)
+
+        if t5_fsdp or dit_fsdp or use_usp:
+            init_on_cpu = False
+
+        if use_usp:
+            from xfuser.core.distributed import \
+                get_sequence_parallel_world_size
+
+            from .distributed.xdit_context_parallel import (usp_attn_forward,usp_dit_forward)
+            for block in self.model.blocks:
+                block.self_attn.forward = types.MethodType(
+                    usp_attn_forward, block.self_attn)
+            self.model.forward = types.MethodType(usp_dit_forward, self.model)
+            self.sp_size = get_sequence_parallel_world_size()
+        else:
+            self.sp_size = 1
+
+        if dist.is_initialized():
+            dist.barrier()
+        if dit_fsdp:
+            self.model = shard_fn(self.model)
+        else:
+            if not init_on_cpu:
+                self.model=self.model.to(self.device)
+
+        self.sample_neg_prompt = config.sample_neg_prompt
+
+
+    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,
+
+                 student_steps=20,
+                 norm=2,
+                 frame_type="all",
+                 channel_type="all",
+                 random_channel=False,
+                 ):
+        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`, *optional*, defaults 5.0):
+                Classifier-free guidance scale. Controls prompt adherence vs. creativity
+            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)
+        """
+        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)
+        if seed >= 0:
+            set_seed(seed)
+        seed_g = torch.Generator(device=self.device)
+        seed_g.manual_seed(seed)
+        noise = torch.randn(
+            16,
+            F//4+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=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=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]
+
+        self.clip.model=self.clip.model.to(self.device)
+        clip_context = self.clip.visual([img[:, None, :, :]])
+        if offload_model:
+            self.clip.model=self.clip.model.cpu()
+        torch.cuda.empty_cache()
+        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 = getattr(self.model, 'no_sync', noop_no_sync)
+
+        # sampling_steps=10
+        # evaluation mode
+        with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync():
+            device = self.device
+            num_inference_steps=sampling_steps
+            self.scheduler.set_timesteps(num_inference_steps, 1.0, shift=5.0)
+
+            # sample videos
+            latents = noise
+            if offload_model:
+                torch.cuda.empty_cache()
+
+            self.model=self.model.to(self.device)
+
+            # pre-process
+            model = _WrappedModel_Wan(self.model, self.scheduler.timesteps, self.num_train_timesteps, context_null, guide_scale)
+            model_kwargs = {
+                'seq_len': max_seq_len,
+                'y': [y],
+                'clip_fea': clip_context,
+            }
+            latents = latents.unsqueeze(0)
+
+            oss_steps=[2, 6, 14, 28, 44, 56, 66, 74, 79, 84, 87, 90, 93, 94, 95, 96]####oss544Pmed96-16
+            # oss_steps=[2, 5, 11, 21, 36, 49, 63, 71, 80, 84, 87, 89, 92, 94, 95, 96]####oss544Pmed96-16
+            latents_oss = infer_OSS(oss_steps, model, latents, context, self.device, model_kwargs=model_kwargs)
+
+            x0_oss = latents_oss
+
+            if offload_model:
+                self.model.cpu()
+                torch.cuda.empty_cache()
+            if self.rank == 0:
+                videos_oss = self.vae.decode(x0_oss)
+
+        del noise, latents
+        # del self.scheduler
+        if offload_model:
+            gc.collect()
+            torch.cuda.synchronize()
+        if dist.is_initialized():
+            dist.barrier()
+
+        return videos_oss[0] if self.rank == 0 else None

wan/image2video_mdinfer_oss_tea.py

@@ -0,0 +1,513 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import sys,os
+sys.path.append('../OSS')
+from OSS.OSS import search_OSS_video, infer_OSS
+from OSS.model_wrap import _WrappedModel_Wan
+import gc
+import logging
+import math
+import os
+import pdb
+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 .modules.clip import CLIPModel
+from .modules.model_infer import WanModel
+from .modules.t5 import T5EncoderModel
+from .modules.vae import WanVAE
+# from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler,get_sampling_sigmas, retrieve_timesteps)
+from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler)
+from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+
+from diffusers import FlowMatchEulerDiscreteScheduler
+
+import inspect
+import math
+from typing import Callable, Dict, List, Optional, Tuple, Union
+
+import torch
+import numpy as np
+import random
+def set_seed(seed):
+    if seed == -1:
+        seed = random.randint(0, 1000000)
+    seed = int(seed)
+    random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+class FlowMatchScheduler():
+
+    def __init__(self, num_inference_steps=100, num_train_timesteps=1000, shift=3.0, sigma_max=1.0, sigma_min=0.003 / 1.002, inverse_timesteps=False, extra_one_step=False, reverse_sigmas=False):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_max = sigma_max
+        self.sigma_min = sigma_min
+        self.inverse_timesteps = inverse_timesteps
+        self.extra_one_step = extra_one_step
+        self.reverse_sigmas = reverse_sigmas
+        self.set_timesteps(num_inference_steps)
+
+    def set_timesteps(self, num_inference_steps=100, denoising_strength=1.0, training=False, shift=None):
+        if shift is not None:
+            self.shift = shift
+        sigma_start = self.sigma_min + (self.sigma_max - self.sigma_min) * denoising_strength
+        if self.extra_one_step:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps + 1)[:-1]
+        else:
+            self.sigmas = torch.linspace(sigma_start, self.sigma_min, num_inference_steps)
+        if self.inverse_timesteps:
+            self.sigmas = torch.flip(self.sigmas, dims=[0])
+        self.sigmas = self.shift * self.sigmas / (1 + (self.shift - 1) * self.sigmas)
+        if self.reverse_sigmas:
+            self.sigmas = 1 - self.sigmas
+        self.timesteps = self.sigmas * self.num_train_timesteps
+        if training:
+            x = self.timesteps
+            y = torch.exp(-2 * ((x - num_inference_steps / 2) / num_inference_steps) ** 2)
+            y_shifted = y - y.min()
+            bsmntw_weighing = y_shifted * (num_inference_steps / y_shifted.sum())
+            self.linear_timesteps_weights = bsmntw_weighing
+
+    def step(self, model_output, timestep, sample, to_final=False):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        if to_final or timestep_id + 1 >= len(self.timesteps):
+            sigma_ = 1 if (self.inverse_timesteps or self.reverse_sigmas) else 0
+        else:
+            sigma_ = self.sigmas[timestep_id + 1]
+        prev_sample = sample + model_output * (sigma_ - sigma)
+        return prev_sample
+
+    def return_to_timestep(self, timestep, sample, sample_stablized):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        model_output = (sample - sample_stablized) / sigma
+        return model_output
+
+    def add_noise(self, original_samples, noise, timestep):
+        if isinstance(timestep, torch.Tensor):
+            timestep = timestep.cpu()
+        timestep_id = torch.argmin((self.timesteps - timestep).abs())
+        sigma = self.sigmas[timestep_id]
+        sample = (1 - sigma) * original_samples + sigma * noise
+        return sample
+
+    def training_target(self, sample, noise, timestep):
+        target = noise - sample
+        return target
+
+    def training_weight(self, timestep):
+        timestep_id = torch.argmin((self.timesteps - timestep.to(self.timesteps.device)).abs())
+        weights = self.linear_timesteps_weights[timestep_id]
+        return weights
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    r"""
+    Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+    custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+    Args:
+        scheduler (`SchedulerMixin`):
+            The scheduler to get timesteps from.
+        num_inference_steps (`int`):
+            The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+            must be `None`.
+        device (`str` or `torch.device`, *optional*):
+            The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+        timesteps (`List[int]`, *optional*):
+            Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+            `num_inference_steps` and `sigmas` must be `None`.
+        sigmas (`List[float]`, *optional*):
+            Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+            `num_inference_steps` and `timesteps` must be `None`.
+
+    Returns:
+        `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+        second element is the number of inference steps.
+    """
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accepts_timesteps:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" timestep schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+        if not accept_sigmas:
+            raise ValueError(
+                f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+                f" sigmas schedules. Please check whether you are using the correct scheduler."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+class WanI2V:
+
+    def __init__(
+        self,
+        config,
+        checkpoint_dir,
+        device_id=0,
+        rank=0,
+        t5_fsdp=False,
+        dit_fsdp=False,
+        use_usp=False,
+        t5_cpu=False,
+        init_on_cpu=True,
+    ):
+        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_usp (`bool`, *optional*, defaults to False):
+                Enable distribution strategy of USP.
+            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.
+        """
+        self.device = torch.device(f"cuda:{device_id}")
+        self.config = config
+        self.rank = rank
+        self.use_usp = use_usp
+        self.t5_cpu = t5_cpu
+        self.scheduler =FlowMatchScheduler(shift=5, sigma_min=0.0, extra_one_step=True)
+        # self.scheduler =FlowMatchScheduler(shift=17, sigma_min=0.0, extra_one_step=True)
+        self.num_train_timesteps = config.num_train_timesteps
+        self.param_dtype = config.param_dtype
+
+        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 = WanVAE(
+            vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
+            device=self.device)
+
+        self.clip = CLIPModel(
+            dtype=config.clip_dtype,
+            device=self.device,
+            checkpoint_path=os.path.join(checkpoint_dir,config.clip_checkpoint),
+            tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer))
+
+        logging.info(f"Creating WanModel from {checkpoint_dir}")
+        self.model = WanModel.from_pretrained(checkpoint_dir)
+        self.model.eval().requires_grad_(False)
+
+        if t5_fsdp or dit_fsdp or use_usp:
+            init_on_cpu = False
+
+        if use_usp:
+            from xfuser.core.distributed import \
+                get_sequence_parallel_world_size
+
+            from .distributed.xdit_context_parallel import (usp_attn_forward,usp_dit_forward)
+            for block in self.model.blocks:
+                block.self_attn.forward = types.MethodType(
+                    usp_attn_forward, block.self_attn)
+            self.model.forward = types.MethodType(usp_dit_forward, self.model)
+            self.sp_size = get_sequence_parallel_world_size()
+        else:
+            self.sp_size = 1
+
+        if dist.is_initialized():
+            dist.barrier()
+        if dit_fsdp:
+            self.model = shard_fn(self.model)
+        else:
+            if not init_on_cpu:
+                self.model=self.model.to(self.device)
+
+        self.sample_neg_prompt = config.sample_neg_prompt
+
+
+    def generate(self,
+                 args,
+                 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,
+
+                 student_steps=20,
+                 norm=2,
+                 frame_type="all",
+                 channel_type="all",
+                 random_channel=False,
+                 ):
+        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`, *optional*, defaults 5.0):
+                Classifier-free guidance scale. Controls prompt adherence vs. creativity
+            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)
+        """
+        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)
+        if seed >= 0:
+            set_seed(seed)
+        seed_g = torch.Generator(device=self.device)
+        seed_g.manual_seed(seed)
+        noise = torch.randn(
+            16,
+            F//4+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=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=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]
+
+        self.clip.model=self.clip.model.to(self.device)
+        clip_context = self.clip.visual([img[:, None, :, :]])
+        if offload_model:
+            self.clip.model=self.clip.model.cpu()
+        torch.cuda.empty_cache()
+        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 = getattr(self.model, 'no_sync', noop_no_sync)
+
+        # sampling_steps=10
+        # evaluation mode
+        with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync():
+            device = self.device
+            num_inference_steps=sampling_steps
+            self.scheduler.set_timesteps(num_inference_steps, 1.0, shift=5.0)
+
+            # sample videos
+            latents = noise
+            if offload_model:
+                torch.cuda.empty_cache()
+
+            self.model=self.model.to(self.device)
+
+
+
+            # arg_c = {
+            #     'context': [context[0]],
+            #     'clip_fea': clip_context,
+            #     'seq_len': max_seq_len,
+            #     'y': [y],
+            # }
+            #
+            # arg_null = {
+            #     'context': context_null,
+            #     'clip_fea': clip_context,
+            #     'seq_len': max_seq_len,
+            #     'y': [y],
+            # }
+
+            # pre-process
+            model = _WrappedModel_Wan(self.model, self.scheduler.timesteps, self.num_train_timesteps, context_null, guide_scale)
+            model_kwargs = {
+                'seq_len': max_seq_len,
+                'y': [y],
+                'clip_fea': clip_context,
+            }
+            B = 1
+            # latents = latents[0].unsqueeze(0)
+            latents = latents.unsqueeze(0)
+
+            oss_steps = search_OSS_video(model, latents, B, context, self.device, teacher_steps=sampling_steps, student_steps=student_steps, norm=norm, model_kwargs=model_kwargs, frame_type=frame_type, channel_type=channel_type, random_channel=random_channel)
+            latents_oss = infer_OSS(oss_steps, model, latents, context, self.device, model_kwargs=model_kwargs)
+
+            with open("%s.txt"%args.save_file,"w")as f:f.write(str(oss_steps))
+            os._exit(2333)
+            # pdb.set_trace()
+            # teacher video
+            teacher_steps = list(range(1, sampling_steps+1))
+            latents_tea = infer_OSS(teacher_steps, model, latents, context, self.device, model_kwargs=model_kwargs)
+
+            x0_oss = latents_oss
+            x0_tea = latents_tea
+
+            if offload_model:
+                self.model.cpu()
+                torch.cuda.empty_cache()
+            if self.rank == 0:
+                videos_oss = self.vae.decode(x0_oss)
+                videos_tea = self.vae.decode(x0_tea)
+
+            # for idx, t in enumerate(tqdm(self.scheduler.timesteps)):
+            #     latent_model_input = [latent.to(self.device)]
+            #     timestep = [t]
+            #
+            #     timestep = torch.stack(timestep).to(self.device)
+            #     noise_pred_cond = self.model(latent_model_input, t=timestep, **arg_c)[0].to(torch.device('cpu') if offload_model else self.device)
+            #     if offload_model:
+            #         torch.cuda.empty_cache()
+            #     noise_pred_uncond = self.model(latent_model_input, t=timestep, **arg_null)[0].to(torch.device('cpu') if offload_model else self.device)
+            #     if offload_model:
+            #         torch.cuda.empty_cache()
+            #     noise_pred = noise_pred_uncond + guide_scale * (noise_pred_cond - noise_pred_uncond)
+            #     # noise_pred = noise_pred_cond
+            #     latent = latent.to(torch.device('cpu') if offload_model else self.device)
+            #
+            #     # latents = self.scheduler.step(noise_pred, self.scheduler.timesteps[progress_id], latents)
+            #     temp_x0 = self.scheduler.step(
+            #         noise_pred.unsqueeze(0),
+            #         self.scheduler.timesteps[idx],
+            #         latent.unsqueeze(0))[0]
+            #     latent = temp_x0.squeeze(0)
+            #
+            #     x0 = [latent.to(self.device)]
+            #     del latent_model_input, timestep
+            #
+            # if offload_model:
+            #     self.model=self.model.cpu()
+            #     torch.cuda.empty_cache()
+            #
+            # if self.rank == 0:
+            #     videos = self.vae.decode(x0)
+
+        del noise, latents
+        # del self.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,16 @@
+from .attention import flash_attention
+from .model import WanModel
+from .t5 import T5Decoder, T5Encoder, T5EncoderModel, T5Model
+from .tokenizers import HuggingfaceTokenizer
+from .vae import WanVAE
+
+__all__ = [
+    'WanVAE',
+    'WanModel',
+    'T5Model',
+    'T5Encoder',
+    'T5Decoder',
+    'T5EncoderModel',
+    'HuggingfaceTokenizer',
+    'flash_attention',
+]

wan/modules/attention.py

@@ -0,0 +1,208 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+
+try:
+    import flash_attn_interface
+    FLASH_ATTN_3_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_3_AVAILABLE = False
+
+try:
+    import flash_attn
+    FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+    FLASH_ATTN_2_AVAILABLE = False
+
+import warnings
+
+__all__ = [
+    'flash_attention',
+    'attention',
+]
+
+# try:
+#     # from sageattention import sageattn_varlen#sage1
+#     from sageattention import sageattn  # sage2
+#
+#     print("using sageattn2")
+#
+#
+#     @torch.compiler.disable()
+#     def sageattn_wrapper(
+#             q, k, v,
+#             attention_length
+#     ):
+#         padding_length = q.shape[0] - attention_length
+#         q = q[:attention_length, :, :].unsqueeze(0)
+#         k = k[:attention_length, :, :].unsqueeze(0)
+#         v = v[:attention_length, :, :].unsqueeze(0)
+#
+#         o = sageattn(q, k, v, tensor_layout="NHD").squeeze(0)
+#         if padding_length > 0:
+#             o = torch.cat([o, torch.empty((padding_length, *o.shape[-2:]), dtype=o.dtype, device=o.device)], 0)
+#
+#         return o
+# except:
+#     sageattn_wrapper=None
+#     traceback.print_exc()
+
+def flash_attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=True,#False
+    dtype=torch.bfloat16,
+    version=None,
+):
+    """
+    q:              [B, Lq, Nq, C1].
+    k:              [B, Lk, Nk, C1].
+    v:              [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+    q_lens:         [B].
+    k_lens:         [B].
+    dropout_p:      float. Dropout probability.
+    softmax_scale:  float. The scaling of QK^T before applying softmax.
+    causal:         bool. Whether to apply causal attention mask.
+    window_size:    (left right). If not (-1, -1), apply sliding window local attention.
+    deterministic:  bool. If True, slightly slower and uses more memory.
+    dtype:          torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+    """
+    half_dtypes = (torch.float16, torch.bfloat16)
+    assert dtype in half_dtypes
+    assert q.device.type == 'cuda' and q.size(-1) <= 256
+
+    # params
+    b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+    def half(x):
+        return x if x.dtype in half_dtypes else x.to(dtype)
+
+    # preprocess query
+    if q_lens is None:
+        q = half(q.flatten(0, 1))
+        q_lens = torch.tensor(
+            [lq] * b, dtype=torch.int32).to(
+                device=q.device, non_blocking=True)
+    else:
+        q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
+
+    # preprocess key, value
+    if k_lens is None:
+        k = half(k.flatten(0, 1))
+        v = half(v.flatten(0, 1))
+        k_lens = torch.tensor(
+            [lk] * b, dtype=torch.int32).to(
+                device=k.device, non_blocking=True)
+    else:
+        k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
+        v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
+
+    q = q.to(v.dtype)
+    k = k.to(v.dtype)
+
+    if q_scale is not None:
+        q = q * q_scale
+
+    # if type(sageattn_wrapper)!=type(None):
+    #     x = sageattn_wrapper(q, k, v, lq).unsqueeze(0)
+    #     return x.type(out_dtype)
+
+    if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+        warnings.warn(
+            'Flash attention 3 is not available, use flash attention 2 instead.'
+        )
+
+    # apply attention
+    if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+        # Note: dropout_p, window_size are not supported in FA3 now.
+        x = flash_attn_interface.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            seqused_q=None,
+            seqused_k=None,
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            deterministic=deterministic)[0].unflatten(0, (b, lq))
+    else:
+        assert FLASH_ATTN_2_AVAILABLE
+        x = flash_attn.flash_attn_varlen_func(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+                0, dtype=torch.int32).to(q.device, non_blocking=True),
+            max_seqlen_q=lq,
+            max_seqlen_k=lk,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic).unflatten(0, (b, lq))
+
+    # output
+    return x.type(out_dtype)
+
+
+def attention(
+    q,
+    k,
+    v,
+    q_lens=None,
+    k_lens=None,
+    dropout_p=0.,
+    softmax_scale=None,
+    q_scale=None,
+    causal=False,
+    window_size=(-1, -1),
+    deterministic=True,#False
+    dtype=torch.bfloat16,
+    fa_version=None,
+):
+    if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+        return flash_attention(
+            q=q,
+            k=k,
+            v=v,
+            q_lens=q_lens,
+            k_lens=k_lens,
+            dropout_p=dropout_p,
+            softmax_scale=softmax_scale,
+            q_scale=q_scale,
+            causal=causal,
+            window_size=window_size,
+            deterministic=deterministic,
+            dtype=dtype,
+            version=fa_version,
+        )
+    else:
+        if q_lens is not None or k_lens is not None:
+            warnings.warn(
+                'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.'
+            )
+        attn_mask = None
+
+        q = q.transpose(1, 2).to(dtype)
+        k = k.transpose(1, 2).to(dtype)
+        v = v.transpose(1, 2).to(dtype)
+
+        out = torch.nn.functional.scaled_dot_product_attention(
+            q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p)
+
+        out = out.transpose(1, 2).contiguous()
+        return out