mirror of
https://github.com/Wan-Video/Wan2.1.git
synced 2025-06-06 23:34:53 +00:00
1207 lines
47 KiB
Python
1207 lines
47 KiB
Python
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||
import argparse
|
||
import logging
|
||
import os
|
||
import sys
|
||
import warnings
|
||
from tqdm import tqdm
|
||
from datetime import datetime
|
||
|
||
warnings.filterwarnings('ignore')
|
||
|
||
import random
|
||
import torch
|
||
import torch.distributed as dist
|
||
from PIL import Image
|
||
import torchvision.transforms.functional as TF
|
||
import torch.cuda.amp as amp
|
||
import numpy as np
|
||
import math
|
||
|
||
import wan
|
||
from wan.configs import MAX_AREA_CONFIGS, SIZE_CONFIGS, SUPPORTED_SIZES, WAN_CONFIGS
|
||
from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander
|
||
from wan.utils.utils import cache_image, cache_video, str2bool
|
||
|
||
import gc
|
||
from contextlib import contextmanager
|
||
from wan.modules.model import sinusoidal_embedding_1d
|
||
from wan.utils.fm_solvers import (FlowDPMSolverMultistepScheduler,
|
||
get_sampling_sigmas, retrieve_timesteps)
|
||
from wan.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
|
||
|
||
|
||
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",
|
||
},
|
||
"flf2v-14B": {
|
||
"prompt":
|
||
"CG动画风格，一只蓝色的小鸟从地面起飞，煽动翅膀。小鸟羽毛细腻，胸前有独特的花纹，背景是蓝天白云，阳光明媚。镜跟随小鸟向上移动，展现出小鸟飞翔的姿态和天空的广阔。近景，仰视视角。",
|
||
"first_frame":
|
||
"examples/flf2v_input_first_frame.png",
|
||
"last_frame":
|
||
"examples/flf2v_input_last_frame.png",
|
||
},
|
||
"vace-1.3B": {
|
||
"src_ref_images":
|
||
'examples/girl.png,examples/snake.png',
|
||
"prompt":
|
||
"在一个欢乐而充满节日气氛的场景中，穿着鲜艳红色春服的小女孩正与她的可爱卡通蛇嬉戏。她的春服上绣着金色吉祥图案，散发着喜庆的气息，脸上洋溢着灿烂的笑容。蛇身呈现出亮眼的绿色，形状圆润，宽大的眼睛让它显得既友善又幽默。小女孩欢快地用手轻轻抚摸着蛇的头部，共同享受着这温馨的时刻。周围五彩斑斓的灯笼和彩带装饰着环境，阳光透过洒在她们身上，营造出一个充满友爱与幸福的新年氛围。"
|
||
},
|
||
"vace-14B": {
|
||
"src_ref_images":
|
||
'examples/girl.png,examples/snake.png',
|
||
"prompt":
|
||
"在一个欢乐而充满节日气氛的场景中，穿着鲜艳红色春服的小女孩正与她的可爱卡通蛇嬉戏。她的春服上绣着金色吉祥图案，散发着喜庆的气息，脸上洋溢着灿烂的笑容。蛇身呈现出亮眼的绿色，形状圆润，宽大的眼睛让它显得既友善又幽默。小女孩欢快地用手轻轻抚摸着蛇的头部，共同享受着这温馨的时刻。周围五彩斑斓的灯笼和彩带装饰着环境，阳光透过洒在她们身上，营造出一个充满友爱与幸福的新年氛围。"
|
||
}
|
||
}
|
||
|
||
def t2v_generate(self,
|
||
input_prompt,
|
||
size=(1280, 720),
|
||
frame_num=81,
|
||
shift=5.0,
|
||
sample_solver='unipc',
|
||
sampling_steps=50,
|
||
guide_scale=5.0,
|
||
n_prompt="",
|
||
seed=-1,
|
||
offload_model=True):
|
||
r"""
|
||
Generates video frames from text prompt using diffusion process.
|
||
|
||
Args:
|
||
input_prompt (`str`):
|
||
Text prompt for content generation
|
||
size (tupele[`int`], *optional*, defaults to (1280,720)):
|
||
Controls video resolution, (width,height).
|
||
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
|
||
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 size)
|
||
- W: Frame width from size)
|
||
"""
|
||
# preprocess
|
||
F = frame_num
|
||
target_shape = (self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1,
|
||
size[1] // self.vae_stride[1],
|
||
size[0] // self.vae_stride[2])
|
||
|
||
seq_len = math.ceil((target_shape[2] * target_shape[3]) /
|
||
(self.patch_size[1] * self.patch_size[2]) *
|
||
target_shape[1] / self.sp_size) * self.sp_size
|
||
|
||
if n_prompt == "":
|
||
n_prompt = self.sample_neg_prompt
|
||
seed = seed if seed >= 0 else random.randint(0, sys.maxsize)
|
||
seed_g = torch.Generator(device=self.device)
|
||
seed_g.manual_seed(seed)
|
||
|
||
if not self.t5_cpu:
|
||
self.text_encoder.model.to(self.device)
|
||
context = self.text_encoder([input_prompt], self.device)
|
||
context_null = self.text_encoder([n_prompt], self.device)
|
||
if offload_model:
|
||
self.text_encoder.model.cpu()
|
||
else:
|
||
context = self.text_encoder([input_prompt], torch.device('cpu'))
|
||
context_null = self.text_encoder([n_prompt], torch.device('cpu'))
|
||
context = [t.to(self.device) for t in context]
|
||
context_null = [t.to(self.device) for t in context_null]
|
||
|
||
noise = [
|
||
torch.randn(
|
||
target_shape[0],
|
||
target_shape[1],
|
||
target_shape[2],
|
||
target_shape[3],
|
||
dtype=torch.float32,
|
||
device=self.device,
|
||
generator=seed_g)
|
||
]
|
||
|
||
@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
|
||
latents = noise
|
||
|
||
arg_c = {'context': context, 'seq_len': seq_len}
|
||
arg_null = {'context': context_null, 'seq_len': seq_len}
|
||
|
||
for _, t in enumerate(tqdm(timesteps)):
|
||
latent_model_input = latents
|
||
timestep = [t]
|
||
|
||
timestep = torch.stack(timestep)
|
||
|
||
self.model.to(self.device)
|
||
noise_pred_cond = self.model(
|
||
latent_model_input, t=timestep, **arg_c)[0]
|
||
noise_pred_uncond = self.model(
|
||
latent_model_input, t=timestep, **arg_null)[0]
|
||
|
||
noise_pred = noise_pred_uncond + guide_scale * (
|
||
noise_pred_cond - noise_pred_uncond)
|
||
|
||
temp_x0 = sample_scheduler.step(
|
||
noise_pred.unsqueeze(0),
|
||
t,
|
||
latents[0].unsqueeze(0),
|
||
return_dict=False,
|
||
generator=seed_g)[0]
|
||
latents = [temp_x0.squeeze(0)]
|
||
|
||
x0 = latents
|
||
if offload_model:
|
||
self.model.cpu()
|
||
torch.cuda.empty_cache()
|
||
if self.rank == 0:
|
||
videos = self.vae.decode(x0)
|
||
|
||
del noise, 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
|
||
|
||
def i2v_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(
|
||
self.vae.model.z_dim,
|
||
(F - 1) // self.vae_stride[0] + 1,
|
||
lat_h,
|
||
lat_w,
|
||
dtype=torch.float32,
|
||
generator=seed_g,
|
||
device=self.device)
|
||
|
||
msk = torch.ones(1, F, lat_h, lat_w, device=self.device)
|
||
msk[:, 1:] = 0
|
||
msk = torch.concat([
|
||
torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]
|
||
],
|
||
dim=1)
|
||
msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)
|
||
msk = msk.transpose(1, 2)[0]
|
||
|
||
if n_prompt == "":
|
||
n_prompt = self.sample_neg_prompt
|
||
|
||
# preprocess
|
||
if not self.t5_cpu:
|
||
self.text_encoder.model.to(self.device)
|
||
context = self.text_encoder([input_prompt], self.device)
|
||
context_null = self.text_encoder([n_prompt], self.device)
|
||
if offload_model:
|
||
self.text_encoder.model.cpu()
|
||
else:
|
||
context = self.text_encoder([input_prompt], torch.device('cpu'))
|
||
context_null = self.text_encoder([n_prompt], torch.device('cpu'))
|
||
context = [t.to(self.device) for t in context]
|
||
context_null = [t.to(self.device) for t in context_null]
|
||
|
||
self.clip.model.to(self.device)
|
||
clip_context = self.clip.visual([img[:, None, :, :]])
|
||
if offload_model:
|
||
self.clip.model.cpu()
|
||
|
||
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():
|
||
|
||
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],
|
||
# 'cond_flag': True,
|
||
}
|
||
|
||
arg_null = {
|
||
'context': context_null,
|
||
'clip_fea': clip_context,
|
||
'seq_len': max_seq_len,
|
||
'y': [y],
|
||
# 'cond_flag': False,
|
||
}
|
||
|
||
if offload_model:
|
||
torch.cuda.empty_cache()
|
||
|
||
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)
|
||
|
||
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)
|
||
|
||
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.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
|
||
|
||
def teacache_forward(
|
||
self,
|
||
x,
|
||
t,
|
||
context,
|
||
seq_len,
|
||
clip_fea=None,
|
||
y=None,
|
||
):
|
||
r"""
|
||
Forward pass through the diffusion model
|
||
|
||
Args:
|
||
x (List[Tensor]):
|
||
List of input video tensors, each with shape [C_in, F, H, W]
|
||
t (Tensor):
|
||
Diffusion timesteps tensor of shape [B]
|
||
context (List[Tensor]):
|
||
List of text embeddings each with shape [L, C]
|
||
seq_len (`int`):
|
||
Maximum sequence length for positional encoding
|
||
clip_fea (Tensor, *optional*):
|
||
CLIP image features for image-to-video mode
|
||
y (List[Tensor], *optional*):
|
||
Conditional video inputs for image-to-video mode, same shape as x
|
||
|
||
Returns:
|
||
List[Tensor]:
|
||
List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
|
||
"""
|
||
|
||
logging.info("via teacache forward process")
|
||
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)
|
||
|
||
if self.enable_teacache:
|
||
modulated_inp = e0 if self.use_ref_steps else e
|
||
# teacache
|
||
if self.cnt%2==0: # even -> conditon
|
||
self.is_even = True
|
||
if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:
|
||
should_calc_even = True
|
||
self.accumulated_rel_l1_distance_even = 0
|
||
else:
|
||
rescale_func = np.poly1d(self.coefficients)
|
||
self.accumulated_rel_l1_distance_even += rescale_func(((modulated_inp-self.previous_e0_even).abs().mean() / self.previous_e0_even.abs().mean()).cpu().item())
|
||
if self.accumulated_rel_l1_distance_even < self.teacache_thresh:
|
||
should_calc_even = False
|
||
else:
|
||
should_calc_even = True
|
||
self.accumulated_rel_l1_distance_even = 0
|
||
self.previous_e0_even = modulated_inp.clone()
|
||
|
||
else: # odd -> unconditon
|
||
self.is_even = False
|
||
if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:
|
||
should_calc_odd = True
|
||
self.accumulated_rel_l1_distance_odd = 0
|
||
else:
|
||
rescale_func = np.poly1d(self.coefficients)
|
||
self.accumulated_rel_l1_distance_odd += rescale_func(((modulated_inp-self.previous_e0_odd).abs().mean() / self.previous_e0_odd.abs().mean()).cpu().item())
|
||
if self.accumulated_rel_l1_distance_odd < self.teacache_thresh:
|
||
should_calc_odd = False
|
||
else:
|
||
should_calc_odd = True
|
||
self.accumulated_rel_l1_distance_odd = 0
|
||
self.previous_e0_odd = modulated_inp.clone()
|
||
|
||
if self.enable_teacache:
|
||
if self.is_even:
|
||
if not should_calc_even:
|
||
logging.info("use residual estimation for this difusion step")
|
||
x += self.previous_residual_even
|
||
else:
|
||
ori_x = x.clone()
|
||
for block in self.blocks:
|
||
x = block(x, **kwargs)
|
||
self.previous_residual_even = x - ori_x
|
||
else:
|
||
if not should_calc_odd:
|
||
logging.info("use residual estimation for thi8s difusion step")
|
||
x += self.previous_residual_odd
|
||
else:
|
||
ori_x = x.clone()
|
||
for block in self.blocks:
|
||
x = block(x, **kwargs)
|
||
self.previous_residual_odd = x - ori_x
|
||
|
||
else:
|
||
for block in self.blocks:
|
||
x = block(x, **kwargs)
|
||
|
||
# head
|
||
x = self.head(x, e)
|
||
|
||
# unpatchify
|
||
x = self.unpatchify(x, grid_sizes)
|
||
self.cnt += 1
|
||
if self.cnt >= self.num_steps:
|
||
self.cnt = 0
|
||
return [u.float() for u in x]
|
||
|
||
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 = 50
|
||
if "i2v" in args.task:
|
||
args.sample_steps = 40
|
||
|
||
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
|
||
elif "flf2v" in args.task or "vace" in args.task:
|
||
args.sample_shift = 16
|
||
|
||
# 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(
|
||
"--src_video",
|
||
type=str,
|
||
default=None,
|
||
help="The file of the source video. Default None.")
|
||
parser.add_argument(
|
||
"--src_mask",
|
||
type=str,
|
||
default=None,
|
||
help="The file of the source mask. Default None.")
|
||
parser.add_argument(
|
||
"--src_ref_images",
|
||
type=str,
|
||
default=None,
|
||
help="The file list of the source reference images. Separated by ','. Default None."
|
||
)
|
||
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="zh",
|
||
choices=["zh", "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="[image to video] The image to generate the video from.")
|
||
parser.add_argument(
|
||
"--first_frame",
|
||
type=str,
|
||
default=None,
|
||
help="[first-last frame to video] The image (first frame) to generate the video from."
|
||
)
|
||
parser.add_argument(
|
||
"--last_frame",
|
||
type=str,
|
||
default=None,
|
||
help="[first-last frame to video] The image (last frame) to generate the video from."
|
||
)
|
||
parser.add_argument(
|
||
"--sample_solver",
|
||
type=str,
|
||
default='unipc',
|
||
choices=['unipc', 'dpm++'],
|
||
help="The solver used to sample.")
|
||
parser.add_argument(
|
||
"--sample_steps", type=int, default=None, help="The sampling steps.")
|
||
parser.add_argument(
|
||
"--sample_shift",
|
||
type=float,
|
||
default=None,
|
||
help="Sampling shift factor for flow matching schedulers.")
|
||
parser.add_argument(
|
||
"--sample_guide_scale",
|
||
type=float,
|
||
default=5.0,
|
||
help="Classifier free guidance scale.")
|
||
|
||
parser.add_argument(
|
||
"--use_ret_steps",
|
||
action="store_true",
|
||
default=False,
|
||
help=" use ret_steps or not")
|
||
|
||
parser.add_argument(
|
||
"--enable_teacache",
|
||
action="store_true",
|
||
default=False,
|
||
help=" use ret_steps or not")
|
||
|
||
#teacache_thresh
|
||
parser.add_argument(
|
||
"--teacache_thresh",
|
||
type=float,
|
||
default= 0.2,
|
||
help="tea_cache threshold")
|
||
|
||
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 (
|
||
init_distributed_environment,
|
||
initialize_model_parallel,
|
||
)
|
||
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 or "flf2v" 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"`{cfg.num_heads=}` cannot be divided evenly by `{args.ulysses_size=}`."
|
||
|
||
logging.info(f"Generation job args: {args}")
|
||
logging.info(f"Generation model config: {cfg}")
|
||
|
||
if dist.is_initialized():
|
||
base_seed = [args.base_seed] if rank == 0 else [None]
|
||
dist.broadcast_object_list(base_seed, src=0)
|
||
args.base_seed = base_seed[0]
|
||
|
||
if "t2v" in args.task or "t2i" in args.task:
|
||
if args.prompt is None:
|
||
args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
|
||
logging.info(f"Input prompt: {args.prompt}")
|
||
if args.use_prompt_extend:
|
||
logging.info("Extending prompt ...")
|
||
if rank == 0:
|
||
prompt_output = prompt_expander(
|
||
args.prompt,
|
||
tar_lang=args.prompt_extend_target_lang,
|
||
seed=args.base_seed)
|
||
if prompt_output.status == False:
|
||
logging.info(
|
||
f"Extending prompt failed: {prompt_output.message}")
|
||
logging.info("Falling back to original prompt.")
|
||
input_prompt = args.prompt
|
||
else:
|
||
input_prompt = prompt_output.prompt
|
||
input_prompt = [input_prompt]
|
||
else:
|
||
input_prompt = [None]
|
||
if dist.is_initialized():
|
||
dist.broadcast_object_list(input_prompt, src=0)
|
||
args.prompt = input_prompt[0]
|
||
logging.info(f"Extended prompt: {args.prompt}")
|
||
|
||
logging.info("Creating WanT2V pipeline.")
|
||
wan_t2v = wan.WanT2V(
|
||
config=cfg,
|
||
checkpoint_dir=args.ckpt_dir,
|
||
device_id=device,
|
||
rank=rank,
|
||
t5_fsdp=args.t5_fsdp,
|
||
dit_fsdp=args.dit_fsdp,
|
||
use_usp=(args.ulysses_size > 1 or args.ring_size > 1),
|
||
t5_cpu=args.t5_cpu,
|
||
)
|
||
|
||
if args.enable_teacache:
|
||
wan_t2v.__class__.generate = t2v_generate
|
||
wan_t2v.model.__class__.enable_teacache = True
|
||
wan_t2v.model.__class__.forward = teacache_forward
|
||
wan_t2v.model.__class__.cnt = 0
|
||
wan_t2v.model.__class__.num_steps = args.sample_steps*2
|
||
wan_t2v.model.__class__.teacache_thresh = args.teacache_thresh
|
||
wan_t2v.model.__class__.accumulated_rel_l1_distance_even = 0
|
||
wan_t2v.model.__class__.accumulated_rel_l1_distance_odd = 0
|
||
wan_t2v.model.__class__.previous_e0_even = None
|
||
wan_t2v.model.__class__.previous_e0_odd = None
|
||
wan_t2v.model.__class__.previous_residual_even = None
|
||
wan_t2v.model.__class__.previous_residual_odd = None
|
||
wan_t2v.model.__class__.use_ref_steps = args.use_ret_steps
|
||
if args.use_ret_steps:
|
||
if '1.3B' in args.ckpt_dir:
|
||
wan_t2v.model.__class__.coefficients = [-5.21862437e+04, 9.23041404e+03, -5.28275948e+02, 1.36987616e+01, -4.99875664e-02]
|
||
if '14B' in args.ckpt_dir:
|
||
wan_t2v.model.__class__.coefficients = [-3.03318725e+05, 4.90537029e+04, -2.65530556e+03, 5.87365115e+01, -3.15583525e-01]
|
||
wan_t2v.model.__class__.ret_steps = 5*2
|
||
wan_t2v.model.__class__.cutoff_steps = args.sample_steps*2
|
||
else:
|
||
if '1.3B' in args.ckpt_dir:
|
||
wan_t2v.model.__class__.coefficients = [2.39676752e+03, -1.31110545e+03, 2.01331979e+02, -8.29855975e+00, 1.37887774e-01]
|
||
if '14B' in args.ckpt_dir:
|
||
wan_t2v.model.__class__.coefficients = [-5784.54975374, 5449.50911966, -1811.16591783, 256.27178429, -13.02252404]
|
||
wan_t2v.model.__class__.ret_steps = 1*2
|
||
wan_t2v.model.__class__.cutoff_steps = args.sample_steps*2 - 2
|
||
|
||
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)
|
||
|
||
elif "i2v" in args.task:
|
||
if args.prompt is None:
|
||
args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
|
||
if args.image is None:
|
||
args.image = EXAMPLE_PROMPT[args.task]["image"]
|
||
logging.info(f"Input prompt: {args.prompt}")
|
||
logging.info(f"Input image: {args.image}")
|
||
|
||
img = Image.open(args.image).convert("RGB")
|
||
if args.use_prompt_extend:
|
||
logging.info("Extending prompt ...")
|
||
if rank == 0:
|
||
prompt_output = prompt_expander(
|
||
args.prompt,
|
||
tar_lang=args.prompt_extend_target_lang,
|
||
image=img,
|
||
seed=args.base_seed)
|
||
if prompt_output.status == False:
|
||
logging.info(
|
||
f"Extending prompt failed: {prompt_output.message}")
|
||
logging.info("Falling back to original prompt.")
|
||
input_prompt = args.prompt
|
||
else:
|
||
input_prompt = prompt_output.prompt
|
||
input_prompt = [input_prompt]
|
||
else:
|
||
input_prompt = [None]
|
||
if dist.is_initialized():
|
||
dist.broadcast_object_list(input_prompt, src=0)
|
||
args.prompt = input_prompt[0]
|
||
logging.info(f"Extended prompt: {args.prompt}")
|
||
|
||
logging.info("Creating WanI2V pipeline.")
|
||
wan_i2v = wan.WanI2V(
|
||
config=cfg,
|
||
checkpoint_dir=args.ckpt_dir,
|
||
device_id=device,
|
||
rank=rank,
|
||
t5_fsdp=args.t5_fsdp,
|
||
dit_fsdp=args.dit_fsdp,
|
||
use_usp=(args.ulysses_size > 1 or args.ring_size > 1),
|
||
t5_cpu=args.t5_cpu,
|
||
)
|
||
|
||
if args.enable_teacache:
|
||
wan_i2v.__class__.generate = i2v_generate
|
||
wan_i2v.model.__class__.enable_teacache = True
|
||
wan_i2v.model.__class__.forward = teacache_forward
|
||
wan_i2v.model.__class__.cnt = 0
|
||
wan_i2v.model.__class__.num_steps = args.sample_steps*2
|
||
wan_i2v.model.__class__.teacache_thresh = args.teacache_thresh
|
||
wan_i2v.model.__class__.accumulated_rel_l1_distance_even = 0
|
||
wan_i2v.model.__class__.accumulated_rel_l1_distance_odd = 0
|
||
wan_i2v.model.__class__.previous_e0_even = None
|
||
wan_i2v.model.__class__.previous_e0_odd = None
|
||
wan_i2v.model.__class__.previous_residual_even = None
|
||
wan_i2v.model.__class__.previous_residual_odd = None
|
||
wan_i2v.model.__class__.use_ref_steps = args.use_ret_steps
|
||
if args.use_ret_steps:
|
||
if '480P' in args.ckpt_dir:
|
||
wan_i2v.model.__class__.coefficients = [ 2.57151496e+05, -3.54229917e+04, 1.40286849e+03, -1.35890334e+01, 1.32517977e-01]
|
||
if '720P' in args.ckpt_dir:
|
||
wan_i2v.model.__class__.coefficients = [ 8.10705460e+03, 2.13393892e+03, -3.72934672e+02, 1.66203073e+01, -4.17769401e-02]
|
||
wan_i2v.model.__class__.ret_steps = 5*2
|
||
wan_i2v.model.__class__.cutoff_steps = args.sample_steps*2
|
||
else:
|
||
if '480P' in args.ckpt_dir:
|
||
wan_i2v.model.__class__.coefficients = [-3.02331670e+02, 2.23948934e+02, -5.25463970e+01, 5.87348440e+00, -2.01973289e-01]
|
||
if '720P' in args.ckpt_dir:
|
||
wan_i2v.model.__class__.coefficients = [-114.36346466, 65.26524496, -18.82220707, 4.91518089, -0.23412683]
|
||
wan_i2v.model.__class__.ret_steps = 1*2
|
||
wan_i2v.model.__class__.cutoff_steps = args.sample_steps*2 - 2
|
||
|
||
logging.info("Generating video ...")
|
||
video = wan_i2v.generate(
|
||
args.prompt,
|
||
img,
|
||
max_area=MAX_AREA_CONFIGS[args.size],
|
||
frame_num=args.frame_num,
|
||
shift=args.sample_shift,
|
||
sample_solver=args.sample_solver,
|
||
sampling_steps=args.sample_steps,
|
||
guide_scale=args.sample_guide_scale,
|
||
seed=args.base_seed,
|
||
offload_model=args.offload_model)
|
||
elif "flf2v" in args.task:
|
||
if args.prompt is None:
|
||
args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
|
||
if args.first_frame is None or args.last_frame is None:
|
||
args.first_frame = EXAMPLE_PROMPT[args.task]["first_frame"]
|
||
args.last_frame = EXAMPLE_PROMPT[args.task]["last_frame"]
|
||
logging.info(f"Input prompt: {args.prompt}")
|
||
logging.info(f"Input first frame: {args.first_frame}")
|
||
logging.info(f"Input last frame: {args.last_frame}")
|
||
first_frame = Image.open(args.first_frame).convert("RGB")
|
||
last_frame = Image.open(args.last_frame).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=[first_frame, last_frame],
|
||
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 WanFLF2V pipeline.")
|
||
wan_flf2v = wan.WanFLF2V(
|
||
config=cfg,
|
||
checkpoint_dir=args.ckpt_dir,
|
||
device_id=device,
|
||
rank=rank,
|
||
t5_fsdp=args.t5_fsdp,
|
||
dit_fsdp=args.dit_fsdp,
|
||
use_usp=(args.ulysses_size > 1 or args.ring_size > 1),
|
||
t5_cpu=args.t5_cpu,
|
||
)
|
||
|
||
logging.info("Generating video ...")
|
||
video = wan_flf2v.generate(
|
||
args.prompt,
|
||
first_frame,
|
||
last_frame,
|
||
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)
|
||
elif "vace" in args.task:
|
||
if args.prompt is None:
|
||
args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
|
||
args.src_video = EXAMPLE_PROMPT[args.task].get("src_video", None)
|
||
args.src_mask = EXAMPLE_PROMPT[args.task].get("src_mask", None)
|
||
args.src_ref_images = EXAMPLE_PROMPT[args.task].get(
|
||
"src_ref_images", None)
|
||
|
||
logging.info(f"Input prompt: {args.prompt}")
|
||
if args.use_prompt_extend and args.use_prompt_extend != 'plain':
|
||
logging.info("Extending prompt ...")
|
||
if rank == 0:
|
||
prompt = prompt_expander.forward(args.prompt)
|
||
logging.info(
|
||
f"Prompt extended from '{args.prompt}' to '{prompt}'")
|
||
input_prompt = [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 VACE pipeline.")
|
||
wan_vace = wan.WanVace(
|
||
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,
|
||
)
|
||
|
||
src_video, src_mask, src_ref_images = wan_vace.prepare_source(
|
||
[args.src_video], [args.src_mask], [
|
||
None if args.src_ref_images is None else
|
||
args.src_ref_images.split(',')
|
||
], args.frame_num, SIZE_CONFIGS[args.size], device)
|
||
|
||
logging.info(f"Generating video...")
|
||
video = wan_vace.generate(
|
||
args.prompt,
|
||
src_video,
|
||
src_mask,
|
||
src_ref_images,
|
||
size=SIZE_CONFIGS[args.size],
|
||
frame_num=args.frame_num,
|
||
shift=args.sample_shift,
|
||
sample_solver=args.sample_solver,
|
||
sampling_steps=args.sample_steps,
|
||
guide_scale=args.sample_guide_scale,
|
||
seed=args.base_seed,
|
||
offload_model=args.offload_model)
|
||
else:
|
||
raise ValueError(f"Unkown task type: {args.task}")
|
||
|
||
if rank == 0:
|
||
if args.save_file is None:
|
||
formatted_time = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||
formatted_prompt = args.prompt.replace(" ", "_").replace("/",
|
||
"_")[:50]
|
||
suffix = '.png' if "t2i" in args.task else '.mp4'
|
||
args.save_file = f"{args.task}_{args.size.replace('*','x') if sys.platform=='win32' else args.size}_{args.ulysses_size}_{args.ring_size}_{formatted_prompt}_{formatted_time}" + suffix
|
||
|
||
if "t2i" in args.task:
|
||
logging.info(f"Saving generated image to {args.save_file}")
|
||
cache_image(
|
||
tensor=video.squeeze(1)[None],
|
||
save_file=args.save_file,
|
||
nrow=1,
|
||
normalize=True,
|
||
value_range=(-1, 1))
|
||
else:
|
||
logging.info(f"Saving generated video to {args.save_file}")
|
||
cache_video(
|
||
tensor=video[None],
|
||
save_file=args.save_file,
|
||
fps=cfg.sample_fps,
|
||
nrow=1,
|
||
normalize=True,
|
||
value_range=(-1, 1))
|
||
logging.info("Finished.")
|
||
|
||
|
||
if __name__ == "__main__":
|
||
args = _parse_args()
|
||
generate(args)
|