nodes/aigc-apps-VideoX-Fun
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Wan animate inference and training (#365)

Browse Source
This commit is contained in:
Bubbliiiing
2025-11-04 13:36:51 +08:00
committed by GitHub
parent ce7ff41bfb
commit df77df019e
23 changed files with 5139 additions and 80 deletions
Download Patch File Download Diff File Expand all files Collapse all files
config/wan2.2/wan_civitai_animate.yaml
+41
View File
@@ -0,0 +1,41 @@
format: civitai
pipeline: Wan
transformer_additional_kwargs:
transformer_low_noise_model_subpath: ./
transformer_combination_type: "single"
dict_mapping:
in_dim: in_channels
dim: hidden_size
vae_kwargs:
vae_type: "AutoencoderKLWan"
vae_subpath: Wan2.1_VAE.pth
temporal_compression_ratio: 4
spatial_compression_ratio: 8
text_encoder_kwargs:
text_encoder_subpath: models_t5_umt5-xxl-enc-bf16.pth
tokenizer_subpath: google/umt5-xxl
text_length: 512
vocab: 256384
dim: 4096
dim_attn: 4096
dim_ffn: 10240
num_heads: 64
num_layers: 24
num_buckets: 32
shared_pos: False
dropout: 0.0
scheduler_kwargs:
scheduler_subpath: null
num_train_timesteps: 1000
shift: 5.0
use_dynamic_shifting: false
base_shift: 0.5
max_shift: 1.15
base_image_seq_len: 256
max_image_seq_len: 4096
image_encoder_kwargs:
image_encoder_subpath: models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
examples/wan2.2/predict_animate.py
+379
View File
@@ -0,0 +1,379 @@
import os
import sys
import numpy as np
import torch
from diffusers import FlowMatchEulerDiscreteScheduler
from omegaconf import OmegaConf
from PIL import Image
current_file_path = os.path.abspath(__file__)
project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
for project_root in project_roots:
sys.path.insert(0, project_root) if project_root not in sys.path else None
from videox_fun.dist import set_multi_gpus_devices, shard_model
from videox_fun.models import (AutoencoderKLWan, AutoencoderKLWan3_8,
AutoTokenizer, CLIPModel,
Wan2_2Transformer3DModel_Animate,
WanT5EncoderModel)
from videox_fun.models.cache_utils import get_teacache_coefficients
from videox_fun.pipeline import Wan2_2AnimatePipeline
from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
convert_weight_dtype_wrapper,
replace_parameters_by_name)
from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
from videox_fun.utils.utils import (filter_kwargs, get_image,
get_image_to_video_latent,
get_video_to_video_latent,
save_videos_grid)
# GPU memory mode, which can be chosen in [model_full_load, model_full_load_and_qfloat8, model_cpu_offload, model_cpu_offload_and_qfloat8, sequential_cpu_offload].
# model_full_load means that the entire model will be moved to the GPU.
#
# model_full_load_and_qfloat8 means that the entire model will be moved to the GPU,
# and the transformer model has been quantized to float8, which can save more GPU memory.
#
# model_cpu_offload means that the entire model will be moved to the CPU after use, which can save some GPU memory.
#
# model_cpu_offload_and_qfloat8 indicates that the entire model will be moved to the CPU after use,
# and the transformer model has been quantized to float8, which can save more GPU memory.
#
# sequential_cpu_offload means that each layer of the model will be moved to the CPU after use,
# resulting in slower speeds but saving a large amount of GPU memory.
GPU_memory_mode = "sequential_cpu_offload"
# Multi GPUs config
# Please ensure that the product of ulysses_degree and ring_degree equals the number of GPUs used.
# For example, if you are using 8 GPUs, you can set ulysses_degree = 2 and ring_degree = 4.
# If you are using 1 GPU, you can set ulysses_degree = 1 and ring_degree = 1.
ulysses_degree = 1
ring_degree = 1
# Use FSDP to save more GPU memory in multi gpus.
fsdp_dit = False
fsdp_text_encoder = True
# Compile will give a speedup in fixed resolution and need a little GPU memory.
# The compile_dit is not compatible with the fsdp_dit and sequential_cpu_offload.
compile_dit = False
# TeaCache config
enable_teacache = True
# Recommended to be set between 0.05 and 0.30. A larger threshold can cache more steps, speeding up the inference process,
# but it may cause slight differences between the generated content and the original content.
teacache_threshold = 0.10
# The number of steps to skip TeaCache at the beginning of the inference process, which can
# reduce the impact of TeaCache on generated video quality.
num_skip_start_steps = 5
# Whether to offload TeaCache tensors to cpu to save a little bit of GPU memory.
teacache_offload = False
# Skip some cfg steps in inference
# Recommended to be set between 0.00 and 0.25
cfg_skip_ratio = 0
# Riflex config
enable_riflex = False
# Index of intrinsic frequency
riflex_k = 6
# Config and model path
config_path = "config/wan2.2/wan_civitai_animate.yaml"
# model path
model_name = "./models/Diffusion_Transformer/Wan2.2-Animate-14B/"
# Choose the sampler in "Flow", "Flow_Unipc", "Flow_DPM++"
sampler_name = "Flow_Unipc"
# [NOTE]: Noise schedule shift parameter. Affects temporal dynamics.
# Used when the sampler is in "Flow_Unipc", "Flow_DPM++".
shift = 5
# Load pretrained model if need
# The transformer_path is used for low noise model, the transformer_high_path is used for high noise model.
transformer_path = None
transformer_high_path = None
vae_path = None
# Load lora model if need
# The lora_path is used for low noise model, the lora_high_path is used for high noise model.
lora_path = None
lora_high_path = None
src_root_path = "asset/wan_animate/replace/process_results/"
src_pose_path = os.path.join(src_root_path, "src_pose.mp4")
src_face_path = os.path.join(src_root_path, "src_face.mp4")
src_ref_path = os.path.join(src_root_path, "src_ref.png")
src_bg_path = os.path.join(src_root_path, "src_bg.mp4")
src_mask_path = os.path.join(src_root_path, "src_mask.mp4")
# Other params
sample_size = [480, 832]
video_length = 81
fps = 16
# Use torch.float16 if GPU does not support torch.bfloat16
# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
weight_dtype = torch.bfloat16
prompt = "视频中的人在做动作"
negative_prompt = "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
guidance_scale = 4.0
seed = 43
num_inference_steps = 20
# The lora_weight is used for low noise model, the lora_high_weight is used for high noise model.
lora_weight = 0.55
lora_high_weight = 0.55
save_path = "samples/wan-videos-animate"
device = set_multi_gpus_devices(ulysses_degree, ring_degree)
config = OmegaConf.load(config_path)
boundary = config['transformer_additional_kwargs'].get('boundary', 0.875)
transformer = Wan2_2Transformer3DModel_Animate.from_pretrained(
os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
torch_dtype=weight_dtype,
)
if config['transformer_additional_kwargs'].get('transformer_combination_type', 'single') == "moe":
transformer_2 = Wan2_2Transformer3DModel.from_pretrained(
os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')),
transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
low_cpu_mem_usage=True,
torch_dtype=weight_dtype,
)
else:
transformer_2 = None
if transformer_path is not None:
print(f"From checkpoint: {transformer_path}")
if transformer_path.endswith("safetensors"):
from safetensors.torch import load_file, safe_open
state_dict = load_file(transformer_path)
else:
state_dict = torch.load(transformer_path, map_location="cpu")
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
m, u = transformer.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
if transformer_2 is not None:
if transformer_high_path is not None:
print(f"From checkpoint: {transformer_high_path}")
if transformer_high_path.endswith("safetensors"):
from safetensors.torch import load_file, safe_open
state_dict = load_file(transformer_high_path)
else:
state_dict = torch.load(transformer_high_path, map_location="cpu")
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
m, u = transformer_2.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
# Get Vae
Chosen_AutoencoderKL = {
"AutoencoderKLWan": AutoencoderKLWan,
"AutoencoderKLWan3_8": AutoencoderKLWan3_8
}[config['vae_kwargs'].get('vae_type', 'AutoencoderKLWan')]
vae = Chosen_AutoencoderKL.from_pretrained(
os.path.join(model_name, config['vae_kwargs'].get('vae_subpath', 'vae')),
additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
).to(weight_dtype)
if vae_path is not None:
print(f"From checkpoint: {vae_path}")
if vae_path.endswith("safetensors"):
from safetensors.torch import load_file, safe_open
state_dict = load_file(vae_path)
else:
state_dict = torch.load(vae_path, map_location="cpu")
state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
m, u = vae.load_state_dict(state_dict, strict=False)
print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
# Get Tokenizer
tokenizer = AutoTokenizer.from_pretrained(
os.path.join(model_name, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')),
)
# Get Text encoder
text_encoder = WanT5EncoderModel.from_pretrained(
os.path.join(model_name, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
low_cpu_mem_usage=True,
torch_dtype=weight_dtype,
)
text_encoder = text_encoder.eval()
# Get Clip Image Encoder
clip_image_encoder = CLIPModel.from_pretrained(
os.path.join(model_name, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')),
).to(weight_dtype)
clip_image_encoder = clip_image_encoder.eval()
# Get Scheduler
Chosen_Scheduler = scheduler_dict = {
"Flow": FlowMatchEulerDiscreteScheduler,
"Flow_Unipc": FlowUniPCMultistepScheduler,
"Flow_DPM++": FlowDPMSolverMultistepScheduler,
}[sampler_name]
if sampler_name == "Flow_Unipc" or sampler_name == "Flow_DPM++":
config['scheduler_kwargs']['shift'] = 1
scheduler = Chosen_Scheduler(
**filter_kwargs(Chosen_Scheduler, OmegaConf.to_container(config['scheduler_kwargs']))
)
# Get Pipeline
pipeline = Wan2_2AnimatePipeline(
transformer=transformer,
transformer_2=transformer_2,
vae=vae,
tokenizer=tokenizer,
text_encoder=text_encoder,
clip_image_encoder=clip_image_encoder,
scheduler=scheduler,
)
if ulysses_degree > 1 or ring_degree > 1:
from functools import partial
transformer.enable_multi_gpus_inference()
if transformer_2 is not None:
transformer_2.enable_multi_gpus_inference()
if fsdp_dit:
shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
pipeline.transformer = shard_fn(pipeline.transformer)
if transformer_2 is not None:
pipeline.transformer_2 = shard_fn(pipeline.transformer_2)
print("Add FSDP DIT")
if fsdp_text_encoder:
shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
pipeline.text_encoder = shard_fn(pipeline.text_encoder)
print("Add FSDP TEXT ENCODER")
if compile_dit:
for i in range(len(pipeline.transformer.blocks)):
pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
if transformer_2 is not None:
for i in range(len(pipeline.transformer_2.blocks)):
pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
print("Add Compile")
if GPU_memory_mode == "sequential_cpu_offload":
replace_parameters_by_name(transformer, ["modulation",], device=device)
transformer.freqs = transformer.freqs.to(device=device)
if transformer_2 is not None:
replace_parameters_by_name(transformer_2, ["modulation",], device=device)
transformer_2.freqs = transformer_2.freqs.to(device=device)
pipeline.enable_sequential_cpu_offload(device=device)
elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
convert_weight_dtype_wrapper(transformer, weight_dtype)
if transformer_2 is not None:
convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
convert_weight_dtype_wrapper(transformer_2, weight_dtype)
pipeline.enable_model_cpu_offload(device=device)
elif GPU_memory_mode == "model_cpu_offload":
pipeline.enable_model_cpu_offload(device=device)
elif GPU_memory_mode == "model_full_load_and_qfloat8":
convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
convert_weight_dtype_wrapper(transformer, weight_dtype)
if transformer_2 is not None:
convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
convert_weight_dtype_wrapper(transformer_2, weight_dtype)
pipeline.to(device=device)
else:
pipeline.to(device=device)
coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
if coefficients is not None:
print(f"Enable TeaCache with threshold {teacache_threshold} and skip the first {num_skip_start_steps} steps.")
pipeline.transformer.enable_teacache(
coefficients, num_inference_steps, teacache_threshold, num_skip_start_steps=num_skip_start_steps, offload=teacache_offload
)
if transformer_2 is not None:
pipeline.transformer_2.share_teacache(transformer=pipeline.transformer)
if cfg_skip_ratio is not None:
print(f"Enable cfg_skip_ratio {cfg_skip_ratio}.")
pipeline.transformer.enable_cfg_skip(cfg_skip_ratio, num_inference_steps)
if transformer_2 is not None:
pipeline.transformer_2.share_cfg_skip(transformer=pipeline.transformer)
generator = torch.Generator(device=device).manual_seed(seed)
if lora_path is not None:
pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
if transformer_2 is not None:
pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2")
with torch.no_grad():
video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
latent_frames = (video_length - 1) // vae.config.temporal_compression_ratio + 1
if enable_riflex:
pipeline.transformer.enable_riflex(k = riflex_k, L_test = latent_frames)
if transformer_2 is not None:
pipeline.transformer_2.enable_riflex(k = riflex_k, L_test = latent_frames)
pose_video, _, _, _ = get_video_to_video_latent(src_pose_path, video_length=video_length, sample_size=sample_size, fps=fps, ref_image=None)
face_video, _, _, _ = get_video_to_video_latent(src_face_path, video_length=video_length, sample_size=[512, 512], fps=fps, ref_image=None)
ref_image = get_image(src_ref_path)
if os.path.exists(src_bg_path):
bg_video, _, _, _ = get_video_to_video_latent(src_bg_path, video_length=video_length, sample_size=sample_size, fps=fps, ref_image=None)
mask_video, _, _, _ = get_video_to_video_latent(src_mask_path, video_length=video_length, sample_size=sample_size, fps=fps, ref_image=None)
mask_video = mask_video[:, :1]
replace_flag = True
else:
bg_video = None
mask_video = None
replace_flag = False
sample = pipeline(
prompt,
num_frames = video_length,
negative_prompt = negative_prompt,
height = sample_size[0],
width = sample_size[1],
generator = generator,
guidance_scale = guidance_scale,
num_inference_steps = num_inference_steps,
boundary = boundary,
pose_video = pose_video,
face_video = face_video,
ref_image = ref_image,
bg_video = bg_video,
mask_video = mask_video,
replace_flag = replace_flag,
shift = shift,
).videos
if lora_path is not None:
pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device, dtype=weight_dtype)
if transformer_2 is not None:
pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, dtype=weight_dtype, sub_transformer_name="transformer_2")
def save_results():
if not os.path.exists(save_path):
os.makedirs(save_path, exist_ok=True)
index = len([path for path in os.listdir(save_path)]) + 1
prefix = str(index).zfill(8)
if video_length == 1:
video_path = os.path.join(save_path, prefix + ".png")
image = sample[0, :, 0]
image = image.transpose(0, 1).transpose(1, 2)
image = (image * 255).numpy().astype(np.uint8)
image = Image.fromarray(image)
image.save(video_path)
else:
video_path = os.path.join(save_path, prefix + ".mp4")
save_videos_grid(sample, video_path, fps=fps)
if ulysses_degree * ring_degree > 1:
import torch.distributed as dist
if dist.get_rank() == 0:
save_results()
else:
save_results()
examples/wan2.2/predict_s2v.py
+7 -2
View File
@@ -110,9 +110,13 @@ fps = 16
# Use torch.float16 if GPU does not support torch.bfloat16
# ome graphics cards, such as v100, 2080ti, do not support torch.bfloat16
weight_dtype = torch.bfloat16
# If you want to generate from text, please set the validation_image_start = None and validation_image_end = None
# The path of the pose control video
control_video = "asset/pose.mp4"
# The path of the reference image
ref_image = "asset/8.png"
# Use ref_image as the first frame
init_first_frame = False
# The path of the audio
audio_path = "asset/talk.wav"
# prompts
@@ -335,7 +339,8 @@ with torch.no_grad():
pose_video = pose_video,
audio_path = audio_path,
shift = shift,
fps = fps
fps = fps,
init_first_frame = init_first_frame
).videos
if lora_path is not None:
scripts/wan2.2/README_TRAIN_Animate.md
Executable
+247
View File
@@ -0,0 +1,247 @@
## Training Code
The default training commands for the different versions are as follows:
We can choose whether to use fsdp in Wan-Animate, which can save a lot of video memory.
The metadata_control.json is a little different from normal json in Wan, you need to add some new paths in json.
- Animate tiem: a control_file_path, a face_file_path and a ref_file_path.
- Replace item: a control_file_path, a face_file_path, a ref_file_path, a mask_file_path and a background_file_path.
You can use
```json
[
{
"file_path": "train/00000001.mp4",
"control_file_path": "control/00000001_src_pose.mp4",
"face_file_path": "face/00000001_src_face.mp4",
"ref_file_path": "ref/00000001_src_ref.png",
"text": "A group of young men in suits and sunglasses are walking down a city street.",
"type": "video"
},
{
"file_path": "train/00000002.mp4",
"control_file_path": "control/00000001_src_pose.mp4",
"face_file_path": "face/00000001_src_face.mp4",
"ref_file_path": "ref/00000001_src_ref.png",
"background_file_path": "bg/00000002_src_bg.png",
"mask_file_path": "mask/00000002_src_mask.mp4",
"text": "视频中的人在做动作",
"type": "video",
"height": 480,
"width": 832
},
.....
]
```
Some parameters in the sh file can be confusing, and they are explained in this document:
- `enable_bucket` is used to enable bucket training. When enabled, the model does not crop the videos at the center, but instead, it trains the videos after grouping them into buckets based on resolution.
- `random_frame_crop` is used for random cropping on video frames to simulate videos with different frame counts.
- `random_hw_adapt` is used to enable automatic height and width scaling for videos. When `random_hw_adapt` is enabled, for training videos, the height and width will be set to `video_sample_size` as the maximum and `512` as the minimum.
- For example, when `random_hw_adapt` is enabled, with `video_sample_n_frames=49`, `video_sample_size=768`, the resolution of video inputs for training is `512x512x49`, `768x768x49`.
- `training_with_video_token_length` specifies training the model according to token length. For training videos, the height and width will be set to `video_sample_size` as the maximum and `256` as the minimum.
- For example, when `training_with_video_token_length` is enabled, with `video_sample_n_frames=49`, `token_sample_size=512`, `video_sample_size=768`, the resolution of video inputs for training is `256x256x49`, `512x512x49`, `768x768x21`.
- The token length for a video with dimensions 512x512 and 49 frames is 13,312. We need to set the `token_sample_size = 512`.
- At 512x512 resolution, the number of video frames is 49 (~= 512 * 512 * 49 / 512 / 512).
- At 768x768 resolution, the number of video frames is 21 (~= 512 * 512 * 49 / 768 / 768).
- At 1024x1024 resolution, the number of video frames is 9 (~= 512 * 512 * 49 / 1024 / 1024).
- These resolutions combined with their corresponding lengths allow the model to generate videos of different sizes.
- `resume_from_checkpoint` is used to set the training should be resumed from a previous checkpoint. Use a path or `"latest"` to automatically select the last available checkpoint.
Wan-Animate without deepspeed:
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-Animate-14B/"
export DATASET_NAME="datasets/internal_datasets/"
export DATASET_META_NAME="datasets/internal_datasets/metadata_control.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_animate.py \
--config_path="config/wan2.2/wan_civitai_animate.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=1024 \
--video_sample_size=256 \
--token_sample_size=512 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-05 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--boundary_type="full" \
--low_vram \
--trainable_modules "."
```
Wan-Animate with deepspeed zero-2:
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-Animate-14B/"
export DATASET_NAME="datasets/internal_datasets/"
export DATASET_META_NAME="datasets/internal_datasets/metadata_control.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --use_deepspeed --deepspeed_config_file config/zero_stage2_config.json --deepspeed_multinode_launcher standard scripts/wan2.2/train_animate.py \
--config_path="config/wan2.2/wan_civitai_animate.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=1024 \
--video_sample_size=256 \
--token_sample_size=512 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-05 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--boundary_type="full" \
--low_vram \
--trainable_modules "."
```
Wan-Animate with deepspeed zero-3:
```sh
python scripts/zero_to_bf16.py output_dir/checkpoint-{our-num-steps} output_dir/checkpoint-{your-num-steps}-outputs --max_shard_size 80GB --safe_serialization
```
Training shell command is as follows:
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-Animate-14B/"
export DATASET_NAME="datasets/internal_datasets/"
export DATASET_META_NAME="datasets/internal_datasets/metadata_control.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --zero_stage 3 --zero3_save_16bit_model true --zero3_init_flag true --use_deepspeed --deepspeed_config_file config/zero_stage3_config.json --deepspeed_multinode_launcher standard scripts/wan2.2/train_animate.py \
--config_path="config/wan2.2/wan_civitai_animate.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=1024 \
--video_sample_size=256 \
--token_sample_size=512 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-05 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--boundary_type="full" \
--low_vram \
--trainable_modules "."
```
Wan-Animate with FSDP:
Wan with FSDP is suitable for 14B Wan at high resolutions. Training shell command is as follows:
```sh
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-Animate-14B/"
export DATASET_NAME="datasets/internal_datasets/"
export DATASET_META_NAME="datasets/internal_datasets/metadata_control.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=WanAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/wan2.2/train_animate.py \
--config_path="config/wan2.2/wan_civitai_animate.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=1024 \
--video_sample_size=256 \
--token_sample_size=512 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-05 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--boundary_type="full" \
--low_vram \
--trainable_modules "."
```
scripts/wan2.2/README_TRAIN_S2V.md
+2 -2
View File
@@ -4,7 +4,7 @@ The default training commands for the different versions are as follows:
We can choose whether to use fsdp in Wan-S2V, which can save a lot of video memory.
The metadata_control.json is a little different from normal json in Wan-S2V, you need to add a audio_path.
The metadata_control.json is a little different from normal json in Wan, you need to add a audio_path.
```json
[
@@ -183,7 +183,7 @@ export DATASET_META_NAME="datasets/internal_datasets/metadata_control.json"
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=AudioAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/wan2.2/train_s2v.py \
accelerate launch --mixed_precision="bf16" --use_fsdp --fsdp_auto_wrap_policy TRANSFORMER_BASED_WRAP --fsdp_transformer_layer_cls_to_wrap=WanS2VAttentionBlock --fsdp_sharding_strategy "FULL_SHARD" --fsdp_state_dict_type=SHARDED_STATE_DICT --fsdp_backward_prefetch "BACKWARD_PRE" --fsdp_cpu_ram_efficient_loading False scripts/wan2.2/train_s2v.py \
--config_path="config/wan2.2/wan_civitai_s2v.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
scripts/wan2.2/train_animate.py
+2050
View File
File diff suppressed because it is too large Load Diff
scripts/wan2.2/train_animate.sh
+42
View File
@@ -0,0 +1,42 @@
export MODEL_NAME="models/Diffusion_Transformer/Wan2.2-Animate-14B/"
export DATASET_NAME="datasets/internal_datasets/"
export DATASET_META_NAME="datasets/internal_datasets/metadata_control.json"
# NCCL_IB_DISABLE=1 and NCCL_P2P_DISABLE=1 are used in multi nodes without RDMA.
# export NCCL_IB_DISABLE=1
# export NCCL_P2P_DISABLE=1
NCCL_DEBUG=INFO
accelerate launch --mixed_precision="bf16" scripts/wan2.2/train_animate.py \
--config_path="config/wan2.2/wan_civitai_animate.yaml" \
--pretrained_model_name_or_path=$MODEL_NAME \
--train_data_dir=$DATASET_NAME \
--train_data_meta=$DATASET_META_NAME \
--image_sample_size=1024 \
--video_sample_size=256 \
--token_sample_size=512 \
--video_sample_stride=2 \
--video_sample_n_frames=81 \
--train_batch_size=1 \
--video_repeat=1 \
--gradient_accumulation_steps=1 \
--dataloader_num_workers=8 \
--num_train_epochs=100 \
--checkpointing_steps=50 \
--learning_rate=2e-05 \
--lr_scheduler="constant_with_warmup" \
--lr_warmup_steps=100 \
--seed=42 \
--output_dir="output_dir" \
--gradient_checkpointing \
--mixed_precision="bf16" \
--adam_weight_decay=3e-2 \
--adam_epsilon=1e-10 \
--vae_mini_batch=1 \
--max_grad_norm=0.05 \
--random_hw_adapt \
--training_with_video_token_length \
--enable_bucket \
--uniform_sampling \
--boundary_type="full" \
--low_vram \
--trainable_modules "."
scripts/wan2.2/train_s2v.py
+1 -33
View File
@@ -90,38 +90,6 @@ def filter_kwargs(cls, kwargs):
filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
return filtered_kwargs
def get_random_downsample_ratio(sample_size, image_ratio=[],
all_choices=False, rng=None):
def _create_special_list(length):
if length == 1:
return [1.0]
if length >= 2:
first_element = 0.75
remaining_sum = 1.0 - first_element
other_elements_value = remaining_sum / (length - 1)
special_list = [first_element] + [other_elements_value] * (length - 1)
return special_list
if sample_size >= 1536:
number_list = [1, 1.25, 1.5, 2, 2.5, 3] + image_ratio
elif sample_size >= 1024:
number_list = [1, 1.25, 1.5, 2] + image_ratio
elif sample_size >= 768:
number_list = [1, 1.25, 1.5] + image_ratio
elif sample_size >= 512:
number_list = [1] + image_ratio
else:
number_list = [1]
if all_choices:
return number_list
number_list_prob = np.array(_create_special_list(len(number_list)))
if rng is None:
return np.random.choice(number_list, p = number_list_prob)
else:
return rng.choice(number_list, p = number_list_prob)
def resize_mask(mask, latent, process_first_frame_only=True):
latent_size = latent.size()
batch_size, channels, num_frames, height, width = mask.shape
@@ -1799,7 +1767,7 @@ def main():
init_first_frame = rng.choice([0, 1], p = [0.50, 0.50])
if init_first_frame or has_motion_pixel_values:
if not has_motion_pixel_values:
motion_pixel_values[:, -6:, :] = ref_pixel_values[:, 0, :]
motion_pixel_values[:, -6:, :] = ref_pixel_values
motion_frames_latents_length = int((args.motion_frames - 1) / sample_n_frames_bucket_interval + 1)
local_pixel_values = torch.cat([motion_pixel_values, pixel_values], dim = 1)
videox_fun/data/__init__.py
+1 -1
View File
@@ -1,7 +1,7 @@
from .dataset_image import CC15M, ImageEditDataset
from .dataset_image_video import (ImageVideoControlDataset, ImageVideoDataset,
ImageVideoSampler)
from .dataset_video import VideoDataset, VideoSpeechDataset, WebVid10M
from .dataset_video import VideoDataset, VideoSpeechDataset, VideoAnimateDataset, WebVid10M
from .utils import (VIDEO_READER_TIMEOUT, Camera, VideoReader_contextmanager,
custom_meshgrid, get_random_mask, get_relative_pose,
get_video_reader_batch, padding_image, process_pose_file,
videox_fun/data/dataset_image_video.py
+1
View File
@@ -595,6 +595,7 @@ class ImageVideoControlDataset(Dataset):
return sample
class ImageVideoSafetensorsDataset(Dataset):
def __init__(
self,
videox_fun/data/dataset_video.py
+328 -3
View File
@@ -554,18 +554,343 @@ class VideoSpeechControlDataset(Dataset):
return sample
class VideoAnimateDataset(Dataset):
def __init__(
self,
ann_path, data_root=None,
video_sample_size=512,
video_sample_stride=4,
video_sample_n_frames=16,
video_repeat=0,
text_drop_ratio=0.1,
enable_bucket=False,
video_length_drop_start=0.1,
video_length_drop_end=0.9,
return_file_name=False,
):
# Loading annotations from files
print(f"loading annotations from {ann_path} ...")
if ann_path.endswith('.csv'):
with open(ann_path, 'r') as csvfile:
dataset = list(csv.DictReader(csvfile))
elif ann_path.endswith('.json'):
dataset = json.load(open(ann_path))
self.data_root = data_root
# It's used to balance num of images and videos.
if video_repeat > 0:
self.dataset = []
for data in dataset:
if data.get('type', 'image') != 'video':
self.dataset.append(data)
for _ in range(video_repeat):
for data in dataset:
if data.get('type', 'image') == 'video':
self.dataset.append(data)
else:
self.dataset = dataset
del dataset
self.length = len(self.dataset)
print(f"data scale: {self.length}")
# TODO: enable bucket training
self.enable_bucket = enable_bucket
self.text_drop_ratio = text_drop_ratio
self.video_length_drop_start = video_length_drop_start
self.video_length_drop_end = video_length_drop_end
# Video params
self.video_sample_stride = video_sample_stride
self.video_sample_n_frames = video_sample_n_frames
self.video_sample_size = tuple(video_sample_size) if not isinstance(video_sample_size, int) else (video_sample_size, video_sample_size)
self.video_transforms = transforms.Compose(
[
transforms.Resize(min(self.video_sample_size)),
transforms.CenterCrop(self.video_sample_size),
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
]
)
self.larger_side_of_image_and_video = min(self.video_sample_size)
def get_batch(self, idx):
data_info = self.dataset[idx % len(self.dataset)]
video_id, text = data_info['file_path'], data_info['text']
if self.data_root is None:
video_dir = video_id
else:
video_dir = os.path.join(self.data_root, video_id)
with VideoReader_contextmanager(video_dir, num_threads=2) as video_reader:
min_sample_n_frames = min(
self.video_sample_n_frames,
int(len(video_reader) * (self.video_length_drop_end - self.video_length_drop_start) // self.video_sample_stride)
)
if min_sample_n_frames == 0:
raise ValueError(f"No Frames in video.")
video_length = int(self.video_length_drop_end * len(video_reader))
clip_length = min(video_length, (min_sample_n_frames - 1) * self.video_sample_stride + 1)
start_idx = random.randint(int(self.video_length_drop_start * video_length), video_length - clip_length) if video_length != clip_length else 0
batch_index = np.linspace(start_idx, start_idx + clip_length - 1, min_sample_n_frames, dtype=int)
try:
sample_args = (video_reader, batch_index)
pixel_values = func_timeout(
VIDEO_READER_TIMEOUT, get_video_reader_batch, args=sample_args
)
resized_frames = []
for i in range(len(pixel_values)):
frame = pixel_values[i]
resized_frame = resize_frame(frame, self.larger_side_of_image_and_video)
resized_frames.append(resized_frame)
pixel_values = np.array(resized_frames)
except FunctionTimedOut:
raise ValueError(f"Read {idx} timeout.")
except Exception as e:
raise ValueError(f"Failed to extract frames from video. Error is {e}.")
if not self.enable_bucket:
pixel_values = torch.from_numpy(pixel_values).permute(0, 3, 1, 2).contiguous()
pixel_values = pixel_values / 255.
del video_reader
else:
pixel_values = pixel_values
if not self.enable_bucket:
pixel_values = self.video_transforms(pixel_values)
# Random use no text generation
if random.random() < self.text_drop_ratio:
text = ''
control_video_id = data_info['control_file_path']
if control_video_id is not None:
if self.data_root is None:
control_video_id = control_video_id
else:
control_video_id = os.path.join(self.data_root, control_video_id)
if control_video_id is not None:
with VideoReader_contextmanager(control_video_id, num_threads=2) as control_video_reader:
try:
sample_args = (control_video_reader, batch_index)
control_pixel_values = func_timeout(
VIDEO_READER_TIMEOUT, get_video_reader_batch, args=sample_args
)
resized_frames = []
for i in range(len(control_pixel_values)):
frame = control_pixel_values[i]
resized_frame = resize_frame(frame, self.larger_side_of_image_and_video)
resized_frames.append(resized_frame)
control_pixel_values = np.array(resized_frames)
except FunctionTimedOut:
raise ValueError(f"Read {idx} timeout.")
except Exception as e:
raise ValueError(f"Failed to extract frames from video. Error is {e}.")
if not self.enable_bucket:
control_pixel_values = torch.from_numpy(control_pixel_values).permute(0, 3, 1, 2).contiguous()
control_pixel_values = control_pixel_values / 255.
del control_video_reader
else:
control_pixel_values = control_pixel_values
if not self.enable_bucket:
control_pixel_values = self.video_transforms(control_pixel_values)
else:
if not self.enable_bucket:
control_pixel_values = torch.zeros_like(pixel_values)
else:
control_pixel_values = np.zeros_like(pixel_values)
face_video_id = data_info['face_file_path']
if face_video_id is not None:
if self.data_root is None:
face_video_id = face_video_id
else:
face_video_id = os.path.join(self.data_root, face_video_id)
if face_video_id is not None:
with VideoReader_contextmanager(face_video_id, num_threads=2) as face_video_reader:
try:
sample_args = (face_video_reader, batch_index)
face_pixel_values = func_timeout(
VIDEO_READER_TIMEOUT, get_video_reader_batch, args=sample_args
)
resized_frames = []
for i in range(len(face_pixel_values)):
frame = face_pixel_values[i]
resized_frame = resize_frame(frame, self.larger_side_of_image_and_video)
resized_frames.append(resized_frame)
face_pixel_values = np.array(resized_frames)
except FunctionTimedOut:
raise ValueError(f"Read {idx} timeout.")
except Exception as e:
raise ValueError(f"Failed to extract frames from video. Error is {e}.")
if not self.enable_bucket:
face_pixel_values = torch.from_numpy(face_pixel_values).permute(0, 3, 1, 2).contiguous()
face_pixel_values = face_pixel_values / 255.
del face_video_reader
else:
face_pixel_values = face_pixel_values
if not self.enable_bucket:
face_pixel_values = self.video_transforms(face_pixel_values)
else:
if not self.enable_bucket:
face_pixel_values = torch.zeros_like(pixel_values)
else:
face_pixel_values = np.zeros_like(pixel_values)
background_video_id = data_info.get('background_file_path', None)
if background_video_id is not None:
if self.data_root is None:
background_video_id = background_video_id
else:
background_video_id = os.path.join(self.data_root, background_video_id)
if background_video_id is not None:
with VideoReader_contextmanager(background_video_id, num_threads=2) as background_video_reader:
try:
sample_args = (background_video_reader, batch_index)
background_pixel_values = func_timeout(
VIDEO_READER_TIMEOUT, get_video_reader_batch, args=sample_args
)
resized_frames = []
for i in range(len(background_pixel_values)):
frame = background_pixel_values[i]
resized_frame = resize_frame(frame, self.larger_side_of_image_and_video)
resized_frames.append(resized_frame)
background_pixel_values = np.array(resized_frames)
except FunctionTimedOut:
raise ValueError(f"Read {idx} timeout.")
except Exception as e:
raise ValueError(f"Failed to extract frames from video. Error is {e}.")
if not self.enable_bucket:
background_pixel_values = torch.from_numpy(background_pixel_values).permute(0, 3, 1, 2).contiguous()
background_pixel_values = background_pixel_values / 255.
del background_video_reader
else:
background_pixel_values = background_pixel_values
if not self.enable_bucket:
background_pixel_values = self.video_transforms(background_pixel_values)
else:
if not self.enable_bucket:
background_pixel_values = torch.ones_like(pixel_values) * 127.5
else:
background_pixel_values = np.ones_like(pixel_values) * 127.5
mask_video_id = data_info.get('mask_file_path', None)
if mask_video_id is not None:
if self.data_root is None:
mask_video_id = mask_video_id
else:
mask_video_id = os.path.join(self.data_root, mask_video_id)
if mask_video_id is not None:
with VideoReader_contextmanager(mask_video_id, num_threads=2) as mask_video_reader:
try:
sample_args = (mask_video_reader, batch_index)
mask = func_timeout(
VIDEO_READER_TIMEOUT, get_video_reader_batch, args=sample_args
)
resized_frames = []
for i in range(len(mask)):
frame = mask[i]
resized_frame = resize_frame(frame, self.larger_side_of_image_and_video)
resized_frames.append(resized_frame)
mask = np.array(resized_frames)
except FunctionTimedOut:
raise ValueError(f"Read {idx} timeout.")
except Exception as e:
raise ValueError(f"Failed to extract frames from video. Error is {e}.")
if not self.enable_bucket:
mask = torch.from_numpy(mask).permute(0, 3, 1, 2).contiguous()
mask = mask / 255.
del mask_video_reader
else:
mask = mask
else:
if not self.enable_bucket:
mask = torch.ones_like(pixel_values)
else:
mask = np.ones_like(pixel_values) * 255
mask = mask[:, :, :, :1]
ref_pixel_values_path = data_info.get('ref_file_path', [])
if self.data_root is not None:
ref_pixel_values_path = os.path.join(self.data_root, ref_pixel_values_path)
ref_pixel_values = Image.open(ref_pixel_values_path).convert('RGB')
if not self.enable_bucket:
raise ValueError("Not enable_bucket is not supported now. ")
else:
ref_pixel_values = np.array(ref_pixel_values)
return pixel_values, control_pixel_values, face_pixel_values, background_pixel_values, mask, ref_pixel_values, text, "video"
def __len__(self):
return self.length
def __getitem__(self, idx):
data_info = self.dataset[idx % len(self.dataset)]
data_type = data_info.get('type', 'image')
while True:
sample = {}
try:
data_info_local = self.dataset[idx % len(self.dataset)]
data_type_local = data_info_local.get('type', 'image')
if data_type_local != data_type:
raise ValueError("data_type_local != data_type")
pixel_values, control_pixel_values, face_pixel_values, background_pixel_values, mask, ref_pixel_values, name, data_type = \
self.get_batch(idx)
sample["pixel_values"] = pixel_values
sample["control_pixel_values"] = control_pixel_values
sample["face_pixel_values"] = face_pixel_values
sample["background_pixel_values"] = background_pixel_values
sample["mask"] = mask
sample["ref_pixel_values"] = ref_pixel_values
sample["clip_pixel_values"] = ref_pixel_values
sample["text"] = name
sample["data_type"] = data_type
sample["idx"] = idx
if len(sample) > 0:
break
except Exception as e:
print(e, self.dataset[idx % len(self.dataset)])
idx = random.randint(0, self.length-1)
return sample
if __name__ == "__main__":
if 1:
dataset = VideoDataset(
json_path="/home/zhoumo.xjq/disk3/datasets/webvidval/results_2M_val.json",
json_path="./webvidval/results_2M_val.json",
sample_size=256,
sample_stride=4, sample_n_frames=16,
)
if 0:
dataset = WebVid10M(
csv_path="/mnt/petrelfs/guoyuwei/projects/datasets/webvid/results_2M_val.csv",
video_folder="/mnt/petrelfs/guoyuwei/projects/datasets/webvid/2M_val",
csv_path="./webvid/results_2M_val.csv",
video_folder="./webvid/2M_val",
sample_size=256,
sample_stride=4, sample_n_frames=16,
is_image=False,
videox_fun/models/__init__.py
+6 -3
View File
@@ -6,7 +6,9 @@ from transformers import (AutoTokenizer, CLIPImageProcessor, CLIPTextModel,
T5EncoderModel, T5Tokenizer, T5TokenizerFast)
try:
from transformers import Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer, Qwen2VLProcessor, Qwen2_5_VLConfig
from transformers import (Qwen2_5_VLConfig,
Qwen2_5_VLForConditionalGeneration,
Qwen2Tokenizer, Qwen2VLProcessor)
except:
Qwen2_5_VLForConditionalGeneration, Qwen2Tokenizer = None, None
Qwen2VLProcessor, Qwen2_5_VLConfig = None, None
@@ -14,9 +16,9 @@ except:
from .cogvideox_transformer3d import CogVideoXTransformer3DModel
from .cogvideox_vae import AutoencoderKLCogVideoX
from .flux_transformer2d import FluxTransformer2DModel
from .fantasytalking_transformer3d import FantasyTalkingTransformer3DModel
from .fantasytalking_audio_encoder import FantasyTalkingAudioEncoder
from .fantasytalking_transformer3d import FantasyTalkingTransformer3DModel
from .flux_transformer2d import FluxTransformer2DModel
from .qwenimage_transformer2d import QwenImageTransformer2DModel
from .qwenimage_vae import AutoencoderKLQwenImage
from .wan_audio_encoder import WanAudioEncoder
@@ -24,6 +26,7 @@ from .wan_image_encoder import CLIPModel
from .wan_text_encoder import WanT5EncoderModel
from .wan_transformer3d import (Wan2_2Transformer3DModel, WanRMSNorm,
WanSelfAttention, WanTransformer3DModel)
from .wan_transformer3d_animate import Wan2_2Transformer3DModel_Animate
from .wan_transformer3d_s2v import Wan2_2Transformer3DModel_S2V
from .wan_transformer3d_vace import VaceWanTransformer3DModel
from .wan_vae import AutoencoderKLWan, AutoencoderKLWan_
videox_fun/models/cache_utils.py
+2 -1
View File
@@ -11,7 +11,8 @@ def get_teacache_coefficients(model_name):
return [2.57151496e+05, -3.54229917e+04, 1.40286849e+03, -1.35890334e+01, 1.32517977e-01]
elif "wan2.1-i2v-14b-720p" in model_name.lower() or "wan2.1-fun-14b" in model_name.lower() or "wan2.2-fun" in model_name.lower() \
or "wan2.2-i2v-a14b" in model_name.lower() or "wan2.2-t2v-a14b" in model_name.lower() or "wan2.2-ti2v-5b" in model_name.lower() \
or "wan2.2-s2v" in model_name.lower() or "wan2.1-vace-14b" in model_name.lower() or "wan2.2-vace-fun" in model_name.lower():
or "wan2.2-s2v" in model_name.lower() or "wan2.1-vace-14b" in model_name.lower() or "wan2.2-vace-fun" in model_name.lower() \
or "wan2.2-animate" in model_name.lower():
return [8.10705460e+03, 2.13393892e+03, -3.72934672e+02, 1.66203073e+01, -4.17769401e-02]
elif "qwen-image" in model_name.lower():
# Copied from https://github.com/chenpipi0807/ComfyUI-TeaCache/blob/main/nodes.py
videox_fun/models/cogvideox_transformer3d.py
+73 -3
View File
@@ -23,8 +23,7 @@ import torch.nn.functional as F
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.attention import Attention, FeedForward
from diffusers.models.attention_processor import (
AttentionProcessor, CogVideoXAttnProcessor2_0,
FusedCogVideoXAttnProcessor2_0)
AttentionProcessor, FusedCogVideoXAttnProcessor2_0)
from diffusers.models.embeddings import (CogVideoXPatchEmbed,
TimestepEmbedding, Timesteps,
get_3d_sincos_pos_embed)
@@ -39,8 +38,79 @@ from ..dist import (get_sequence_parallel_rank,
get_sequence_parallel_world_size, get_sp_group,
xFuserLongContextAttention)
from ..dist.cogvideox_xfuser import CogVideoXMultiGPUsAttnProcessor2_0
from .attention_utils import attention
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
class CogVideoXAttnProcessor2_0:
r"""
Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on
query and key vectors, but does not include spatial normalization.
"""
def __init__(self):
if not hasattr(F, "scaled_dot_product_attention"):
raise ImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
def __call__(
self,
attn,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor,
attention_mask: torch.Tensor = None,
image_rotary_emb: torch.Tensor = None,
) -> torch.Tensor:
text_seq_length = encoder_hidden_states.size(1)
hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
batch_size, sequence_length, _ = hidden_states.shape
if attention_mask is not None:
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
value = attn.to_v(hidden_states)
inner_dim = key.shape[-1]
head_dim = inner_dim // attn.heads
query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
if attn.norm_q is not None:
query = attn.norm_q(query)
if attn.norm_k is not None:
key = attn.norm_k(key)
# Apply RoPE if needed
if image_rotary_emb is not None:
from diffusers.models.embeddings import apply_rotary_emb
query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
if not attn.is_cross_attention:
key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
query = query.transpose(1, 2)
key = key.transpose(1, 2)
value = value.transpose(1, 2)
hidden_states = attention(
query, key, value, attn_mask=attention_mask, dropout_p=0.0, causal=False
)
hidden_states = hidden_states.reshape(batch_size, -1, attn.heads * head_dim)
# linear proj
hidden_states = attn.to_out[0](hidden_states)
# dropout
hidden_states = attn.to_out[1](hidden_states)
encoder_hidden_states, hidden_states = hidden_states.split(
[text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
)
return hidden_states, encoder_hidden_states
class CogVideoXPatchEmbed(nn.Module):
videox_fun/models/wan_animate_adapter.py
+383
View File
@@ -0,0 +1,383 @@
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import math
from typing import Optional, Tuple
import torch
import torch.nn.functional as F
from einops import rearrange
from torch import nn
try:
from flash_attn import flash_attn_func, flash_attn_qkvpacked_func
except ImportError:
flash_attn_func = None
MEMORY_LAYOUT = {
"flash": (
lambda x: x.view(x.shape[0] * x.shape[1], *x.shape[2:]),
lambda x: x,
),
"torch": (
lambda x: x.transpose(1, 2),
lambda x: x.transpose(1, 2),
),
"vanilla": (
lambda x: x.transpose(1, 2),
lambda x: x.transpose(1, 2),
),
}
def attention(
q,
k,
v,
mode="flash",
drop_rate=0,
attn_mask=None,
causal=False,
max_seqlen_q=None,
batch_size=1,
):
"""
Perform QKV self attention.
Args:
q (torch.Tensor): Query tensor with shape [b, s, a, d], where a is the number of heads.
k (torch.Tensor): Key tensor with shape [b, s1, a, d]
v (torch.Tensor): Value tensor with shape [b, s1, a, d]
mode (str): Attention mode. Choose from 'self_flash', 'cross_flash', 'torch', and 'vanilla'.
drop_rate (float): Dropout rate in attention map. (default: 0)
attn_mask (torch.Tensor): Attention mask with shape [b, s1] (cross_attn), or [b, a, s, s1] (torch or vanilla).
(default: None)
causal (bool): Whether to use causal attention. (default: False)
cu_seqlens_q (torch.Tensor): dtype torch.int32. The cumulative sequence lengths of the sequences in the batch,
used to index into q.
cu_seqlens_kv (torch.Tensor): dtype torch.int32. The cumulative sequence lengths of the sequences in the batch,
used to index into kv.
max_seqlen_q (int): The maximum sequence length in the batch of q.
max_seqlen_kv (int): The maximum sequence length in the batch of k and v.
Returns:
torch.Tensor: Output tensor after self attention with shape [b, s, ad]
"""
pre_attn_layout, post_attn_layout = MEMORY_LAYOUT[mode]
if mode == "torch":
if attn_mask is not None and attn_mask.dtype != torch.bool:
attn_mask = attn_mask.to(q.dtype)
x = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=drop_rate, is_causal=causal)
elif mode == "flash":
x = flash_attn_func(
q,
k,
v,
)
x = x.view(batch_size, max_seqlen_q, x.shape[-2], x.shape[-1]) # reshape x to [b, s, a, d]
elif mode == "vanilla":
scale_factor = 1 / math.sqrt(q.size(-1))
b, a, s, _ = q.shape
s1 = k.size(2)
attn_bias = torch.zeros(b, a, s, s1, dtype=q.dtype, device=q.device)
if causal:
# Only applied to self attention
assert attn_mask is None, "Causal mask and attn_mask cannot be used together"
temp_mask = torch.ones(b, a, s, s, dtype=torch.bool, device=q.device).tril(diagonal=0)
attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
attn_bias.to(q.dtype)
if attn_mask is not None:
if attn_mask.dtype == torch.bool:
attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
else:
attn_bias += attn_mask
attn = (q @ k.transpose(-2, -1)) * scale_factor
attn += attn_bias
attn = attn.softmax(dim=-1)
attn = torch.dropout(attn, p=drop_rate, train=True)
x = attn @ v
else:
raise NotImplementedError(f"Unsupported attention mode: {mode}")
x = post_attn_layout(x)
b, s, a, d = x.shape
out = x.reshape(b, s, -1)
return out
class CausalConv1d(nn.Module):
def __init__(self, chan_in, chan_out, kernel_size=3, stride=1, dilation=1, pad_mode="replicate", **kwargs):
super().__init__()
self.pad_mode = pad_mode
padding = (kernel_size - 1, 0) # T
self.time_causal_padding = padding
self.conv = nn.Conv1d(chan_in, chan_out, kernel_size, stride=stride, dilation=dilation, **kwargs)
def forward(self, x):
x = F.pad(x, self.time_causal_padding, mode=self.pad_mode)
return self.conv(x)
class FaceEncoder(nn.Module):
def __init__(self, in_dim: int, hidden_dim: int, num_heads=int, dtype=None, device=None):
factory_kwargs = {"dtype": dtype, "device": device}
super().__init__()
self.num_heads = num_heads
self.conv1_local = CausalConv1d(in_dim, 1024 * num_heads, 3, stride=1)
self.norm1 = nn.LayerNorm(hidden_dim // 8, elementwise_affine=False, eps=1e-6, **factory_kwargs)
self.act = nn.SiLU()
self.conv2 = CausalConv1d(1024, 1024, 3, stride=2)
self.conv3 = CausalConv1d(1024, 1024, 3, stride=2)
self.out_proj = nn.Linear(1024, hidden_dim)
self.norm1 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
self.norm2 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
self.norm3 = nn.LayerNorm(1024, elementwise_affine=False, eps=1e-6, **factory_kwargs)
self.padding_tokens = nn.Parameter(torch.zeros(1, 1, 1, hidden_dim))
def forward(self, x):
x = rearrange(x, "b t c -> b c t")
b, c, t = x.shape
x = self.conv1_local(x)
x = rearrange(x, "b (n c) t -> (b n) t c", n=self.num_heads)
x = self.norm1(x)
x = self.act(x)
x = rearrange(x, "b t c -> b c t")
x = self.conv2(x)
x = rearrange(x, "b c t -> b t c")
x = self.norm2(x)
x = self.act(x)
x = rearrange(x, "b t c -> b c t")
x = self.conv3(x)
x = rearrange(x, "b c t -> b t c")
x = self.norm3(x)
x = self.act(x)
x = self.out_proj(x)
x = rearrange(x, "(b n) t c -> b t n c", b=b)
padding = self.padding_tokens.repeat(b, x.shape[1], 1, 1)
x = torch.cat([x, padding], dim=-2)
x_local = x.clone()
return x_local
class RMSNorm(nn.Module):
def __init__(
self,
dim: int,
elementwise_affine=True,
eps: float = 1e-6,
device=None,
dtype=None,
):
"""
Initialize the RMSNorm normalization layer.
Args:
dim (int): The dimension of the input tensor.
eps (float, optional): A small value added to the denominator for numerical stability. Default is 1e-6.
Attributes:
eps (float): A small value added to the denominator for numerical stability.
weight (nn.Parameter): Learnable scaling parameter.
"""
factory_kwargs = {"device": device, "dtype": dtype}
super().__init__()
self.eps = eps
if elementwise_affine:
self.weight = nn.Parameter(torch.ones(dim, **factory_kwargs))
def _norm(self, x):
"""
Apply the RMSNorm normalization to the input tensor.
Args:
x (torch.Tensor): The input tensor.
Returns:
torch.Tensor: The normalized tensor.
"""
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
def forward(self, x):
"""
Forward pass through the RMSNorm layer.
Args:
x (torch.Tensor): The input tensor.
Returns:
torch.Tensor: The output tensor after applying RMSNorm.
"""
output = self._norm(x.float()).type_as(x)
if hasattr(self, "weight"):
output = output * self.weight
return output
def get_norm_layer(norm_layer):
"""
Get the normalization layer.
Args:
norm_layer (str): The type of normalization layer.
Returns:
norm_layer (nn.Module): The normalization layer.
"""
if norm_layer == "layer":
return nn.LayerNorm
elif norm_layer == "rms":
return RMSNorm
else:
raise NotImplementedError(f"Norm layer {norm_layer} is not implemented")
class FaceAdapter(nn.Module):
def __init__(
self,
hidden_dim: int,
heads_num: int,
qk_norm: bool = True,
qk_norm_type: str = "rms",
num_adapter_layers: int = 1,
dtype=None,
device=None,
):
factory_kwargs = {"dtype": dtype, "device": device}
super().__init__()
self.hidden_size = hidden_dim
self.heads_num = heads_num
self.fuser_blocks = nn.ModuleList(
[
FaceBlock(
self.hidden_size,
self.heads_num,
qk_norm=qk_norm,
qk_norm_type=qk_norm_type,
**factory_kwargs,
)
for _ in range(num_adapter_layers)
]
)
def forward(
self,
x: torch.Tensor,
motion_embed: torch.Tensor,
idx: int,
freqs_cis_q: Tuple[torch.Tensor, torch.Tensor] = None,
freqs_cis_k: Tuple[torch.Tensor, torch.Tensor] = None,
) -> torch.Tensor:
return self.fuser_blocks[idx](x, motion_embed, freqs_cis_q, freqs_cis_k)
class FaceBlock(nn.Module):
def __init__(
self,
hidden_size: int,
heads_num: int,
qk_norm: bool = True,
qk_norm_type: str = "rms",
qk_scale: float = None,
dtype: Optional[torch.dtype] = None,
device: Optional[torch.device] = None,
):
factory_kwargs = {"device": device, "dtype": dtype}
super().__init__()
self.deterministic = False
self.hidden_size = hidden_size
self.heads_num = heads_num
head_dim = hidden_size // heads_num
self.scale = qk_scale or head_dim**-0.5
self.linear1_kv = nn.Linear(hidden_size, hidden_size * 2, **factory_kwargs)
self.linear1_q = nn.Linear(hidden_size, hidden_size, **factory_kwargs)
self.linear2 = nn.Linear(hidden_size, hidden_size, **factory_kwargs)
qk_norm_layer = get_norm_layer(qk_norm_type)
self.q_norm = (
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity()
)
self.k_norm = (
qk_norm_layer(head_dim, elementwise_affine=True, eps=1e-6, **factory_kwargs) if qk_norm else nn.Identity()
)
self.pre_norm_feat = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
self.pre_norm_motion = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6, **factory_kwargs)
def forward(
self,
x: torch.Tensor,
motion_vec: torch.Tensor,
motion_mask: Optional[torch.Tensor] = None,
use_context_parallel=False,
) -> torch.Tensor:
B, T, N, C = motion_vec.shape
T_comp = T
x_motion = self.pre_norm_motion(motion_vec)
x_feat = self.pre_norm_feat(x)
kv = self.linear1_kv(x_motion)
q = self.linear1_q(x_feat)
k, v = rearrange(kv, "B L N (K H D) -> K B L N H D", K=2, H=self.heads_num)
q = rearrange(q, "B S (H D) -> B S H D", H=self.heads_num)
# Apply QK-Norm if needed.
q = self.q_norm(q).to(v)
k = self.k_norm(k).to(v)
k = rearrange(k, "B L N H D -> (B L) N H D")
v = rearrange(v, "B L N H D -> (B L) N H D")
# if use_context_parallel:
# q = gather_forward(q, dim=1)
q = rearrange(q, "B (L S) H D -> (B L) S H D", L=T_comp)
# Compute attention.
attn = attention(
q,
k,
v,
max_seqlen_q=q.shape[1],
batch_size=q.shape[0],
)
attn = rearrange(attn, "(B L) S C -> B (L S) C", L=T_comp)
# if use_context_parallel:
# attn = torch.chunk(attn, get_world_size(), dim=1)[get_rank()]
output = self.linear2(attn)
if motion_mask is not None:
output = output * rearrange(motion_mask, "B T H W -> B (T H W)").unsqueeze(-1)
return output
videox_fun/models/wan_animate_motion_encoder.py
+309
View File
@@ -0,0 +1,309 @@
# Modified from ``https://github.com/wyhsirius/LIA``
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import math
import torch
import torch.nn as nn
from torch.nn import functional as F
def custom_qr(input_tensor):
original_dtype = input_tensor.dtype
if original_dtype == torch.bfloat16:
q, r = torch.linalg.qr(input_tensor.to(torch.float32))
return q.to(original_dtype), r.to(original_dtype)
return torch.linalg.qr(input_tensor)
def fused_leaky_relu(input, bias, negative_slope=0.2, scale=2 ** 0.5):
return F.leaky_relu(input + bias, negative_slope) * scale
def upfirdn2d_native(input, kernel, up_x, up_y, down_x, down_y, pad_x0, pad_x1, pad_y0, pad_y1):
_, minor, in_h, in_w = input.shape
kernel_h, kernel_w = kernel.shape
out = input.view(-1, minor, in_h, 1, in_w, 1)
out = F.pad(out, [0, up_x - 1, 0, 0, 0, up_y - 1, 0, 0])
out = out.view(-1, minor, in_h * up_y, in_w * up_x)
out = F.pad(out, [max(pad_x0, 0), max(pad_x1, 0), max(pad_y0, 0), max(pad_y1, 0)])
out = out[:, :, max(-pad_y0, 0): out.shape[2] - max(-pad_y1, 0),
max(-pad_x0, 0): out.shape[3] - max(-pad_x1, 0), ]
out = out.reshape([-1, 1, in_h * up_y + pad_y0 + pad_y1, in_w * up_x + pad_x0 + pad_x1])
w = torch.flip(kernel, [0, 1]).view(1, 1, kernel_h, kernel_w)
out = F.conv2d(out, w)
out = out.reshape(-1, minor, in_h * up_y + pad_y0 + pad_y1 - kernel_h + 1,
in_w * up_x + pad_x0 + pad_x1 - kernel_w + 1, )
return out[:, :, ::down_y, ::down_x]
def upfirdn2d(input, kernel, up=1, down=1, pad=(0, 0)):
return upfirdn2d_native(input, kernel, up, up, down, down, pad[0], pad[1], pad[0], pad[1])
def make_kernel(k):
k = torch.tensor(k, dtype=torch.float32)
if k.ndim == 1:
k = k[None, :] * k[:, None]
k /= k.sum()
return k
class FusedLeakyReLU(nn.Module):
def __init__(self, channel, negative_slope=0.2, scale=2 ** 0.5):
super().__init__()
self.bias = nn.Parameter(torch.zeros(1, channel, 1, 1))
self.negative_slope = negative_slope
self.scale = scale
def forward(self, input):
out = fused_leaky_relu(input, self.bias, self.negative_slope, self.scale)
return out
class Blur(nn.Module):
def __init__(self, kernel, pad, upsample_factor=1):
super().__init__()
kernel = make_kernel(kernel)
if upsample_factor > 1:
kernel = kernel * (upsample_factor ** 2)
self.register_buffer('kernel', kernel)
self.pad = pad
def forward(self, input):
return upfirdn2d(input, self.kernel, pad=self.pad)
class ScaledLeakyReLU(nn.Module):
def __init__(self, negative_slope=0.2):
super().__init__()
self.negative_slope = negative_slope
def forward(self, input):
return F.leaky_relu(input, negative_slope=self.negative_slope)
class EqualConv2d(nn.Module):
def __init__(self, in_channel, out_channel, kernel_size, stride=1, padding=0, bias=True):
super().__init__()
self.weight = nn.Parameter(torch.randn(out_channel, in_channel, kernel_size, kernel_size))
self.scale = 1 / math.sqrt(in_channel * kernel_size ** 2)
self.stride = stride
self.padding = padding
if bias:
self.bias = nn.Parameter(torch.zeros(out_channel))
else:
self.bias = None
def forward(self, input):
return F.conv2d(input, self.weight * self.scale, bias=self.bias, stride=self.stride, padding=self.padding)
def __repr__(self):
return (
f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]},'
f' {self.weight.shape[2]}, stride={self.stride}, padding={self.padding})'
)
class EqualLinear(nn.Module):
def __init__(self, in_dim, out_dim, bias=True, bias_init=0, lr_mul=1, activation=None):
super().__init__()
self.weight = nn.Parameter(torch.randn(out_dim, in_dim).div_(lr_mul))
if bias:
self.bias = nn.Parameter(torch.zeros(out_dim).fill_(bias_init))
else:
self.bias = None
self.activation = activation
self.scale = (1 / math.sqrt(in_dim)) * lr_mul
self.lr_mul = lr_mul
def forward(self, input):
if self.activation:
out = F.linear(input, self.weight * self.scale)
out = fused_leaky_relu(out, self.bias * self.lr_mul)
else:
out = F.linear(input, self.weight * self.scale, bias=self.bias * self.lr_mul)
return out
def __repr__(self):
return (f'{self.__class__.__name__}({self.weight.shape[1]}, {self.weight.shape[0]})')
class ConvLayer(nn.Sequential):
def __init__(
self,
in_channel,
out_channel,
kernel_size,
downsample=False,
blur_kernel=[1, 3, 3, 1],
bias=True,
activate=True,
):
layers = []
if downsample:
factor = 2
p = (len(blur_kernel) - factor) + (kernel_size - 1)
pad0 = (p + 1) // 2
pad1 = p // 2
layers.append(Blur(blur_kernel, pad=(pad0, pad1)))
stride = 2
self.padding = 0
else:
stride = 1
self.padding = kernel_size // 2
layers.append(EqualConv2d(in_channel, out_channel, kernel_size, padding=self.padding, stride=stride,
bias=bias and not activate))
if activate:
if bias:
layers.append(FusedLeakyReLU(out_channel))
else:
layers.append(ScaledLeakyReLU(0.2))
super().__init__(*layers)
class ResBlock(nn.Module):
def __init__(self, in_channel, out_channel, blur_kernel=[1, 3, 3, 1]):
super().__init__()
self.conv1 = ConvLayer(in_channel, in_channel, 3)
self.conv2 = ConvLayer(in_channel, out_channel, 3, downsample=True)
self.skip = ConvLayer(in_channel, out_channel, 1, downsample=True, activate=False, bias=False)
def forward(self, input):
out = self.conv1(input)
out = self.conv2(out)
skip = self.skip(input)
out = (out + skip) / math.sqrt(2)
return out
class EncoderApp(nn.Module):
def __init__(self, size, w_dim=512):
super(EncoderApp, self).__init__()
channels = {
4: 512,
8: 512,
16: 512,
32: 512,
64: 256,
128: 128,
256: 64,
512: 32,
1024: 16
}
self.w_dim = w_dim
log_size = int(math.log(size, 2))
self.convs = nn.ModuleList()
self.convs.append(ConvLayer(3, channels[size], 1))
in_channel = channels[size]
for i in range(log_size, 2, -1):
out_channel = channels[2 ** (i - 1)]
self.convs.append(ResBlock(in_channel, out_channel))
in_channel = out_channel
self.convs.append(EqualConv2d(in_channel, self.w_dim, 4, padding=0, bias=False))
def forward(self, x):
res = []
h = x
for conv in self.convs:
h = conv(h)
res.append(h)
return res[-1].squeeze(-1).squeeze(-1), res[::-1][2:]
class Encoder(nn.Module):
def __init__(self, size, dim=512, dim_motion=20):
super(Encoder, self).__init__()
# appearance netmork
self.net_app = EncoderApp(size, dim)
# motion network
fc = [EqualLinear(dim, dim)]
for i in range(3):
fc.append(EqualLinear(dim, dim))
fc.append(EqualLinear(dim, dim_motion))
self.fc = nn.Sequential(*fc)
def enc_app(self, x):
h_source = self.net_app(x)
return h_source
def enc_motion(self, x):
h, _ = self.net_app(x)
h_motion = self.fc(h)
return h_motion
class Direction(nn.Module):
def __init__(self, motion_dim):
super(Direction, self).__init__()
self.weight = nn.Parameter(torch.randn(512, motion_dim))
def forward(self, input):
weight = self.weight + 1e-8
Q, R = custom_qr(weight)
if input is None:
return Q
else:
input_diag = torch.diag_embed(input) # alpha, diagonal matrix
out = torch.matmul(input_diag, Q.T)
out = torch.sum(out, dim=1)
return out
class Synthesis(nn.Module):
def __init__(self, motion_dim):
super(Synthesis, self).__init__()
self.direction = Direction(motion_dim)
class Generator(nn.Module):
def __init__(self, size, style_dim=512, motion_dim=20):
super().__init__()
self.enc = Encoder(size, style_dim, motion_dim)
self.dec = Synthesis(motion_dim)
def get_motion(self, img):
#motion_feat = self.enc.enc_motion(img)
motion_feat = torch.utils.checkpoint.checkpoint((self.enc.enc_motion), img, use_reentrant=True)
with torch.cuda.amp.autocast(dtype=torch.float32):
motion = self.dec.direction(motion_feat)
return motion
videox_fun/models/wan_transformer3d_animate.py
+299
View File
@@ -0,0 +1,299 @@
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
import math
import types
from copy import deepcopy
from typing import List
import numpy as np
import torch
import torch.cuda.amp as amp
import torch.nn as nn
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.loaders import PeftAdapterMixin
from diffusers.models.modeling_utils import ModelMixin
from diffusers.utils import is_torch_version, logging
from einops import rearrange
from .attention_utils import attention
from .wan_animate_adapter import FaceAdapter, FaceEncoder
from .wan_animate_motion_encoder import Generator
from .wan_transformer3d import (Head, MLPProj, WanAttentionBlock, WanLayerNorm,
WanRMSNorm, WanSelfAttention,
WanTransformer3DModel, rope_apply,
sinusoidal_embedding_1d)
from ..utils import cfg_skip
class Wan2_2Transformer3DModel_Animate(WanTransformer3DModel):
# _no_split_modules = ['WanAnimateAttentionBlock']
_supports_gradient_checkpointing = True
@register_to_config
def __init__(
self,
patch_size=(1, 2, 2),
text_len=512,
in_dim=36,
dim=5120,
ffn_dim=13824,
freq_dim=256,
text_dim=4096,
out_dim=16,
num_heads=40,
num_layers=40,
window_size=(-1, -1),
qk_norm=True,
cross_attn_norm=True,
eps=1e-6,
motion_encoder_dim=512,
use_context_parallel=False,
use_img_emb=True
):
model_type = "i2v" # TODO: Hard code for both preview and official versions.
super().__init__(model_type, patch_size, text_len, in_dim, dim, ffn_dim, freq_dim, text_dim, out_dim,
num_heads, num_layers, window_size, qk_norm, cross_attn_norm, eps)
self.motion_encoder_dim = motion_encoder_dim
self.use_context_parallel = use_context_parallel
self.use_img_emb = use_img_emb
self.pose_patch_embedding = nn.Conv3d(
16, dim, kernel_size=patch_size, stride=patch_size
)
# initialize weights
self.init_weights()
self.motion_encoder = Generator(size=512, style_dim=512, motion_dim=20)
self.face_adapter = FaceAdapter(
heads_num=self.num_heads,
hidden_dim=self.dim,
num_adapter_layers=self.num_layers // 5,
)
self.face_encoder = FaceEncoder(
in_dim=motion_encoder_dim,
hidden_dim=self.dim,
num_heads=4,
)
def after_patch_embedding(self, x: List[torch.Tensor], pose_latents, face_pixel_values):
pose_latents = [self.pose_patch_embedding(u.unsqueeze(0)) for u in pose_latents]
for x_, pose_latents_ in zip(x, pose_latents):
x_[:, :, 1:] += pose_latents_
b,c,T,h,w = face_pixel_values.shape
face_pixel_values = rearrange(face_pixel_values, "b c t h w -> (b t) c h w")
encode_bs = 8
face_pixel_values_tmp = []
for i in range(math.ceil(face_pixel_values.shape[0]/encode_bs)):
face_pixel_values_tmp.append(self.motion_encoder.get_motion(face_pixel_values[i*encode_bs:(i+1)*encode_bs]))
motion_vec = torch.cat(face_pixel_values_tmp)
motion_vec = rearrange(motion_vec, "(b t) c -> b t c", t=T)
motion_vec = self.face_encoder(motion_vec)
B, L, H, C = motion_vec.shape
pad_face = torch.zeros(B, 1, H, C).type_as(motion_vec)
motion_vec = torch.cat([pad_face, motion_vec], dim=1)
return x, motion_vec
def after_transformer_block(self, block_idx, x, motion_vec, motion_masks=None):
if block_idx % 5 == 0:
adapter_args = [x, motion_vec, motion_masks, self.use_context_parallel]
residual_out = self.face_adapter.fuser_blocks[block_idx // 5](*adapter_args)
x = residual_out + x
return x
@cfg_skip()
def forward(
self,
x,
t,
clip_fea,
context,
seq_len,
y=None,
pose_latents=None,
face_pixel_values=None,
cond_flag=True
):
# params
device = self.patch_embedding.weight.device
dtype = x.dtype
if self.freqs.device != device and torch.device(type="meta") != device:
self.freqs = self.freqs.to(device)
if y is not None:
x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
# embeddings
x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
x, motion_vec = self.after_patch_embedding(x, pose_latents, face_pixel_values)
grid_sizes = torch.stack(
[torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
x = [u.flatten(2).transpose(1, 2) for u in x]
seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
assert seq_lens.max() <= seq_len
x = torch.cat([
torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],
dim=1) for u in x
])
# time embeddings
with amp.autocast(dtype=torch.float32):
e = self.time_embedding(
sinusoidal_embedding_1d(self.freq_dim, t).float()
)
e0 = self.time_projection(e).unflatten(1, (6, self.dim))
assert e.dtype == torch.float32 and e0.dtype == torch.float32
# context
context_lens = None
context = self.text_embedding(
torch.stack([
torch.cat(
[u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
for u in context
]))
if self.use_img_emb:
context_clip = self.img_emb(clip_fea) # bs x 257 x dim
context = torch.concat([context_clip, context], dim=1)
# Context Parallel
if self.sp_world_size > 1:
x = torch.chunk(x, self.sp_world_size, dim=1)[self.sp_world_rank]
if t.dim() != 1:
e0 = torch.chunk(e0, self.sp_world_size, dim=1)[self.sp_world_rank]
e = torch.chunk(e, self.sp_world_size, dim=1)[self.sp_world_rank]
# TeaCache
if self.teacache is not None:
if cond_flag:
if t.dim() != 1:
modulated_inp = e0[0][:, -1, :]
else:
modulated_inp = e0[0]
skip_flag = self.teacache.cnt < self.teacache.num_skip_start_steps
if skip_flag:
self.should_calc = True
self.teacache.accumulated_rel_l1_distance = 0
else:
if cond_flag:
rel_l1_distance = self.teacache.compute_rel_l1_distance(self.teacache.previous_modulated_input, modulated_inp)
self.teacache.accumulated_rel_l1_distance += self.teacache.rescale_func(rel_l1_distance)
if self.teacache.accumulated_rel_l1_distance < self.teacache.rel_l1_thresh:
self.should_calc = False
else:
self.should_calc = True
self.teacache.accumulated_rel_l1_distance = 0
self.teacache.previous_modulated_input = modulated_inp
self.teacache.should_calc = self.should_calc
else:
self.should_calc = self.teacache.should_calc
# TeaCache
if self.teacache is not None:
if not self.should_calc:
previous_residual = self.teacache.previous_residual_cond if cond_flag else self.teacache.previous_residual_uncond
x = x + previous_residual.to(x.device)[-x.size()[0]:,]
else:
ori_x = x.clone().cpu() if self.teacache.offload else x.clone()
for idx, block in enumerate(self.blocks):
if torch.is_grad_enabled() and self.gradient_checkpointing:
def create_custom_forward(module):
def custom_forward(*inputs):
return module(*inputs)
return custom_forward
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
x = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
x,
e0,
seq_lens,
grid_sizes,
self.freqs,
context,
context_lens,
dtype,
t,
**ckpt_kwargs,
)
x = self.after_transformer_block(idx, x, motion_vec)
else:
# arguments
kwargs = dict(
e=e0,
seq_lens=seq_lens,
grid_sizes=grid_sizes,
freqs=self.freqs,
context=context,
context_lens=context_lens,
dtype=dtype,
t=t
)
x = block(x, **kwargs)
x = self.after_transformer_block(idx, x, motion_vec)
if cond_flag:
self.teacache.previous_residual_cond = x.cpu() - ori_x if self.teacache.offload else x - ori_x
else:
self.teacache.previous_residual_uncond = x.cpu() - ori_x if self.teacache.offload else x - ori_x
else:
for idx, block in enumerate(self.blocks):
if torch.is_grad_enabled() and self.gradient_checkpointing:
def create_custom_forward(module):
def custom_forward(*inputs):
return module(*inputs)
return custom_forward
ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
x = torch.utils.checkpoint.checkpoint(
create_custom_forward(block),
x,
e0,
seq_lens,
grid_sizes,
self.freqs,
context,
context_lens,
dtype,
t,
**ckpt_kwargs,
)
x = self.after_transformer_block(idx, x, motion_vec)
else:
# arguments
kwargs = dict(
e=e0,
seq_lens=seq_lens,
grid_sizes=grid_sizes,
freqs=self.freqs,
context=context,
context_lens=context_lens,
dtype=dtype,
t=t
)
x = block(x, **kwargs)
x = self.after_transformer_block(idx, x, motion_vec)
# head
x = self.head(x, e)
# Context Parallel
if self.sp_world_size > 1:
x = self.all_gather(x.contiguous(), dim=1)
# unpatchify
x = self.unpatchify(x, grid_sizes)
x = torch.stack(x)
return x
videox_fun/models/wan_transformer3d_s2v.py
+2
View File
@@ -596,6 +596,8 @@ class Wan2_2Transformer3DModel_S2V(Wan2_2Transformer3DModel):
"""
device = self.patch_embedding.weight.device
dtype = x.dtype
if self.freqs.device != device and torch.device(type="meta") != device:
self.freqs = self.freqs.to(device)
add_last_motion = self.add_last_motion * add_last_motion
# Embeddings
videox_fun/models/wan_transformer3d_vace.py
+1 -1
View File
@@ -216,8 +216,8 @@ class VaceWanTransformer3DModel(WanTransformer3DModel):
# if self.model_type == 'i2v':
# assert clip_fea is not None and y is not None
# params
dtype = x.dtype
device = self.patch_embedding.weight.device
dtype = x.dtype
if self.freqs.device != device and torch.device(type="meta") != device:
self.freqs = self.freqs.to(device)
videox_fun/pipeline/__init__.py
+7 -2
View File
@@ -1,13 +1,14 @@
from .pipeline_cogvideox_fun import CogVideoXFunPipeline
from .pipeline_cogvideox_fun_control import CogVideoXFunControlPipeline
from .pipeline_cogvideox_fun_inpaint import CogVideoXFunInpaintPipeline
from .pipeline_flux import FluxPipeline
from .pipeline_fantasy_talking import FantasyTalkingPipeline
from .pipeline_flux import FluxPipeline
from .pipeline_qwenimage import QwenImagePipeline
from .pipeline_qwenimage_edit import QwenImageEditPipeline
from .pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from .pipeline_wan import WanPipeline
from .pipeline_wan2_2 import Wan2_2Pipeline
from .pipeline_wan2_2_animate import Wan2_2AnimatePipeline
from .pipeline_wan2_2_fun_control import Wan2_2FunControlPipeline
from .pipeline_wan2_2_fun_inpaint import Wan2_2FunInpaintPipeline
from .pipeline_wan2_2_s2v import Wan2_2S2VPipeline
@@ -38,6 +39,7 @@ if importlib.util.find_spec("paifuser") is not None:
WanFunControlPipeline.__call__ = sparse_reset(WanFunControlPipeline.__call__)
WanI2VPipeline.__call__ = sparse_reset(WanI2VPipeline.__call__)
WanPipeline.__call__ = sparse_reset(WanPipeline.__call__)
WanVacePipeline.__call__ = sparse_reset(WanVacePipeline.__call__)
# Phantom
WanFunPhantomPipeline.__call__ = sparse_reset(WanFunPhantomPipeline.__call__)
@@ -48,4 +50,7 @@ if importlib.util.find_spec("paifuser") is not None:
Wan2_2FunControlPipeline.__call__ = sparse_reset(Wan2_2FunControlPipeline.__call__)
Wan2_2Pipeline.__call__ = sparse_reset(Wan2_2Pipeline.__call__)
Wan2_2I2VPipeline.__call__ = sparse_reset(Wan2_2I2VPipeline.__call__)
Wan2_2TI2VPipeline.__call__ = sparse_reset(Wan2_2TI2VPipeline.__call__)
Wan2_2TI2VPipeline.__call__ = sparse_reset(Wan2_2TI2VPipeline.__call__)
Wan2_2S2VPipeline.__call__ = sparse_reset(Wan2_2S2VPipeline.__call__)
Wan2_2VaceFunPipeline.__call__ = sparse_reset(Wan2_2VaceFunPipeline.__call__)
Wan2_2AnimatePipeline.__call__ = sparse_reset(Wan2_2AnimatePipeline.__call__)
videox_fun/pipeline/pipeline_wan2_2_animate.py
+929
View File
@@ -0,0 +1,929 @@
import inspect
import math
from copy import deepcopy
from dataclasses import dataclass
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
import numpy as np
import copy
import torch
import cv2
import torch.nn.functional as F
from einops import rearrange
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
from diffusers.image_processor import VaeImageProcessor
from diffusers.pipelines.pipeline_utils import DiffusionPipeline
from diffusers.utils import BaseOutput, logging, replace_example_docstring
from diffusers.utils.torch_utils import randn_tensor
from diffusers.video_processor import VideoProcessor
from decord import VideoReader
from ..models import (AutoencoderKLWan, AutoTokenizer, CLIPModel,
WanT5EncoderModel, Wan2_2Transformer3DModel_Animate)
from ..utils.fm_solvers import (FlowDPMSolverMultistepScheduler,
get_sampling_sigmas)
from ..utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
EXAMPLE_DOC_STRING = """
Examples:
```python
pass
```
"""
# 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,
):
"""
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
@dataclass
class WanPipelineOutput(BaseOutput):
r"""
Output class for CogVideo pipelines.
Args:
video (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
`(batch_size, num_frames, channels, height, width)`.
"""
videos: torch.Tensor
class Wan2_2AnimatePipeline(DiffusionPipeline):
r"""
Pipeline for text-to-video generation using Wan.
This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
"""
_optional_components = ["transformer_2", "clip_image_encoder"]
model_cpu_offload_seq = "text_encoder->clip_image_encoder->transformer_2->transformer->vae"
_callback_tensor_inputs = [
"latents",
"prompt_embeds",
"negative_prompt_embeds",
]
def __init__(
self,
tokenizer: AutoTokenizer,
text_encoder: WanT5EncoderModel,
vae: AutoencoderKLWan,
transformer: Wan2_2Transformer3DModel_Animate,
transformer_2: Wan2_2Transformer3DModel_Animate = None,
clip_image_encoder: CLIPModel = None,
scheduler: FlowMatchEulerDiscreteScheduler = None,
):
super().__init__()
self.register_modules(
tokenizer=tokenizer, text_encoder=text_encoder, vae=vae, transformer=transformer,
transformer_2=transformer_2, clip_image_encoder=clip_image_encoder, scheduler=scheduler
)
self.video_processor = VideoProcessor(vae_scale_factor=self.vae.spatial_compression_ratio)
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae.spatial_compression_ratio)
self.mask_processor = VaeImageProcessor(
vae_scale_factor=self.vae.spatial_compression_ratio, do_normalize=False, do_binarize=True, do_convert_grayscale=True
)
def _get_t5_prompt_embeds(
self,
prompt: Union[str, List[str]] = None,
num_videos_per_prompt: int = 1,
max_sequence_length: int = 512,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
):
device = device or self._execution_device
dtype = dtype or self.text_encoder.dtype
prompt = [prompt] if isinstance(prompt, str) else prompt
batch_size = len(prompt)
text_inputs = self.tokenizer(
prompt,
padding="max_length",
max_length=max_sequence_length,
truncation=True,
add_special_tokens=True,
return_tensors="pt",
)
text_input_ids = text_inputs.input_ids
prompt_attention_mask = text_inputs.attention_mask
untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])
logger.warning(
"The following part of your input was truncated because `max_sequence_length` is set to "
f" {max_sequence_length} tokens: {removed_text}"
)
seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long()
prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask.to(device))[0]
prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
# duplicate text embeddings for each generation per prompt, using mps friendly method
_, seq_len, _ = prompt_embeds.shape
prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
return [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
def encode_prompt(
self,
prompt: Union[str, List[str]],
negative_prompt: Optional[Union[str, List[str]]] = None,
do_classifier_free_guidance: bool = True,
num_videos_per_prompt: int = 1,
prompt_embeds: Optional[torch.Tensor] = None,
negative_prompt_embeds: Optional[torch.Tensor] = None,
max_sequence_length: int = 512,
device: Optional[torch.device] = None,
dtype: Optional[torch.dtype] = None,
):
r"""
Encodes the prompt into text encoder hidden states.
Args:
prompt (`str` or `List[str]`, *optional*):
prompt to be encoded
negative_prompt (`str` or `List[str]`, *optional*):
The prompt or prompts not to guide the image generation. If not defined, one has to pass
`negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
less than `1`).
do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
Whether to use classifier free guidance or not.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
provided, text embeddings will be generated from `prompt` input argument.
negative_prompt_embeds (`torch.Tensor`, *optional*):
Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
argument.
device: (`torch.device`, *optional*):
torch device
dtype: (`torch.dtype`, *optional*):
torch dtype
"""
device = device or self._execution_device
prompt = [prompt] if isinstance(prompt, str) else prompt
if prompt is not None:
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
if prompt_embeds is None:
prompt_embeds = self._get_t5_prompt_embeds(
prompt=prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
if do_classifier_free_guidance and negative_prompt_embeds is None:
negative_prompt = negative_prompt or ""
negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
f" {type(prompt)}."
)
elif batch_size != len(negative_prompt):
raise ValueError(
f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
" the batch size of `prompt`."
)
negative_prompt_embeds = self._get_t5_prompt_embeds(
prompt=negative_prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
device=device,
dtype=dtype,
)
return prompt_embeds, negative_prompt_embeds
def prepare_latents(
self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None
):
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
shape = (
batch_size,
num_channels_latents,
(num_frames - 1) // self.vae.temporal_compression_ratio + 1,
height // self.vae.spatial_compression_ratio,
width // self.vae.spatial_compression_ratio,
)
if latents is None:
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
else:
latents = latents.to(device)
# scale the initial noise by the standard deviation required by the scheduler
if hasattr(self.scheduler, "init_noise_sigma"):
latents = latents * self.scheduler.init_noise_sigma
return latents
def padding_resize(self, img_ori, height=512, width=512, padding_color=(0, 0, 0), interpolation=cv2.INTER_LINEAR):
ori_height = img_ori.shape[0]
ori_width = img_ori.shape[1]
channel = img_ori.shape[2]
img_pad = np.zeros((height, width, channel))
if channel == 1:
img_pad[:, :, 0] = padding_color[0]
else:
img_pad[:, :, 0] = padding_color[0]
img_pad[:, :, 1] = padding_color[1]
img_pad[:, :, 2] = padding_color[2]
if (ori_height / ori_width) > (height / width):
new_width = int(height / ori_height * ori_width)
img = cv2.resize(img_ori, (new_width, height), interpolation=interpolation)
padding = int((width - new_width) / 2)
if len(img.shape) == 2:
img = img[:, :, np.newaxis]
img_pad[:, padding: padding + new_width, :] = img
else:
new_height = int(width / ori_width * ori_height)
img = cv2.resize(img_ori, (width, new_height), interpolation=interpolation)
padding = int((height - new_height) / 2)
if len(img.shape) == 2:
img = img[:, :, np.newaxis]
img_pad[padding: padding + new_height, :, :] = img
img_pad = np.uint8(img_pad)
return img_pad
def inputs_padding(self, x, target_len):
ndim = x.ndim
if ndim == 4:
f = x.shape[0]
if target_len <= f:
return [deepcopy(x[i]) for i in range(target_len)]
idx = 0
flip = False
target_array = []
while len(target_array) < target_len:
target_array.append(deepcopy(x[idx]))
if flip:
idx -= 1
else:
idx += 1
if idx == 0 or idx == f - 1:
flip = not flip
return target_array[:target_len]
elif ndim == 5:
b, c, f, h, w = x.shape
if target_len <= f:
return x[:, :, :target_len, :, :]
indices = []
idx = 0
flip = False
while len(indices) < target_len:
indices.append(idx)
if flip:
idx -= 1
else:
idx += 1
if idx == 0 or idx == f - 1:
flip = not flip
indices = indices[:target_len]
if isinstance(x, torch.Tensor):
indices_tensor = torch.tensor(indices, device=x.device, dtype=torch.long)
return x[:, :, indices_tensor, :, :]
else:
indices_array = np.array(indices)
return x[:, :, indices_array, :, :]
else:
raise ValueError(f"Unsupported input dimension: {ndim}. Expected 4D or 5D.")
def get_valid_len(self, real_len, clip_len=81, overlap=1):
real_clip_len = clip_len - overlap
last_clip_num = (real_len - overlap) % real_clip_len
if last_clip_num == 0:
extra = 0
else:
extra = real_clip_len - last_clip_num
target_len = real_len + extra
return target_len
def prepare_source(self, src_pose_path, src_face_path, src_ref_path):
pose_video_reader = VideoReader(src_pose_path)
pose_len = len(pose_video_reader)
pose_idxs = list(range(pose_len))
pose_video = pose_video_reader.get_batch(pose_idxs).asnumpy()
face_video_reader = VideoReader(src_face_path)
face_len = len(face_video_reader)
face_idxs = list(range(face_len))
face_video = face_video_reader.get_batch(face_idxs).asnumpy()
height, width = pose_video[0].shape[:2]
ref_image = cv2.imread(src_ref_path)[..., ::-1]
ref_image = self.padding_resize(ref_image, height=height, width=width)
return pose_video, face_video, ref_image
def prepare_source_for_replace(self, src_bg_path, src_mask_path):
bg_video_reader = VideoReader(src_bg_path)
bg_len = len(bg_video_reader)
bg_idxs = list(range(bg_len))
bg_video = bg_video_reader.get_batch(bg_idxs).asnumpy()
mask_video_reader = VideoReader(src_mask_path)
mask_len = len(mask_video_reader)
mask_idxs = list(range(mask_len))
mask_video = mask_video_reader.get_batch(mask_idxs).asnumpy()
mask_video = mask_video[:, :, :, 0] / 255
return bg_video, mask_video
def get_i2v_mask(self, lat_t, lat_h, lat_w, mask_len=1, mask_pixel_values=None, device="cuda"):
if mask_pixel_values is None:
msk = torch.zeros(1, (lat_t-1) * 4 + 1, lat_h, lat_w, device=device)
else:
msk = mask_pixel_values.clone()
msk[:, :mask_len] = 1
msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1)
msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)
msk = msk.transpose(1, 2)
return msk
def prepare_control_latents(
self, control, control_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance
):
# resize the control to latents shape as we concatenate the control to the latents
# we do that before converting to dtype to avoid breaking in case we're using cpu_offload
# and half precision
if control is not None:
control = control.to(device=device, dtype=dtype)
bs = 1
new_control = []
for i in range(0, control.shape[0], bs):
control_bs = control[i : i + bs]
control_bs = self.vae.encode(control_bs)[0]
control_bs = control_bs.mode()
new_control.append(control_bs)
control = torch.cat(new_control, dim = 0)
if control_image is not None:
control_image = control_image.to(device=device, dtype=dtype)
bs = 1
new_control_pixel_values = []
for i in range(0, control_image.shape[0], bs):
control_pixel_values_bs = control_image[i : i + bs]
control_pixel_values_bs = self.vae.encode(control_pixel_values_bs)[0]
control_pixel_values_bs = control_pixel_values_bs.mode()
new_control_pixel_values.append(control_pixel_values_bs)
control_image_latents = torch.cat(new_control_pixel_values, dim = 0)
else:
control_image_latents = None
return control, control_image_latents
def decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
frames = self.vae.decode(latents.to(self.vae.dtype)).sample
frames = (frames / 2 + 0.5).clamp(0, 1)
# we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
# frames = frames.cpu().float().numpy()
return frames
# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
def prepare_extra_step_kwargs(self, generator, eta):
# prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
# eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
# eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
# and should be between [0, 1]
accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
extra_step_kwargs = {}
if accepts_eta:
extra_step_kwargs["eta"] = eta
# check if the scheduler accepts generator
accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
if accepts_generator:
extra_step_kwargs["generator"] = generator
return extra_step_kwargs
# Copied from diffusers.pipelines.latte.pipeline_latte.LattePipeline.check_inputs
def check_inputs(
self,
prompt,
height,
width,
negative_prompt,
callback_on_step_end_tensor_inputs,
prompt_embeds=None,
negative_prompt_embeds=None,
):
if height % 8 != 0 or width % 8 != 0:
raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
if callback_on_step_end_tensor_inputs is not None and not all(
k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
):
raise ValueError(
f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
)
if prompt is not None and prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
" only forward one of the two."
)
elif prompt is None and prompt_embeds is None:
raise ValueError(
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
)
elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
if prompt is not None and negative_prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
)
if negative_prompt is not None and negative_prompt_embeds is not None:
raise ValueError(
f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
)
if prompt_embeds is not None and negative_prompt_embeds is not None:
if prompt_embeds.shape != negative_prompt_embeds.shape:
raise ValueError(
"`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
f" {negative_prompt_embeds.shape}."
)
@property
def guidance_scale(self):
return self._guidance_scale
@property
def num_timesteps(self):
return self._num_timesteps
@property
def attention_kwargs(self):
return self._attention_kwargs
@property
def interrupt(self):
return self._interrupt
@torch.no_grad()
@replace_example_docstring(EXAMPLE_DOC_STRING)
def __call__(
self,
prompt: Optional[Union[str, List[str]]] = None,
negative_prompt: Optional[Union[str, List[str]]] = None,
height: int = 480,
width: int = 720,
clip_len=77,
num_frames: int = 49,
num_inference_steps: int = 50,
pose_video = None,
face_video = None,
ref_image = None,
bg_video = None,
mask_video = None,
replace_flag = True,
timesteps: Optional[List[int]] = None,
guidance_scale: float = 6,
num_videos_per_prompt: int = 1,
eta: float = 0.0,
generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
latents: Optional[torch.FloatTensor] = None,
prompt_embeds: Optional[torch.FloatTensor] = None,
negative_prompt_embeds: Optional[torch.FloatTensor] = None,
output_type: str = "numpy",
return_dict: bool = False,
callback_on_step_end: Optional[
Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
] = None,
attention_kwargs: Optional[Dict[str, Any]] = None,
callback_on_step_end_tensor_inputs: List[str] = ["latents"],
max_sequence_length: int = 512,
boundary: float = 0.875,
comfyui_progressbar: bool = False,
shift: int = 5,
refert_num = 1,
) -> Union[WanPipelineOutput, Tuple]:
"""
Function invoked when calling the pipeline for generation.
Args:
Examples:
Returns:
"""
if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
num_videos_per_prompt = 1
# 1. Check inputs. Raise error if not correct
self.check_inputs(
prompt,
height,
width,
negative_prompt,
callback_on_step_end_tensor_inputs,
prompt_embeds,
negative_prompt_embeds,
)
self._guidance_scale = guidance_scale
self._attention_kwargs = attention_kwargs
self._interrupt = False
# 2. Default call parameters
if prompt is not None and isinstance(prompt, str):
batch_size = 1
elif prompt is not None and isinstance(prompt, list):
batch_size = len(prompt)
else:
batch_size = prompt_embeds.shape[0]
device = self._execution_device
weight_dtype = self.text_encoder.dtype
# here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
# of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
# corresponds to doing no classifier free guidance.
do_classifier_free_guidance = guidance_scale > 1.0
# 3. Encode input prompt
prompt_embeds, negative_prompt_embeds = self.encode_prompt(
prompt,
negative_prompt,
do_classifier_free_guidance,
num_videos_per_prompt=num_videos_per_prompt,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
max_sequence_length=max_sequence_length,
device=device,
)
if do_classifier_free_guidance:
in_prompt_embeds = negative_prompt_embeds + prompt_embeds
else:
in_prompt_embeds = prompt_embeds
if comfyui_progressbar:
from comfy.utils import ProgressBar
pbar = ProgressBar(num_inference_steps + 1)
# 4. Prepare latents
if pose_video is not None:
video_length = pose_video.shape[2]
pose_video = self.image_processor.preprocess(rearrange(pose_video, "b c f h w -> (b f) c h w"), height=height, width=width)
pose_video = pose_video.to(dtype=torch.float32)
pose_video = rearrange(pose_video, "(b f) c h w -> b c f h w", f=video_length)
else:
pose_video = None
if face_video is not None:
video_length = face_video.shape[2]
face_video = self.image_processor.preprocess(rearrange(face_video, "b c f h w -> (b f) c h w"))
face_video = face_video.to(dtype=torch.float32)
face_video = rearrange(face_video, "(b f) c h w -> b c f h w", f=video_length)
else:
face_video = None
real_frame_len = pose_video.size()[2]
target_len = self.get_valid_len(real_frame_len, clip_len, overlap=refert_num)
print('real frames: {} target frames: {}'.format(real_frame_len, target_len))
pose_video = self.inputs_padding(pose_video, target_len).to(device, weight_dtype)
face_video = self.inputs_padding(face_video, target_len).to(device, weight_dtype)
ref_image = self.padding_resize(np.array(ref_image), height=height, width=width)
ref_image = torch.tensor(ref_image / 127.5 - 1).unsqueeze(0).permute([3, 0, 1, 2]).unsqueeze(0).to(device, weight_dtype)
if replace_flag:
if bg_video is not None:
video_length = bg_video.shape[2]
bg_video = self.image_processor.preprocess(rearrange(bg_video, "b c f h w -> (b f) c h w"), height=height, width=width)
bg_video = bg_video.to(dtype=torch.float32)
bg_video = rearrange(bg_video, "(b f) c h w -> b c f h w", f=video_length)
else:
bg_video = None
bg_video = self.inputs_padding(bg_video, target_len).to(device, weight_dtype)
mask_video = self.inputs_padding(mask_video, target_len).to(device, weight_dtype)
if comfyui_progressbar:
pbar.update(1)
# 5. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
target_shape = (self.vae.latent_channels, (num_frames - 1) // self.vae.temporal_compression_ratio + 1, width // self.vae.spatial_compression_ratio, height // self.vae.spatial_compression_ratio)
seq_len = math.ceil((target_shape[2] * target_shape[3]) / (self.transformer.config.patch_size[1] * self.transformer.config.patch_size[2]) * target_shape[1])
# 6. Denoising loop
start = 0
end = clip_len
all_out_frames = []
copy_timesteps = copy.deepcopy(timesteps)
copy_latents = copy.deepcopy(latents)
bs = pose_video.size()[0]
while True:
if start + refert_num >= pose_video.size()[2]:
break
# Prepare timesteps
if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, copy_timesteps, mu=1)
elif isinstance(self.scheduler, FlowUniPCMultistepScheduler):
self.scheduler.set_timesteps(num_inference_steps, device=device, shift=shift)
timesteps = self.scheduler.timesteps
elif isinstance(self.scheduler, FlowDPMSolverMultistepScheduler):
sampling_sigmas = get_sampling_sigmas(num_inference_steps, shift)
timesteps, _ = retrieve_timesteps(
self.scheduler,
device=device,
sigmas=sampling_sigmas)
else:
timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, copy_timesteps)
self._num_timesteps = len(timesteps)
latent_channels = self.transformer.config.in_channels
latents = self.prepare_latents(
batch_size * num_videos_per_prompt,
latent_channels,
num_frames,
height,
width,
weight_dtype,
device,
generator,
copy_latents,
)
if start == 0:
mask_reft_len = 0
else:
mask_reft_len = refert_num
conditioning_pixel_values = pose_video[:, :, start:end]
face_pixel_values = face_video[:, :, start:end]
ref_pixel_values = ref_image.clone().detach()
if start > 0:
refer_t_pixel_values = out_frames[:, :, -refert_num:].clone().detach()
refer_t_pixel_values = (refer_t_pixel_values - 0.5) / 0.5
else:
refer_t_pixel_values = torch.zeros(bs, 3, refert_num, height, width)
refer_t_pixel_values = refer_t_pixel_values.to(device=device, dtype=weight_dtype)
pose_latents, ref_latents = self.prepare_control_latents(
conditioning_pixel_values,
ref_pixel_values,
batch_size,
height,
width,
weight_dtype,
device,
generator,
do_classifier_free_guidance
)
mask_ref = self.get_i2v_mask(1, target_shape[-1], target_shape[-2], 1, device=device)
y_ref = torch.concat([mask_ref, ref_latents], dim=1).to(device=device, dtype=weight_dtype)
if mask_reft_len > 0:
if replace_flag:
# Image.fromarray(np.array((refer_t_pixel_values[0, :, 0].permute(1,2,0) * 0.5 + 0.5).float().cpu().numpy() *255, np.uint8)).save("1.jpg")
bg_pixel_values = bg_video[:, :, start:end]
y_reft = self.vae.encode(
torch.concat(
[
refer_t_pixel_values[:, :, :mask_reft_len],
bg_pixel_values[:, :, mask_reft_len:]
], dim=2
).to(device=device, dtype=weight_dtype)
)[0].mode()
mask_pixel_values = 1 - mask_video[:, :, start:end]
mask_pixel_values = rearrange(mask_pixel_values, "b c t h w -> (b t) c h w")
mask_pixel_values = F.interpolate(mask_pixel_values, size=(target_shape[-1], target_shape[-2]), mode='nearest')
mask_pixel_values = rearrange(mask_pixel_values, "(b t) c h w -> b c t h w", b = bs)[:, 0]
msk_reft = self.get_i2v_mask(
int((clip_len - 1) // self.vae.temporal_compression_ratio + 1), target_shape[-1], target_shape[-2], mask_reft_len, mask_pixel_values=mask_pixel_values, device=device
)
else:
refer_t_pixel_values = rearrange(refer_t_pixel_values[:, :, :mask_reft_len], "b c t h w -> (b t) c h w")
refer_t_pixel_values = F.interpolate(refer_t_pixel_values, size=(height, width), mode="bicubic")
refer_t_pixel_values = rearrange(refer_t_pixel_values, "(b t) c h w -> b c t h w", b = bs)
y_reft = self.vae.encode(
torch.concat(
[
refer_t_pixel_values,
torch.zeros(bs, 3, clip_len - mask_reft_len, height, width).to(device=device, dtype=weight_dtype),
], dim=2,
).to(device=device, dtype=weight_dtype)
)[0].mode()
msk_reft = self.get_i2v_mask(
int((clip_len - 1) // self.vae.temporal_compression_ratio + 1), target_shape[-1], target_shape[-2], mask_reft_len, device=device
)
else:
if replace_flag:
bg_pixel_values = bg_video[:, :, start:end]
y_reft = self.vae.encode(
bg_pixel_values.to(device=device, dtype=weight_dtype)
)[0].mode()
mask_pixel_values = 1 - mask_video[:, :, start:end]
mask_pixel_values = rearrange(mask_pixel_values, "b c t h w -> (b t) c h w")
mask_pixel_values = F.interpolate(mask_pixel_values, size=(target_shape[-1], target_shape[-2]), mode='nearest')
mask_pixel_values = rearrange(mask_pixel_values, "(b t) c h w -> b c t h w", b = bs)[:, 0]
msk_reft = self.get_i2v_mask(
int((clip_len - 1) // self.vae.temporal_compression_ratio + 1), target_shape[-1], target_shape[-2], mask_reft_len, mask_pixel_values=mask_pixel_values, device=device
)
else:
y_reft = self.vae.encode(
torch.zeros(1, 3, clip_len - mask_reft_len, height, width).to(device=device, dtype=weight_dtype)
)[0].mode()
msk_reft = self.get_i2v_mask(
int((clip_len - 1) // self.vae.temporal_compression_ratio + 1), target_shape[-1], target_shape[-2], mask_reft_len, device=device
)
y_reft = torch.concat([msk_reft, y_reft], dim=1).to(device=device, dtype=weight_dtype)
y = torch.concat([y_ref, y_reft], dim=2)
clip_context = self.clip_image_encoder([ref_pixel_values[0, :, :, :]]).to(device=device, dtype=weight_dtype)
num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
self.transformer.num_inference_steps = num_inference_steps
with self.progress_bar(total=num_inference_steps) as progress_bar:
for i, t in enumerate(timesteps):
self.transformer.current_steps = i
if self.interrupt:
continue
latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
if hasattr(self.scheduler, "scale_model_input"):
latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
y_in = torch.cat([y] * 2) if do_classifier_free_guidance else y
clip_context_input = (
torch.cat([clip_context] * 2) if do_classifier_free_guidance else clip_context
)
pose_latents_input = (
torch.cat([pose_latents] * 2) if do_classifier_free_guidance else pose_latents
)
face_pixel_values_input = (
torch.cat([torch.ones_like(face_pixel_values) * -1] + [face_pixel_values]) if do_classifier_free_guidance else face_pixel_values
)
# broadcast to batch dimension in a way that's compatible with ONNX/Core ML
timestep = t.expand(latent_model_input.shape[0])
if self.transformer_2 is not None:
if t >= boundary * self.scheduler.config.num_train_timesteps:
local_transformer = self.transformer_2
else:
local_transformer = self.transformer
else:
local_transformer = self.transformer
# predict noise model_output
with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=device):
noise_pred = local_transformer(
x=latent_model_input,
context=in_prompt_embeds,
t=timestep,
seq_len=seq_len,
y=y_in,
clip_fea=clip_context_input,
pose_latents=pose_latents_input,
face_pixel_values=face_pixel_values_input,
)
# Perform guidance
if do_classifier_free_guidance:
if self.transformer_2 is not None and (isinstance(self.guidance_scale, (list, tuple))):
sample_guide_scale = self.guidance_scale[1] if t >= self.transformer_2.config.boundary * self.scheduler.config.num_train_timesteps else self.guidance_scale[0]
else:
sample_guide_scale = self.guidance_scale
noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + sample_guide_scale * (noise_pred_text - noise_pred_uncond)
# Compute the previous noisy sample x_t -> x_t-1
latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
if callback_on_step_end is not None:
callback_kwargs = {}
for k in callback_on_step_end_tensor_inputs:
callback_kwargs[k] = locals()[k]
callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
latents = callback_outputs.pop("latents", latents)
prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
progress_bar.update()
if comfyui_progressbar:
pbar.update(1)
out_frames = self.decode_latents(latents[:, :, 1:])
if start != 0:
out_frames = out_frames[:, :, refert_num:]
all_out_frames.append(out_frames.cpu())
start += clip_len - refert_num
end += clip_len - refert_num
videos = torch.cat(all_out_frames, dim=2)[:, :, :real_frame_len]
# Offload all models
self.maybe_free_model_hooks()
return WanPipelineOutput(videos=videos.float().cpu())
videox_fun/pipeline/pipeline_wan2_2_s2v.py
+1 -1
View File
@@ -648,7 +648,7 @@ class Wan2_2S2VPipeline(DiffusionPipeline):
drop_first_motion = self.drop_first_motion
if init_first_frame:
drop_first_motion = False
motion_latents[:, :, -6:] = ref_image[:, :, 0]
motion_latents[:, :, -6:] = ref_image
motion_latents = self.vae.encode(motion_latents)[0].mode()
# Get pose cond input if need
videox_fun/utils/utils.py
+28 -28
View File
@@ -299,34 +299,6 @@ def get_video_to_video_latent(input_video_path, video_length, sample_size, fps=N
ref_image = ref_image.unsqueeze(0).permute([3, 0, 1, 2]).unsqueeze(0) / 255
return input_video, input_video_mask, ref_image, clip_image
def padding_image(images, new_width, new_height):
new_image = Image.new('RGB', (new_width, new_height), (255, 255, 255))
aspect_ratio = images.width / images.height
if new_width / new_height > 1:
if aspect_ratio > new_width / new_height:
new_img_width = new_width
new_img_height = int(new_img_width / aspect_ratio)
else:
new_img_height = new_height
new_img_width = int(new_img_height * aspect_ratio)
else:
if aspect_ratio > new_width / new_height:
new_img_width = new_width
new_img_height = int(new_img_width / aspect_ratio)
else:
new_img_height = new_height
new_img_width = int(new_img_height * aspect_ratio)
resized_img = images.resize((new_img_width, new_img_height))
paste_x = (new_width - new_img_width) // 2
paste_y = (new_height - new_img_height) // 2
new_image.paste(resized_img, (paste_x, paste_y))
return new_image
def get_image_latent(ref_image=None, sample_size=None, padding=False):
if ref_image is not None:
if isinstance(ref_image, str):
@@ -358,6 +330,34 @@ def get_image(ref_image=None):
return ref_image
def padding_image(images, new_width, new_height):
new_image = Image.new('RGB', (new_width, new_height), (255, 255, 255))
aspect_ratio = images.width / images.height
if new_width / new_height > 1:
if aspect_ratio > new_width / new_height:
new_img_width = new_width
new_img_height = int(new_img_width / aspect_ratio)
else:
new_img_height = new_height
new_img_width = int(new_img_height * aspect_ratio)
else:
if aspect_ratio > new_width / new_height:
new_img_width = new_width
new_img_height = int(new_img_width / aspect_ratio)
else:
new_img_height = new_height
new_img_width = int(new_img_height * aspect_ratio)
resized_img = images.resize((new_img_width, new_img_height))
paste_x = (new_width - new_img_width) // 2
paste_y = (new_height - new_img_height) // 2
new_image.paste(resized_img, (paste_x, paste_y))
return new_image
def timer(func):
def wrapper(*args, **kwargs):
start_time = time.time()