mirror of
https://github.com/Wan-Video/Wan2.2.git
synced 2026-08-28 17:43:23 +08:00
add codes of wan2.2-s2v
This commit is contained in:
qi wang
+6
-1
@@ -1,2 +1,7 @@
|
||||
__pycache__/
|
||||
.DS_Store
|
||||
.DS_Store
|
||||
.vscode*
|
||||
tmp_examples*
|
||||
new_checkpoint*
|
||||
batch_test*
|
||||
nohup*
|
||||
@@ -34,6 +34,7 @@ We are excited to introduce **Wan2.2**, a major upgrade to our foundational vide
|
||||
|
||||
## 🔥 Latest News!!
|
||||
|
||||
* Aug 26, 2025: 🎵 We introduce **[Wan2.2-S2V-14B](https://humanaigc.github.io/wan-s2v-webpage)**, an audio-driven cinematic video generation model, including [inference code](#run-speech-to-video-generation), [model weights](#model-download), and [technical report]()! Now you can try it on [wan.video](https://wan.video/), [ModelScope Gradio](https://www.modelscope.cn/studios/Wan-AI/Wan2.2-S2V) or [HuggingFace Gradio](https://huggingface.co/spaces/Wan-AI/Wan2.2-S2V)!
|
||||
* Jul 28, 2025: 👋 We have open a [HF space](https://huggingface.co/spaces/Wan-AI/Wan-2.2-5B) using the TI2V-5B model. Enjoy!
|
||||
* Jul 28, 2025: 👋 Wan2.2 has been integrated into ComfyUI ([CN](https://docs.comfy.org/zh-CN/tutorials/video/wan/wan2_2) | [EN](https://docs.comfy.org/tutorials/video/wan/wan2_2)). Enjoy!
|
||||
* Jul 28, 2025: 👋 Wan2.2's T2V, I2V and TI2V have been integrated into Diffusers ([T2V-A14B](https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B-Diffusers) | [I2V-A14B](https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B-Diffusers) | [TI2V-5B](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B-Diffusers)). Feel free to give it a try!
|
||||
@@ -63,6 +64,11 @@ If your research or project builds upon [**Wan2.1**](https://github.com/Wan-Vide
|
||||
- [x] Checkpoints of the 5B model
|
||||
- [x] ComfyUI integration
|
||||
- [x] Diffusers integration
|
||||
- Wan2.2-S2V Speech-to-Video
|
||||
- [x] Inference code of Wan2.2-S2V
|
||||
- [x] Checkpoints of Wan2.2-S2V-14B
|
||||
- [ ] ComfyUI integration
|
||||
- [ ] Diffusers integration
|
||||
|
||||
## Run Wan2.2
|
||||
|
||||
@@ -88,6 +94,8 @@ pip install -r requirements.txt
|
||||
| T2V-A14B | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B) 🤖 [ModelScope](https://modelscope.cn/models/Wan-AI/Wan2.2-T2V-A14B) | Text-to-Video MoE model, supports 480P & 720P |
|
||||
| I2V-A14B | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.2-I2V-A14B) 🤖 [ModelScope](https://modelscope.cn/models/Wan-AI/Wan2.2-I2V-A14B) | Image-to-Video MoE model, supports 480P & 720P |
|
||||
| TI2V-5B | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.2-TI2V-5B) 🤖 [ModelScope](https://modelscope.cn/models/Wan-AI/Wan2.2-TI2V-5B) | High-compression VAE, T2V+I2V, supports 720P |
|
||||
| S2V-14B | 🤗 [Huggingface](https://huggingface.co/Wan-AI/Wan2.2-S2V-14B) 🤖 [ModelScope](https://modelscope.cn/models/Wan-AI/Wan2.2-S2V-14B) | Speech-to-Video model, supports 480P & 720P |
|
||||
|
||||
|
||||
|
||||
> 💡Note:
|
||||
@@ -228,8 +236,38 @@ torchrun --nproc_per_node=8 generate.py --task ti2v-5B --size 1280*704 --ckpt_di
|
||||
|
||||
> The process of prompt extension can be referenced [here](#2-using-prompt-extention).
|
||||
|
||||
#### Run Speech-to-Video Generation
|
||||
|
||||
This repository supports the `Wan2.2-S2V-14B` Speech-to-Video model and can simultaneously support video generation at 480P and 720P resolutions.
|
||||
|
||||
- Single-GPU Speech-to-Video inference
|
||||
|
||||
```sh
|
||||
python generate.py --task s2v-14B --size 1024*704 --ckpt_dir ./Wan2.2-S2V-14B/ --offload_model True --convert_model_dtype --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard." --image "examples/i2v_input.JPG" --audio "examples/talk.wav"
|
||||
# Without setting --num_clip, the generated video length will automatically adjust based on the input audio length
|
||||
```
|
||||
|
||||
> 💡 This command can run on a GPU with at least 80GB VRAM.
|
||||
|
||||
- Multi-GPU inference using FSDP + DeepSpeed Ulysses
|
||||
|
||||
```sh
|
||||
torchrun --nproc_per_node=8 generate.py --task s2v-14B --size 1024*704 --ckpt_dir ./Wan2.2-S2V-14B/ --dit_fsdp --t5_fsdp --ulysses_size 8 --prompt "Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard." --image "examples/i2v_input.JPG" --audio "examples/talk.wav"
|
||||
```
|
||||
|
||||
- Pose + Audio driven generation
|
||||
|
||||
```sh
|
||||
torchrun --nproc_per_node=8 generate.py --task s2v-14B --size 1024*704 --ckpt_dir ./Wan2.2-S2V-14B/ --dit_fsdp --t5_fsdp --ulysses_size 8 --prompt "a person is singing" --image "examples/pose.png" --audio "examples/sing.MP3" --pose_video "./examples/pose.mp4"
|
||||
```
|
||||
|
||||
> 💡For the Speech-to-Video task, the `size` parameter represents the area of the generated video, with the aspect ratio following that of the original input image.
|
||||
|
||||
> 💡The model can generate videos from audio input combined with reference image and optional text prompt.
|
||||
|
||||
> 💡The `--pose_video` parameter enables pose-driven generation, allowing the model to follow specific pose sequences while generating videos synchronized with audio input.
|
||||
|
||||
> 💡The `--num_clip` parameter controls the number of video clips generated, useful for quick preview with shorter generation time.
|
||||
|
||||
## Computational Efficiency on Different GPUs
|
||||
|
||||
|
||||
Binary file not shown.
Binary file not shown.
|
After Width: | Height: | Size: 858 KiB |
Binary file not shown.
Binary file not shown.
|
After Width: | Height: | Size: 804 KiB |
Binary file not shown.
Binary file not shown.
+78
-4
@@ -18,7 +18,7 @@ import wan
|
||||
from wan.configs import MAX_AREA_CONFIGS, SIZE_CONFIGS, SUPPORTED_SIZES, WAN_CONFIGS
|
||||
from wan.distributed.util import init_distributed_group
|
||||
from wan.utils.prompt_extend import DashScopePromptExpander, QwenPromptExpander
|
||||
from wan.utils.utils import save_video, str2bool
|
||||
from wan.utils.utils import merge_video_audio, save_video, str2bool
|
||||
|
||||
EXAMPLE_PROMPT = {
|
||||
"t2v-A14B": {
|
||||
@@ -35,6 +35,14 @@ EXAMPLE_PROMPT = {
|
||||
"prompt":
|
||||
"Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage.",
|
||||
},
|
||||
"s2v-14B": {
|
||||
"prompt":
|
||||
"Summer beach vacation style, a white cat wearing sunglasses sits on a surfboard. The fluffy-furred feline gazes directly at the camera with a relaxed expression. Blurred beach scenery forms the background featuring crystal-clear waters, distant green hills, and a blue sky dotted with white clouds. The cat assumes a naturally relaxed posture, as if savoring the sea breeze and warm sunlight. A close-up shot highlights the feline's intricate details and the refreshing atmosphere of the seaside.",
|
||||
"image":
|
||||
"examples/i2v_input.JPG",
|
||||
"audio":
|
||||
"examples/talk.wav",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -48,6 +56,8 @@ def _validate_args(args):
|
||||
args.prompt = EXAMPLE_PROMPT[args.task]["prompt"]
|
||||
if args.image is None and "image" in EXAMPLE_PROMPT[args.task]:
|
||||
args.image = EXAMPLE_PROMPT[args.task]["image"]
|
||||
if args.audio is None and "audio" in EXAMPLE_PROMPT[args.task]:
|
||||
args.audio = EXAMPLE_PROMPT[args.task]["audio"]
|
||||
|
||||
if args.task == "i2v-A14B":
|
||||
assert args.image is not None, "Please specify the image path for i2v."
|
||||
@@ -69,9 +79,10 @@ def _validate_args(args):
|
||||
args.base_seed = args.base_seed if args.base_seed >= 0 else random.randint(
|
||||
0, sys.maxsize)
|
||||
# Size check
|
||||
assert args.size in SUPPORTED_SIZES[
|
||||
args.
|
||||
task], f"Unsupport size {args.size} for task {args.task}, supported sizes are: {', '.join(SUPPORTED_SIZES[args.task])}"
|
||||
if not 's2v' in args.task:
|
||||
assert args.size in SUPPORTED_SIZES[
|
||||
args.
|
||||
task], f"Unsupport size {args.size} for task {args.task}, supported sizes are: {', '.join(SUPPORTED_SIZES[args.task])}"
|
||||
|
||||
|
||||
def _parse_args():
|
||||
@@ -194,6 +205,36 @@ def _parse_args():
|
||||
default=False,
|
||||
help="Whether to convert model paramerters dtype.")
|
||||
|
||||
# following args only works for s2v
|
||||
parser.add_argument(
|
||||
"--num_clip",
|
||||
type=int,
|
||||
default=None,
|
||||
help="Number of video clips to generate, the whole video will not exceed the length of audio."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to the audio file, e.g. wav, mp3")
|
||||
parser.add_argument(
|
||||
"--pose_video",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Provide Dw-pose sequence to do Pose Driven")
|
||||
parser.add_argument(
|
||||
"--start_from_ref",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="whether set the reference image as the starting point for generation"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--infer_frames",
|
||||
type=int,
|
||||
default=80,
|
||||
help="Number of frames per clip, 48 or 80 or others (must be multiple of 4) for 14B s2v"
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
_validate_args(args)
|
||||
@@ -353,6 +394,37 @@ def generate(args):
|
||||
guide_scale=args.sample_guide_scale,
|
||||
seed=args.base_seed,
|
||||
offload_model=args.offload_model)
|
||||
elif "s2v" in args.task:
|
||||
logging.info("Creating WanS2V pipeline.")
|
||||
wan_s2v = wan.WanS2V(
|
||||
config=cfg,
|
||||
checkpoint_dir=args.ckpt_dir,
|
||||
device_id=device,
|
||||
rank=rank,
|
||||
t5_fsdp=args.t5_fsdp,
|
||||
dit_fsdp=args.dit_fsdp,
|
||||
use_sp=(args.ulysses_size > 1),
|
||||
t5_cpu=args.t5_cpu,
|
||||
convert_model_dtype=args.convert_model_dtype,
|
||||
)
|
||||
logging.info(f"Generating video ...")
|
||||
video = wan_s2v.generate(
|
||||
input_prompt=args.prompt,
|
||||
ref_image_path=args.image,
|
||||
audio_path=args.audio,
|
||||
num_repeat=args.num_clip,
|
||||
pose_video=args.pose_video,
|
||||
max_area=MAX_AREA_CONFIGS[args.size],
|
||||
infer_frames=args.infer_frames,
|
||||
shift=args.sample_shift,
|
||||
sample_solver=args.sample_solver,
|
||||
sampling_steps=args.sample_steps,
|
||||
guide_scale=args.sample_guide_scale,
|
||||
seed=args.base_seed,
|
||||
offload_model=args.offload_model,
|
||||
init_first_frame=args.start_from_ref,
|
||||
)
|
||||
|
||||
else:
|
||||
logging.info("Creating WanI2V pipeline.")
|
||||
wan_i2v = wan.WanI2V(
|
||||
@@ -396,6 +468,8 @@ def generate(args):
|
||||
nrow=1,
|
||||
normalize=True,
|
||||
value_range=(-1, 1))
|
||||
if "s2v" in args.task:
|
||||
merge_video_audio(video_path=args.save_file, audio_path=args.audio)
|
||||
del video
|
||||
|
||||
torch.cuda.synchronize()
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||||
from . import configs, distributed, modules
|
||||
from .image2video import WanI2V
|
||||
from .speech2video import WanS2V
|
||||
from .text2video import WanT2V
|
||||
from .textimage2video import WanTI2V
|
||||
|
||||
@@ -5,6 +5,7 @@ import os
|
||||
os.environ['TOKENIZERS_PARALLELISM'] = 'false'
|
||||
|
||||
from .wan_i2v_A14B import i2v_A14B
|
||||
from .wan_s2v_14B import s2v_14B
|
||||
from .wan_t2v_A14B import t2v_A14B
|
||||
from .wan_ti2v_5B import ti2v_5B
|
||||
|
||||
@@ -12,6 +13,7 @@ WAN_CONFIGS = {
|
||||
't2v-A14B': t2v_A14B,
|
||||
'i2v-A14B': i2v_A14B,
|
||||
'ti2v-5B': ti2v_5B,
|
||||
's2v-14B': s2v_14B,
|
||||
}
|
||||
|
||||
SIZE_CONFIGS = {
|
||||
@@ -20,7 +22,9 @@ SIZE_CONFIGS = {
|
||||
'480*832': (480, 832),
|
||||
'832*480': (832, 480),
|
||||
'704*1280': (704, 1280),
|
||||
'1280*704': (1280, 704)
|
||||
'1280*704': (1280, 704),
|
||||
'1024*704': (1024, 704),
|
||||
'704*1024': (704, 1024),
|
||||
}
|
||||
|
||||
MAX_AREA_CONFIGS = {
|
||||
@@ -30,10 +34,14 @@ MAX_AREA_CONFIGS = {
|
||||
'832*480': 832 * 480,
|
||||
'704*1280': 704 * 1280,
|
||||
'1280*704': 1280 * 704,
|
||||
'1024*704': 1024 * 704,
|
||||
'704*1024': 704 * 1024,
|
||||
}
|
||||
|
||||
SUPPORTED_SIZES = {
|
||||
't2v-A14B': ('720*1280', '1280*720', '480*832', '832*480'),
|
||||
'i2v-A14B': ('720*1280', '1280*720', '480*832', '832*480'),
|
||||
'ti2v-5B': ('704*1280', '1280*704'),
|
||||
's2v-14B': ('720*1280', '1280*720', '480*832', '832*480', '1024*704',
|
||||
'704*1024', '704*1280', '1280*704'),
|
||||
}
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||||
from easydict import EasyDict
|
||||
|
||||
from .shared_config import wan_shared_cfg
|
||||
|
||||
#------------------------ Wan S2V 14B ------------------------#
|
||||
|
||||
s2v_14B = EasyDict(__name__='Config: Wan S2V 14B')
|
||||
s2v_14B.update(wan_shared_cfg)
|
||||
|
||||
# t5
|
||||
s2v_14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
|
||||
s2v_14B.t5_tokenizer = 'google/umt5-xxl'
|
||||
|
||||
# vae
|
||||
s2v_14B.vae_checkpoint = 'Wan2.1_VAE.pth'
|
||||
s2v_14B.vae_stride = (4, 8, 8)
|
||||
|
||||
# wav2vec
|
||||
s2v_14B.wav2vec = "wav2vec2-large-xlsr-53-english"
|
||||
|
||||
s2v_14B.num_heads = 40
|
||||
# transformer
|
||||
s2v_14B.transformer = EasyDict(
|
||||
__name__="Config: Transformer config for WanModel_S2V")
|
||||
s2v_14B.transformer.patch_size = (1, 2, 2)
|
||||
s2v_14B.transformer.dim = 5120
|
||||
s2v_14B.transformer.ffn_dim = 13824
|
||||
s2v_14B.transformer.freq_dim = 256
|
||||
s2v_14B.transformer.num_heads = 40
|
||||
s2v_14B.transformer.num_layers = 40
|
||||
s2v_14B.transformer.window_size = (-1, -1)
|
||||
s2v_14B.transformer.qk_norm = True
|
||||
s2v_14B.transformer.cross_attn_norm = True
|
||||
s2v_14B.transformer.eps = 1e-6
|
||||
s2v_14B.transformer.enable_adain = True
|
||||
s2v_14B.transformer.adain_mode = "attn_norm"
|
||||
s2v_14B.transformer.audio_inject_layers = [
|
||||
0, 4, 8, 12, 16, 20, 24, 27, 30, 33, 36, 39
|
||||
]
|
||||
s2v_14B.transformer.zero_init = True
|
||||
s2v_14B.transformer.zero_timestep = True
|
||||
s2v_14B.transformer.enable_motioner = False
|
||||
s2v_14B.transformer.add_last_motion = True
|
||||
s2v_14B.transformer.trainable_token = False
|
||||
s2v_14B.transformer.enable_tsm = False
|
||||
s2v_14B.transformer.enable_framepack = True
|
||||
s2v_14B.transformer.framepack_drop_mode = 'padd'
|
||||
s2v_14B.transformer.audio_dim = 1024
|
||||
|
||||
s2v_14B.transformer.motion_frames = 73
|
||||
s2v_14B.transformer.cond_dim = 16
|
||||
|
||||
# inference
|
||||
s2v_14B.sample_neg_prompt = "画面模糊，最差质量，画面模糊，细节模糊不清，情绪激动剧烈，手快速抖动，字幕，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走"
|
||||
s2v_14B.drop_first_motion = True
|
||||
s2v_14B.sample_shift = 3
|
||||
s2v_14B.sample_steps = 40
|
||||
s2v_14B.sample_guide_scale = 4.5
|
||||
@@ -356,7 +356,7 @@ class WanModel(ModelMixin, ConfigMixin):
|
||||
|
||||
super().__init__()
|
||||
|
||||
assert model_type in ['t2v', 'i2v', 'ti2v']
|
||||
assert model_type in ['t2v', 'i2v', 'ti2v', 's2v']
|
||||
self.model_type = model_type
|
||||
|
||||
self.patch_size = patch_size
|
||||
|
||||
@@ -0,0 +1,5 @@
|
||||
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||||
from .audio_encoder import AudioEncoder
|
||||
from .model_s2v import WanModel_S2V
|
||||
|
||||
__all__ = ['WanModel_S2V', 'AudioEncoder']
|
||||
@@ -0,0 +1,189 @@
|
||||
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||||
import math
|
||||
|
||||
import librosa
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor
|
||||
|
||||
|
||||
def get_sample_indices(original_fps,
|
||||
total_frames,
|
||||
target_fps,
|
||||
num_sample,
|
||||
fixed_start=None):
|
||||
required_duration = num_sample / target_fps
|
||||
required_origin_frames = int(np.ceil(required_duration * original_fps))
|
||||
if required_duration > total_frames / original_fps:
|
||||
raise ValueError("required_duration must be less than video length")
|
||||
|
||||
if not fixed_start is None and fixed_start >= 0:
|
||||
start_frame = fixed_start
|
||||
else:
|
||||
max_start = total_frames - required_origin_frames
|
||||
if max_start < 0:
|
||||
raise ValueError("video length is too short")
|
||||
start_frame = np.random.randint(0, max_start + 1)
|
||||
start_time = start_frame / original_fps
|
||||
|
||||
end_time = start_time + required_duration
|
||||
time_points = np.linspace(start_time, end_time, num_sample, endpoint=False)
|
||||
|
||||
frame_indices = np.round(np.array(time_points) * original_fps).astype(int)
|
||||
frame_indices = np.clip(frame_indices, 0, total_frames - 1)
|
||||
return frame_indices
|
||||
|
||||
|
||||
def linear_interpolation(features, input_fps, output_fps, output_len=None):
|
||||
"""
|
||||
features: shape=[1, T, 512]
|
||||
input_fps: fps for audio, f_a
|
||||
output_fps: fps for video, f_m
|
||||
output_len: video length
|
||||
"""
|
||||
features = features.transpose(1, 2) # [1, 512, T]
|
||||
seq_len = features.shape[2] / float(input_fps) # T/f_a
|
||||
if output_len is None:
|
||||
output_len = int(seq_len * output_fps) # f_m*T/f_a
|
||||
output_features = F.interpolate(
|
||||
features, size=output_len, align_corners=True,
|
||||
mode='linear') # [1, 512, output_len]
|
||||
return output_features.transpose(1, 2) # [1, output_len, 512]
|
||||
|
||||
|
||||
class AudioEncoder():
|
||||
|
||||
def __init__(self, device='cpu', model_id="facebook/wav2vec2-base-960h"):
|
||||
# load pretrained model
|
||||
self.processor = Wav2Vec2Processor.from_pretrained(model_id)
|
||||
self.model = Wav2Vec2ForCTC.from_pretrained(model_id)
|
||||
|
||||
self.model = self.model.to(device)
|
||||
|
||||
self.video_rate = 30
|
||||
|
||||
def extract_audio_feat(self,
|
||||
audio_path,
|
||||
return_all_layers=False,
|
||||
dtype=torch.float32):
|
||||
audio_input, sample_rate = librosa.load(audio_path, sr=16000)
|
||||
|
||||
input_values = self.processor(
|
||||
audio_input, sampling_rate=sample_rate,
|
||||
return_tensors="pt").input_values
|
||||
|
||||
# INFERENCE
|
||||
|
||||
# retrieve logits & take argmax
|
||||
res = self.model(
|
||||
input_values.to(self.model.device), output_hidden_states=True)
|
||||
if return_all_layers:
|
||||
feat = torch.cat(res.hidden_states)
|
||||
else:
|
||||
feat = res.hidden_states[-1]
|
||||
feat = linear_interpolation(
|
||||
feat, input_fps=50, output_fps=self.video_rate)
|
||||
|
||||
z = feat.to(dtype) # Encoding for the motion
|
||||
return z
|
||||
|
||||
def get_audio_embed_bucket(self,
|
||||
audio_embed,
|
||||
stride=2,
|
||||
batch_frames=12,
|
||||
m=2):
|
||||
num_layers, audio_frame_num, audio_dim = audio_embed.shape
|
||||
|
||||
if num_layers > 1:
|
||||
return_all_layers = True
|
||||
else:
|
||||
return_all_layers = False
|
||||
|
||||
min_batch_num = int(audio_frame_num / (batch_frames * stride)) + 1
|
||||
|
||||
bucket_num = min_batch_num * batch_frames
|
||||
batch_idx = [stride * i for i in range(bucket_num)]
|
||||
batch_audio_eb = []
|
||||
for bi in batch_idx:
|
||||
if bi < audio_frame_num:
|
||||
audio_sample_stride = 2
|
||||
chosen_idx = list(
|
||||
range(bi - m * audio_sample_stride,
|
||||
bi + (m + 1) * audio_sample_stride,
|
||||
audio_sample_stride))
|
||||
chosen_idx = [0 if c < 0 else c for c in chosen_idx]
|
||||
chosen_idx = [
|
||||
audio_frame_num - 1 if c >= audio_frame_num else c
|
||||
for c in chosen_idx
|
||||
]
|
||||
|
||||
if return_all_layers:
|
||||
frame_audio_embed = audio_embed[:, chosen_idx].flatten(
|
||||
start_dim=-2, end_dim=-1)
|
||||
else:
|
||||
frame_audio_embed = audio_embed[0][chosen_idx].flatten()
|
||||
else:
|
||||
frame_audio_embed = \
|
||||
torch.zeros([audio_dim * (2 * m + 1)], device=audio_embed.device) if not return_all_layers \
|
||||
else torch.zeros([num_layers, audio_dim * (2 * m + 1)], device=audio_embed.device)
|
||||
batch_audio_eb.append(frame_audio_embed)
|
||||
batch_audio_eb = torch.cat([c.unsqueeze(0) for c in batch_audio_eb],
|
||||
dim=0)
|
||||
|
||||
return batch_audio_eb, min_batch_num
|
||||
|
||||
def get_audio_embed_bucket_fps(self,
|
||||
audio_embed,
|
||||
fps=16,
|
||||
batch_frames=81,
|
||||
m=0):
|
||||
num_layers, audio_frame_num, audio_dim = audio_embed.shape
|
||||
|
||||
if num_layers > 1:
|
||||
return_all_layers = True
|
||||
else:
|
||||
return_all_layers = False
|
||||
|
||||
scale = self.video_rate / fps
|
||||
|
||||
min_batch_num = int(audio_frame_num / (batch_frames * scale)) + 1
|
||||
|
||||
bucket_num = min_batch_num * batch_frames
|
||||
padd_audio_num = math.ceil(min_batch_num * batch_frames / fps *
|
||||
self.video_rate) - audio_frame_num
|
||||
batch_idx = get_sample_indices(
|
||||
original_fps=self.video_rate,
|
||||
total_frames=audio_frame_num + padd_audio_num,
|
||||
target_fps=fps,
|
||||
num_sample=bucket_num,
|
||||
fixed_start=0)
|
||||
batch_audio_eb = []
|
||||
audio_sample_stride = int(self.video_rate / fps)
|
||||
for bi in batch_idx:
|
||||
if bi < audio_frame_num:
|
||||
|
||||
chosen_idx = list(
|
||||
range(bi - m * audio_sample_stride,
|
||||
bi + (m + 1) * audio_sample_stride,
|
||||
audio_sample_stride))
|
||||
chosen_idx = [0 if c < 0 else c for c in chosen_idx]
|
||||
chosen_idx = [
|
||||
audio_frame_num - 1 if c >= audio_frame_num else c
|
||||
for c in chosen_idx
|
||||
]
|
||||
|
||||
if return_all_layers:
|
||||
frame_audio_embed = audio_embed[:, chosen_idx].flatten(
|
||||
start_dim=-2, end_dim=-1)
|
||||
else:
|
||||
frame_audio_embed = audio_embed[0][chosen_idx].flatten()
|
||||
else:
|
||||
frame_audio_embed = \
|
||||
torch.zeros([audio_dim * (2 * m + 1)], device=audio_embed.device) if not return_all_layers \
|
||||
else torch.zeros([num_layers, audio_dim * (2 * m + 1)], device=audio_embed.device)
|
||||
batch_audio_eb.append(frame_audio_embed)
|
||||
batch_audio_eb = torch.cat([c.unsqueeze(0) for c in batch_audio_eb],
|
||||
dim=0)
|
||||
|
||||
return batch_audio_eb, min_batch_num
|
||||
@@ -0,0 +1,105 @@
|
||||
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||||
import math
|
||||
from typing import Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.cuda.amp as amp
|
||||
import torch.nn as nn
|
||||
from diffusers.models.attention import AdaLayerNorm
|
||||
|
||||
from ..model import WanAttentionBlock, WanCrossAttention
|
||||
from .auxi_blocks import MotionEncoder_tc
|
||||
|
||||
|
||||
class CausalAudioEncoder(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
dim=5120,
|
||||
num_layers=25,
|
||||
out_dim=2048,
|
||||
video_rate=8,
|
||||
num_token=4,
|
||||
need_global=False):
|
||||
super().__init__()
|
||||
self.encoder = MotionEncoder_tc(
|
||||
in_dim=dim,
|
||||
hidden_dim=out_dim,
|
||||
num_heads=num_token,
|
||||
need_global=need_global)
|
||||
weight = torch.ones((1, num_layers, 1, 1)) * 0.01
|
||||
|
||||
self.weights = torch.nn.Parameter(weight)
|
||||
self.act = torch.nn.SiLU()
|
||||
|
||||
def forward(self, features):
|
||||
with amp.autocast(dtype=torch.float32):
|
||||
# features B * num_layers * dim * video_length
|
||||
weights = self.act(self.weights)
|
||||
weights_sum = weights.sum(dim=1, keepdims=True)
|
||||
weighted_feat = ((features * weights) / weights_sum).sum(
|
||||
dim=1) # b dim f
|
||||
weighted_feat = weighted_feat.permute(0, 2, 1) # b f dim
|
||||
res = self.encoder(weighted_feat) # b f n dim
|
||||
|
||||
return res # b f n dim
|
||||
|
||||
|
||||
class AudioCrossAttention(WanCrossAttention):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
|
||||
|
||||
class AudioInjector_WAN(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
all_modules,
|
||||
all_modules_names,
|
||||
dim=2048,
|
||||
num_heads=32,
|
||||
inject_layer=[0, 27],
|
||||
root_net=None,
|
||||
enable_adain=False,
|
||||
adain_dim=2048,
|
||||
need_adain_ont=False):
|
||||
super().__init__()
|
||||
num_injector_layers = len(inject_layer)
|
||||
self.injected_block_id = {}
|
||||
audio_injector_id = 0
|
||||
for mod_name, mod in zip(all_modules_names, all_modules):
|
||||
if isinstance(mod, WanAttentionBlock):
|
||||
for inject_id in inject_layer:
|
||||
if f'transformer_blocks.{inject_id}' in mod_name:
|
||||
self.injected_block_id[inject_id] = audio_injector_id
|
||||
audio_injector_id += 1
|
||||
|
||||
self.injector = nn.ModuleList([
|
||||
AudioCrossAttention(
|
||||
dim=dim,
|
||||
num_heads=num_heads,
|
||||
qk_norm=True,
|
||||
) for _ in range(audio_injector_id)
|
||||
])
|
||||
self.injector_pre_norm_feat = nn.ModuleList([
|
||||
nn.LayerNorm(
|
||||
dim,
|
||||
elementwise_affine=False,
|
||||
eps=1e-6,
|
||||
) for _ in range(audio_injector_id)
|
||||
])
|
||||
self.injector_pre_norm_vec = nn.ModuleList([
|
||||
nn.LayerNorm(
|
||||
dim,
|
||||
elementwise_affine=False,
|
||||
eps=1e-6,
|
||||
) for _ in range(audio_injector_id)
|
||||
])
|
||||
if enable_adain:
|
||||
self.injector_adain_layers = nn.ModuleList([
|
||||
AdaLayerNorm(
|
||||
output_dim=dim * 2, embedding_dim=adain_dim, chunk_dim=1)
|
||||
for _ in range(audio_injector_id)
|
||||
])
|
||||
if need_adain_ont:
|
||||
self.injector_adain_output_layers = nn.ModuleList(
|
||||
[nn.Linear(dim, dim) for _ in range(audio_injector_id)])
|
||||
@@ -0,0 +1,242 @@
|
||||
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||||
import importlib.metadata
|
||||
import math
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models import ModelMixin
|
||||
from diffusers.utils import is_torch_version, logging
|
||||
from einops import rearrange
|
||||
|
||||
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 with shape [(bxs), a, d]
|
||||
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
|
||||
|
||||
# TODO: Maybe force q and k to be float32 to avoid numerical overflow
|
||||
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 MotionEncoder_tc(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
in_dim: int,
|
||||
hidden_dim: int,
|
||||
num_heads=int,
|
||||
need_global=True,
|
||||
dtype=None,
|
||||
device=None):
|
||||
factory_kwargs = {"dtype": dtype, "device": device}
|
||||
super().__init__()
|
||||
|
||||
self.num_heads = num_heads
|
||||
self.need_global = need_global
|
||||
self.conv1_local = CausalConv1d(
|
||||
in_dim, hidden_dim // 4 * num_heads, 3, stride=1)
|
||||
if need_global:
|
||||
self.conv1_global = CausalConv1d(
|
||||
in_dim, hidden_dim // 4, 3, stride=1)
|
||||
self.norm1 = nn.LayerNorm(
|
||||
hidden_dim // 4,
|
||||
elementwise_affine=False,
|
||||
eps=1e-6,
|
||||
**factory_kwargs)
|
||||
self.act = nn.SiLU()
|
||||
self.conv2 = CausalConv1d(hidden_dim // 4, hidden_dim // 2, 3, stride=2)
|
||||
self.conv3 = CausalConv1d(hidden_dim // 2, hidden_dim, 3, stride=2)
|
||||
|
||||
if need_global:
|
||||
self.final_linear = nn.Linear(hidden_dim, hidden_dim,
|
||||
**factory_kwargs)
|
||||
|
||||
self.norm1 = nn.LayerNorm(
|
||||
hidden_dim // 4,
|
||||
elementwise_affine=False,
|
||||
eps=1e-6,
|
||||
**factory_kwargs)
|
||||
|
||||
self.norm2 = nn.LayerNorm(
|
||||
hidden_dim // 2,
|
||||
elementwise_affine=False,
|
||||
eps=1e-6,
|
||||
**factory_kwargs)
|
||||
|
||||
self.norm3 = nn.LayerNorm(
|
||||
hidden_dim, 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')
|
||||
x_ori = x.clone()
|
||||
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 = 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()
|
||||
|
||||
if not self.need_global:
|
||||
return x_local
|
||||
|
||||
x = self.conv1_global(x_ori)
|
||||
x = rearrange(x, 'b c t -> b t c')
|
||||
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.final_linear(x)
|
||||
x = rearrange(x, '(b n) t c -> b t n c', b=b)
|
||||
|
||||
return x, x_local
|
||||
@@ -0,0 +1,906 @@
|
||||
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||||
import math
|
||||
import types
|
||||
from copy import deepcopy
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.cuda.amp as amp
|
||||
import torch.nn as nn
|
||||
from diffusers.configuration_utils import ConfigMixin, register_to_config
|
||||
from diffusers.models.modeling_utils import ModelMixin
|
||||
from einops import rearrange
|
||||
|
||||
from ...distributed.sequence_parallel import (
|
||||
distributed_attention,
|
||||
gather_forward,
|
||||
get_rank,
|
||||
get_world_size,
|
||||
)
|
||||
from ..model import (
|
||||
Head,
|
||||
WanAttentionBlock,
|
||||
WanLayerNorm,
|
||||
WanModel,
|
||||
WanSelfAttention,
|
||||
flash_attention,
|
||||
rope_params,
|
||||
sinusoidal_embedding_1d,
|
||||
)
|
||||
from .audio_utils import AudioInjector_WAN, CausalAudioEncoder
|
||||
from .motioner import FramePackMotioner, MotionerTransformers
|
||||
from .s2v_utils import rope_precompute
|
||||
|
||||
|
||||
def zero_module(module):
|
||||
"""
|
||||
Zero out the parameters of a module and return it.
|
||||
"""
|
||||
for p in module.parameters():
|
||||
p.detach().zero_()
|
||||
return module
|
||||
|
||||
|
||||
def torch_dfs(model: nn.Module, parent_name='root'):
|
||||
module_names, modules = [], []
|
||||
current_name = parent_name if parent_name else 'root'
|
||||
module_names.append(current_name)
|
||||
modules.append(model)
|
||||
|
||||
for name, child in model.named_children():
|
||||
if parent_name:
|
||||
child_name = f'{parent_name}.{name}'
|
||||
else:
|
||||
child_name = name
|
||||
child_modules, child_names = torch_dfs(child, child_name)
|
||||
module_names += child_names
|
||||
modules += child_modules
|
||||
return modules, module_names
|
||||
|
||||
|
||||
@amp.autocast(enabled=False)
|
||||
def rope_apply(x, grid_sizes, freqs, start=None):
|
||||
n, c = x.size(2), x.size(3) // 2
|
||||
# loop over samples
|
||||
output = []
|
||||
for i, _ in enumerate(x):
|
||||
s = x.size(1)
|
||||
x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(
|
||||
s, n, -1, 2))
|
||||
freqs_i = freqs[i, :s]
|
||||
# apply rotary embedding
|
||||
x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
|
||||
x_i = torch.cat([x_i, x[i, s:]])
|
||||
# append to collection
|
||||
output.append(x_i)
|
||||
return torch.stack(output).float()
|
||||
|
||||
|
||||
@amp.autocast(enabled=False)
|
||||
def rope_apply_usp(x, grid_sizes, freqs):
|
||||
s, n, c = x.size(1), x.size(2), x.size(3) // 2
|
||||
# loop over samples
|
||||
output = []
|
||||
for i, _ in enumerate(x):
|
||||
s = x.size(1)
|
||||
# precompute multipliers
|
||||
x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(
|
||||
s, n, -1, 2))
|
||||
freqs_i = freqs[i]
|
||||
freqs_i_rank = freqs_i
|
||||
x_i = torch.view_as_real(x_i * freqs_i_rank).flatten(2)
|
||||
x_i = torch.cat([x_i, x[i, s:]])
|
||||
# append to collection
|
||||
output.append(x_i)
|
||||
return torch.stack(output).float()
|
||||
|
||||
|
||||
def sp_attn_forward_s2v(self,
|
||||
x,
|
||||
seq_lens,
|
||||
grid_sizes,
|
||||
freqs,
|
||||
dtype=torch.bfloat16):
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
half_dtypes = (torch.float16, torch.bfloat16)
|
||||
|
||||
def half(x):
|
||||
return x if x.dtype in half_dtypes else x.to(dtype)
|
||||
|
||||
# query, key, value function
|
||||
def qkv_fn(x):
|
||||
q = self.norm_q(self.q(x)).view(b, s, n, d)
|
||||
k = self.norm_k(self.k(x)).view(b, s, n, d)
|
||||
v = self.v(x).view(b, s, n, d)
|
||||
return q, k, v
|
||||
|
||||
q, k, v = qkv_fn(x)
|
||||
q = rope_apply_usp(q, grid_sizes, freqs)
|
||||
k = rope_apply_usp(k, grid_sizes, freqs)
|
||||
|
||||
x = distributed_attention(
|
||||
half(q),
|
||||
half(k),
|
||||
half(v),
|
||||
seq_lens,
|
||||
window_size=self.window_size,
|
||||
)
|
||||
|
||||
# output
|
||||
x = x.flatten(2)
|
||||
x = self.o(x)
|
||||
return x
|
||||
|
||||
|
||||
class Head_S2V(Head):
|
||||
|
||||
def forward(self, x, e):
|
||||
"""
|
||||
Args:
|
||||
x(Tensor): Shape [B, L1, C]
|
||||
e(Tensor): Shape [B, L1, C]
|
||||
"""
|
||||
assert e.dtype == torch.float32
|
||||
with amp.autocast(dtype=torch.float32):
|
||||
e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
|
||||
x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))
|
||||
return x
|
||||
|
||||
|
||||
class WanS2VSelfAttention(WanSelfAttention):
|
||||
|
||||
def forward(self, x, seq_lens, grid_sizes, freqs):
|
||||
"""
|
||||
Args:
|
||||
x(Tensor): Shape [B, L, num_heads, C / num_heads]
|
||||
seq_lens(Tensor): Shape [B]
|
||||
grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
|
||||
freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
|
||||
"""
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
|
||||
# query, key, value function
|
||||
def qkv_fn(x):
|
||||
q = self.norm_q(self.q(x)).view(b, s, n, d)
|
||||
k = self.norm_k(self.k(x)).view(b, s, n, d)
|
||||
v = self.v(x).view(b, s, n, d)
|
||||
return q, k, v
|
||||
|
||||
q, k, v = qkv_fn(x)
|
||||
|
||||
x = flash_attention(
|
||||
q=rope_apply(q, grid_sizes, freqs),
|
||||
k=rope_apply(k, grid_sizes, freqs),
|
||||
v=v,
|
||||
k_lens=seq_lens,
|
||||
window_size=self.window_size)
|
||||
|
||||
# output
|
||||
x = x.flatten(2)
|
||||
x = self.o(x)
|
||||
return x
|
||||
|
||||
|
||||
class WanS2VAttentionBlock(WanAttentionBlock):
|
||||
|
||||
def __init__(self,
|
||||
dim,
|
||||
ffn_dim,
|
||||
num_heads,
|
||||
window_size=(-1, -1),
|
||||
qk_norm=True,
|
||||
cross_attn_norm=False,
|
||||
eps=1e-6):
|
||||
super().__init__(dim, ffn_dim, num_heads, window_size, qk_norm,
|
||||
cross_attn_norm, eps)
|
||||
self.self_attn = WanS2VSelfAttention(dim, num_heads, window_size,
|
||||
qk_norm, eps)
|
||||
|
||||
def forward(self, x, e, seq_lens, grid_sizes, freqs, context, context_lens):
|
||||
assert e[0].dtype == torch.float32
|
||||
seg_idx = e[1].item()
|
||||
seg_idx = min(max(0, seg_idx), x.size(1))
|
||||
seg_idx = [0, seg_idx, x.size(1)]
|
||||
e = e[0]
|
||||
modulation = self.modulation.unsqueeze(2)
|
||||
with amp.autocast(dtype=torch.float32):
|
||||
e = (modulation + e).chunk(6, dim=1)
|
||||
assert e[0].dtype == torch.float32
|
||||
|
||||
e = [element.squeeze(1) for element in e]
|
||||
norm_x = self.norm1(x).float()
|
||||
parts = []
|
||||
for i in range(2):
|
||||
parts.append(norm_x[:, seg_idx[i]:seg_idx[i + 1]] *
|
||||
(1 + e[1][:, i:i + 1]) + e[0][:, i:i + 1])
|
||||
norm_x = torch.cat(parts, dim=1)
|
||||
# self-attention
|
||||
y = self.self_attn(norm_x, seq_lens, grid_sizes, freqs)
|
||||
with amp.autocast(dtype=torch.float32):
|
||||
z = []
|
||||
for i in range(2):
|
||||
z.append(y[:, seg_idx[i]:seg_idx[i + 1]] * e[2][:, i:i + 1])
|
||||
y = torch.cat(z, dim=1)
|
||||
x = x + y
|
||||
# cross-attention & ffn function
|
||||
def cross_attn_ffn(x, context, context_lens, e):
|
||||
x = x + self.cross_attn(self.norm3(x), context, context_lens)
|
||||
norm2_x = self.norm2(x).float()
|
||||
parts = []
|
||||
for i in range(2):
|
||||
parts.append(norm2_x[:, seg_idx[i]:seg_idx[i + 1]] *
|
||||
(1 + e[4][:, i:i + 1]) + e[3][:, i:i + 1])
|
||||
norm2_x = torch.cat(parts, dim=1)
|
||||
y = self.ffn(norm2_x)
|
||||
with amp.autocast(dtype=torch.float32):
|
||||
z = []
|
||||
for i in range(2):
|
||||
z.append(y[:, seg_idx[i]:seg_idx[i + 1]] * e[5][:, i:i + 1])
|
||||
y = torch.cat(z, dim=1)
|
||||
x = x + y
|
||||
return x
|
||||
|
||||
x = cross_attn_ffn(x, context, context_lens, e)
|
||||
return x
|
||||
|
||||
|
||||
class WanModel_S2V(ModelMixin, ConfigMixin):
|
||||
ignore_for_config = [
|
||||
'args', 'kwargs', 'patch_size', 'cross_attn_norm', 'qk_norm',
|
||||
'text_dim', 'window_size'
|
||||
]
|
||||
_no_split_modules = ['WanS2VAttentionBlock']
|
||||
|
||||
@register_to_config
|
||||
def __init__(
|
||||
self,
|
||||
cond_dim=0,
|
||||
audio_dim=5120,
|
||||
num_audio_token=4,
|
||||
enable_adain=False,
|
||||
adain_mode="attn_norm",
|
||||
audio_inject_layers=[0, 4, 8, 12, 16, 20, 24, 27],
|
||||
zero_init=False,
|
||||
zero_timestep=False,
|
||||
enable_motioner=True,
|
||||
add_last_motion=True,
|
||||
enable_tsm=False,
|
||||
trainable_token_pos_emb=False,
|
||||
motion_token_num=1024,
|
||||
enable_framepack=False, # Mutually exclusive with enable_motioner
|
||||
framepack_drop_mode="drop",
|
||||
model_type='s2v',
|
||||
patch_size=(1, 2, 2),
|
||||
text_len=512,
|
||||
in_dim=16,
|
||||
dim=2048,
|
||||
ffn_dim=8192,
|
||||
freq_dim=256,
|
||||
text_dim=4096,
|
||||
out_dim=16,
|
||||
num_heads=16,
|
||||
num_layers=32,
|
||||
window_size=(-1, -1),
|
||||
qk_norm=True,
|
||||
cross_attn_norm=True,
|
||||
eps=1e-6,
|
||||
*args,
|
||||
**kwargs):
|
||||
super().__init__()
|
||||
|
||||
assert model_type == 's2v'
|
||||
self.model_type = model_type
|
||||
|
||||
self.patch_size = patch_size
|
||||
self.text_len = text_len
|
||||
self.in_dim = in_dim
|
||||
self.dim = dim
|
||||
self.ffn_dim = ffn_dim
|
||||
self.freq_dim = freq_dim
|
||||
self.text_dim = text_dim
|
||||
self.out_dim = out_dim
|
||||
self.num_heads = num_heads
|
||||
self.num_layers = num_layers
|
||||
self.window_size = window_size
|
||||
self.qk_norm = qk_norm
|
||||
self.cross_attn_norm = cross_attn_norm
|
||||
self.eps = eps
|
||||
|
||||
# embeddings
|
||||
self.patch_embedding = nn.Conv3d(
|
||||
in_dim, dim, kernel_size=patch_size, stride=patch_size)
|
||||
self.text_embedding = nn.Sequential(
|
||||
nn.Linear(text_dim, dim), nn.GELU(approximate='tanh'),
|
||||
nn.Linear(dim, dim))
|
||||
|
||||
self.time_embedding = nn.Sequential(
|
||||
nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim))
|
||||
self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6))
|
||||
|
||||
# blocks
|
||||
self.blocks = nn.ModuleList([
|
||||
WanS2VAttentionBlock(dim, ffn_dim, num_heads, window_size, qk_norm,
|
||||
cross_attn_norm, eps)
|
||||
for _ in range(num_layers)
|
||||
])
|
||||
|
||||
# head
|
||||
self.head = Head_S2V(dim, out_dim, patch_size, eps)
|
||||
|
||||
# buffers (don't use register_buffer otherwise dtype will be changed in to())
|
||||
assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
|
||||
d = dim // num_heads
|
||||
self.freqs = torch.cat([
|
||||
rope_params(1024, d - 4 * (d // 6)),
|
||||
rope_params(1024, 2 * (d // 6)),
|
||||
rope_params(1024, 2 * (d // 6))
|
||||
],
|
||||
dim=1)
|
||||
|
||||
# initialize weights
|
||||
self.init_weights()
|
||||
|
||||
self.use_context_parallel = False # will modify in _configure_model func
|
||||
|
||||
if cond_dim > 0:
|
||||
self.cond_encoder = nn.Conv3d(
|
||||
cond_dim,
|
||||
self.dim,
|
||||
kernel_size=self.patch_size,
|
||||
stride=self.patch_size)
|
||||
self.enbale_adain = enable_adain
|
||||
self.casual_audio_encoder = CausalAudioEncoder(
|
||||
dim=audio_dim,
|
||||
out_dim=self.dim,
|
||||
num_token=num_audio_token,
|
||||
need_global=enable_adain)
|
||||
all_modules, all_modules_names = torch_dfs(
|
||||
self.blocks, parent_name="root.transformer_blocks")
|
||||
self.audio_injector = AudioInjector_WAN(
|
||||
all_modules,
|
||||
all_modules_names,
|
||||
dim=self.dim,
|
||||
num_heads=self.num_heads,
|
||||
inject_layer=audio_inject_layers,
|
||||
root_net=self,
|
||||
enable_adain=enable_adain,
|
||||
adain_dim=self.dim,
|
||||
need_adain_ont=adain_mode != "attn_norm",
|
||||
)
|
||||
self.adain_mode = adain_mode
|
||||
|
||||
self.trainable_cond_mask = nn.Embedding(3, self.dim)
|
||||
|
||||
if zero_init:
|
||||
self.zero_init_weights()
|
||||
|
||||
self.zero_timestep = zero_timestep # Whether to assign 0 value timestep to ref/motion
|
||||
|
||||
# init motioner
|
||||
if enable_motioner and enable_framepack:
|
||||
raise ValueError(
|
||||
"enable_motioner and enable_framepack are mutually exclusive, please set one of them to False"
|
||||
)
|
||||
self.enable_motioner = enable_motioner
|
||||
self.add_last_motion = add_last_motion
|
||||
if enable_motioner:
|
||||
motioner_dim = 2048
|
||||
self.motioner = MotionerTransformers(
|
||||
patch_size=(2, 4, 4),
|
||||
dim=motioner_dim,
|
||||
ffn_dim=motioner_dim,
|
||||
freq_dim=256,
|
||||
out_dim=16,
|
||||
num_heads=16,
|
||||
num_layers=13,
|
||||
window_size=(-1, -1),
|
||||
qk_norm=True,
|
||||
cross_attn_norm=False,
|
||||
eps=1e-6,
|
||||
motion_token_num=motion_token_num,
|
||||
enable_tsm=enable_tsm,
|
||||
motion_stride=4,
|
||||
expand_ratio=2,
|
||||
trainable_token_pos_emb=trainable_token_pos_emb,
|
||||
)
|
||||
self.zip_motion_out = torch.nn.Sequential(
|
||||
WanLayerNorm(motioner_dim),
|
||||
zero_module(nn.Linear(motioner_dim, self.dim)))
|
||||
|
||||
self.trainable_token_pos_emb = trainable_token_pos_emb
|
||||
if trainable_token_pos_emb:
|
||||
d = self.dim // self.num_heads
|
||||
x = torch.zeros([1, motion_token_num, self.num_heads, d])
|
||||
x[..., ::2] = 1
|
||||
|
||||
gride_sizes = [[
|
||||
torch.tensor([0, 0, 0]).unsqueeze(0).repeat(1, 1),
|
||||
torch.tensor([
|
||||
1, self.motioner.motion_side_len,
|
||||
self.motioner.motion_side_len
|
||||
]).unsqueeze(0).repeat(1, 1),
|
||||
torch.tensor([
|
||||
1, self.motioner.motion_side_len,
|
||||
self.motioner.motion_side_len
|
||||
]).unsqueeze(0).repeat(1, 1),
|
||||
]]
|
||||
token_freqs = rope_apply(x, gride_sizes, self.freqs)
|
||||
token_freqs = token_freqs[0, :,
|
||||
0].reshape(motion_token_num, -1, 2)
|
||||
token_freqs = token_freqs * 0.01
|
||||
self.token_freqs = torch.nn.Parameter(token_freqs)
|
||||
|
||||
self.enable_framepack = enable_framepack
|
||||
if enable_framepack:
|
||||
self.frame_packer = FramePackMotioner(
|
||||
inner_dim=self.dim,
|
||||
num_heads=self.num_heads,
|
||||
zip_frame_buckets=[1, 2, 16],
|
||||
drop_mode=framepack_drop_mode)
|
||||
|
||||
def zero_init_weights(self):
|
||||
with torch.no_grad():
|
||||
self.trainable_cond_mask = zero_module(self.trainable_cond_mask)
|
||||
if hasattr(self, "cond_encoder"):
|
||||
self.cond_encoder = zero_module(self.cond_encoder)
|
||||
|
||||
for i in range(self.audio_injector.injector.__len__()):
|
||||
self.audio_injector.injector[i].o = zero_module(
|
||||
self.audio_injector.injector[i].o)
|
||||
if self.enbale_adain:
|
||||
self.audio_injector.injector_adain_layers[
|
||||
i].linear = zero_module(
|
||||
self.audio_injector.injector_adain_layers[i].linear)
|
||||
|
||||
def process_motion(self, motion_latents, drop_motion_frames=False):
|
||||
if drop_motion_frames or motion_latents[0].shape[1] == 0:
|
||||
return [], []
|
||||
self.lat_motion_frames = motion_latents[0].shape[1]
|
||||
mot = [self.patch_embedding(m.unsqueeze(0)) for m in motion_latents]
|
||||
batch_size = len(mot)
|
||||
|
||||
mot_remb = []
|
||||
flattern_mot = []
|
||||
for bs in range(batch_size):
|
||||
height, width = mot[bs].shape[3], mot[bs].shape[4]
|
||||
flat_mot = mot[bs].flatten(2).transpose(1, 2).contiguous()
|
||||
motion_grid_sizes = [[
|
||||
torch.tensor([-self.lat_motion_frames, 0,
|
||||
0]).unsqueeze(0).repeat(1, 1),
|
||||
torch.tensor([0, height, width]).unsqueeze(0).repeat(1, 1),
|
||||
torch.tensor([self.lat_motion_frames, height,
|
||||
width]).unsqueeze(0).repeat(1, 1)
|
||||
]]
|
||||
motion_rope_emb = rope_precompute(
|
||||
flat_mot.detach().view(1, flat_mot.shape[1], self.num_heads,
|
||||
self.dim // self.num_heads),
|
||||
motion_grid_sizes,
|
||||
self.freqs,
|
||||
start=None)
|
||||
mot_remb.append(motion_rope_emb)
|
||||
flattern_mot.append(flat_mot)
|
||||
return flattern_mot, mot_remb
|
||||
|
||||
def process_motion_frame_pack(self,
|
||||
motion_latents,
|
||||
drop_motion_frames=False,
|
||||
add_last_motion=2):
|
||||
flattern_mot, mot_remb = self.frame_packer(motion_latents,
|
||||
add_last_motion)
|
||||
if drop_motion_frames:
|
||||
return [m[:, :0] for m in flattern_mot
|
||||
], [m[:, :0] for m in mot_remb]
|
||||
else:
|
||||
return flattern_mot, mot_remb
|
||||
|
||||
def process_motion_transformer_motioner(self,
|
||||
motion_latents,
|
||||
drop_motion_frames=False,
|
||||
add_last_motion=True):
|
||||
batch_size, height, width = len(
|
||||
motion_latents), motion_latents[0].shape[2] // self.patch_size[
|
||||
1], motion_latents[0].shape[3] // self.patch_size[2]
|
||||
|
||||
freqs = self.freqs
|
||||
device = self.patch_embedding.weight.device
|
||||
if freqs.device != device:
|
||||
freqs = freqs.to(device)
|
||||
if self.trainable_token_pos_emb:
|
||||
with amp.autocast(dtype=torch.float64):
|
||||
token_freqs = self.token_freqs.to(torch.float64)
|
||||
token_freqs = token_freqs / token_freqs.norm(
|
||||
dim=-1, keepdim=True)
|
||||
freqs = [freqs, torch.view_as_complex(token_freqs)]
|
||||
|
||||
if not drop_motion_frames and add_last_motion:
|
||||
last_motion_latent = [u[:, -1:] for u in motion_latents]
|
||||
last_mot = [
|
||||
self.patch_embedding(m.unsqueeze(0)) for m in last_motion_latent
|
||||
]
|
||||
last_mot = [m.flatten(2).transpose(1, 2) for m in last_mot]
|
||||
last_mot = torch.cat(last_mot)
|
||||
gride_sizes = [[
|
||||
torch.tensor([-1, 0, 0]).unsqueeze(0).repeat(batch_size, 1),
|
||||
torch.tensor([0, height,
|
||||
width]).unsqueeze(0).repeat(batch_size, 1),
|
||||
torch.tensor([1, height,
|
||||
width]).unsqueeze(0).repeat(batch_size, 1)
|
||||
]]
|
||||
else:
|
||||
last_mot = torch.zeros([batch_size, 0, self.dim],
|
||||
device=motion_latents[0].device,
|
||||
dtype=motion_latents[0].dtype)
|
||||
gride_sizes = []
|
||||
|
||||
zip_motion = self.motioner(motion_latents)
|
||||
zip_motion = self.zip_motion_out(zip_motion)
|
||||
if drop_motion_frames:
|
||||
zip_motion = zip_motion * 0.0
|
||||
zip_motion_grid_sizes = [[
|
||||
torch.tensor([-1, 0, 0]).unsqueeze(0).repeat(batch_size, 1),
|
||||
torch.tensor([
|
||||
0, self.motioner.motion_side_len, self.motioner.motion_side_len
|
||||
]).unsqueeze(0).repeat(batch_size, 1),
|
||||
torch.tensor(
|
||||
[1 if not self.trainable_token_pos_emb else -1, height,
|
||||
width]).unsqueeze(0).repeat(batch_size, 1),
|
||||
]]
|
||||
|
||||
mot = torch.cat([last_mot, zip_motion], dim=1)
|
||||
gride_sizes = gride_sizes + zip_motion_grid_sizes
|
||||
|
||||
motion_rope_emb = rope_precompute(
|
||||
mot.detach().view(batch_size, mot.shape[1], self.num_heads,
|
||||
self.dim // self.num_heads),
|
||||
gride_sizes,
|
||||
freqs,
|
||||
start=None)
|
||||
return [m.unsqueeze(0) for m in mot
|
||||
], [r.unsqueeze(0) for r in motion_rope_emb]
|
||||
|
||||
def inject_motion(self,
|
||||
x,
|
||||
seq_lens,
|
||||
rope_embs,
|
||||
mask_input,
|
||||
motion_latents,
|
||||
drop_motion_frames=False,
|
||||
add_last_motion=True):
|
||||
# inject the motion frames token to the hidden states
|
||||
if self.enable_motioner:
|
||||
mot, mot_remb = self.process_motion_transformer_motioner(
|
||||
motion_latents,
|
||||
drop_motion_frames=drop_motion_frames,
|
||||
add_last_motion=add_last_motion)
|
||||
elif self.enable_framepack:
|
||||
mot, mot_remb = self.process_motion_frame_pack(
|
||||
motion_latents,
|
||||
drop_motion_frames=drop_motion_frames,
|
||||
add_last_motion=add_last_motion)
|
||||
else:
|
||||
mot, mot_remb = self.process_motion(
|
||||
motion_latents, drop_motion_frames=drop_motion_frames)
|
||||
|
||||
if len(mot) > 0:
|
||||
x = [torch.cat([u, m], dim=1) for u, m in zip(x, mot)]
|
||||
seq_lens = seq_lens + torch.tensor([r.size(1) for r in mot],
|
||||
dtype=torch.long)
|
||||
rope_embs = [
|
||||
torch.cat([u, m], dim=1) for u, m in zip(rope_embs, mot_remb)
|
||||
]
|
||||
mask_input = [
|
||||
torch.cat([
|
||||
m, 2 * torch.ones([1, u.shape[1] - m.shape[1]],
|
||||
device=m.device,
|
||||
dtype=m.dtype)
|
||||
],
|
||||
dim=1) for m, u in zip(mask_input, x)
|
||||
]
|
||||
return x, seq_lens, rope_embs, mask_input
|
||||
|
||||
def after_transformer_block(self, block_idx, hidden_states):
|
||||
if block_idx in self.audio_injector.injected_block_id.keys():
|
||||
audio_attn_id = self.audio_injector.injected_block_id[block_idx]
|
||||
audio_emb = self.merged_audio_emb # b f n c
|
||||
num_frames = audio_emb.shape[1]
|
||||
|
||||
if self.use_context_parallel:
|
||||
hidden_states = gather_forward(hidden_states, dim=1)
|
||||
|
||||
input_hidden_states = hidden_states[:, :self.
|
||||
original_seq_len].clone(
|
||||
) # b (f h w) c
|
||||
input_hidden_states = rearrange(
|
||||
input_hidden_states, "b (t n) c -> (b t) n c", t=num_frames)
|
||||
|
||||
if self.enbale_adain and self.adain_mode == "attn_norm":
|
||||
audio_emb_global = self.audio_emb_global
|
||||
audio_emb_global = rearrange(audio_emb_global,
|
||||
"b t n c -> (b t) n c")
|
||||
adain_hidden_states = self.audio_injector.injector_adain_layers[
|
||||
audio_attn_id](
|
||||
input_hidden_states, temb=audio_emb_global[:, 0])
|
||||
attn_hidden_states = adain_hidden_states
|
||||
else:
|
||||
attn_hidden_states = self.audio_injector.injector_pre_norm_feat[
|
||||
audio_attn_id](
|
||||
input_hidden_states)
|
||||
audio_emb = rearrange(
|
||||
audio_emb, "b t n c -> (b t) n c", t=num_frames)
|
||||
attn_audio_emb = audio_emb
|
||||
residual_out = self.audio_injector.injector[audio_attn_id](
|
||||
x=attn_hidden_states,
|
||||
context=attn_audio_emb,
|
||||
context_lens=torch.ones(
|
||||
attn_hidden_states.shape[0],
|
||||
dtype=torch.long,
|
||||
device=attn_hidden_states.device) * attn_audio_emb.shape[1])
|
||||
residual_out = rearrange(
|
||||
residual_out, "(b t) n c -> b (t n) c", t=num_frames)
|
||||
hidden_states[:, :self.
|
||||
original_seq_len] = hidden_states[:, :self.
|
||||
original_seq_len] + residual_out
|
||||
|
||||
if self.use_context_parallel:
|
||||
hidden_states = torch.chunk(
|
||||
hidden_states, get_world_size(), dim=1)[get_rank()]
|
||||
|
||||
return hidden_states
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
t,
|
||||
context,
|
||||
seq_len,
|
||||
ref_latents,
|
||||
motion_latents,
|
||||
cond_states,
|
||||
audio_input=None,
|
||||
motion_frames=[17, 5],
|
||||
add_last_motion=2,
|
||||
drop_motion_frames=False,
|
||||
*extra_args,
|
||||
**extra_kwargs):
|
||||
"""
|
||||
x: A list of videos each with shape [C, T, H, W].
|
||||
t: [B].
|
||||
context: A list of text embeddings each with shape [L, C].
|
||||
seq_len: A list of video token lens, no need for this model.
|
||||
ref_latents A list of reference image for each video with shape [C, 1, H, W].
|
||||
motion_latents A list of motion frames for each video with shape [C, T_m, H, W].
|
||||
cond_states A list of condition frames (i.e. pose) each with shape [C, T, H, W].
|
||||
audio_input The input audio embedding [B, num_wav2vec_layer, C_a, T_a].
|
||||
motion_frames The number of motion frames and motion latents frames encoded by vae, i.e. [17, 5]
|
||||
add_last_motion For the motioner, if add_last_motion > 0, it means that the most recent frame (i.e., the last frame) will be added.
|
||||
For frame packing, the behavior depends on the value of add_last_motion:
|
||||
add_last_motion = 0: Only the farthest part of the latent (i.e., clean_latents_4x) is included.
|
||||
add_last_motion = 1: Both clean_latents_2x and clean_latents_4x are included.
|
||||
add_last_motion = 2: All motion-related latents are used.
|
||||
drop_motion_frames Bool, whether drop the motion frames info
|
||||
"""
|
||||
add_last_motion = self.add_last_motion * add_last_motion
|
||||
audio_input = torch.cat([
|
||||
audio_input[..., 0:1].repeat(1, 1, 1, motion_frames[0]), audio_input
|
||||
],
|
||||
dim=-1)
|
||||
audio_emb_res = self.casual_audio_encoder(audio_input)
|
||||
if self.enbale_adain:
|
||||
audio_emb_global, audio_emb = audio_emb_res
|
||||
self.audio_emb_global = audio_emb_global[:,
|
||||
motion_frames[1]:].clone()
|
||||
else:
|
||||
audio_emb = audio_emb_res
|
||||
self.merged_audio_emb = audio_emb[:, motion_frames[1]:, :]
|
||||
|
||||
device = self.patch_embedding.weight.device
|
||||
|
||||
# embeddings
|
||||
x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
|
||||
# cond states
|
||||
cond = [self.cond_encoder(c.unsqueeze(0)) for c in cond_states]
|
||||
x = [x_ + pose for x_, pose in zip(x, cond)]
|
||||
|
||||
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)
|
||||
|
||||
original_grid_sizes = deepcopy(grid_sizes)
|
||||
grid_sizes = [[torch.zeros_like(grid_sizes), grid_sizes, grid_sizes]]
|
||||
|
||||
# ref and motion
|
||||
self.lat_motion_frames = motion_latents[0].shape[1]
|
||||
|
||||
ref = [self.patch_embedding(r.unsqueeze(0)) for r in ref_latents]
|
||||
batch_size = len(ref)
|
||||
height, width = ref[0].shape[3], ref[0].shape[4]
|
||||
ref_grid_sizes = [[
|
||||
torch.tensor([30, 0, 0]).unsqueeze(0).repeat(batch_size,
|
||||
1), # the start index
|
||||
torch.tensor([31, height,
|
||||
width]).unsqueeze(0).repeat(batch_size,
|
||||
1), # the end index
|
||||
torch.tensor([1, height, width]).unsqueeze(0).repeat(batch_size, 1),
|
||||
] # the range
|
||||
]
|
||||
|
||||
ref = [r.flatten(2).transpose(1, 2) for r in ref] # r: 1 c f h w
|
||||
self.original_seq_len = seq_lens[0]
|
||||
|
||||
seq_lens = seq_lens + torch.tensor([r.size(1) for r in ref],
|
||||
dtype=torch.long)
|
||||
|
||||
grid_sizes = grid_sizes + ref_grid_sizes
|
||||
|
||||
x = [torch.cat([u, r], dim=1) for u, r in zip(x, ref)]
|
||||
|
||||
# Initialize masks to indicate noisy latent, ref latent, and motion latent.
|
||||
# However, at this point, only the first two (noisy and ref latents) are marked;
|
||||
# the marking of motion latent will be implemented inside `inject_motion`.
|
||||
mask_input = [
|
||||
torch.zeros([1, u.shape[1]], dtype=torch.long, device=x[0].device)
|
||||
for u in x
|
||||
]
|
||||
for i in range(len(mask_input)):
|
||||
mask_input[i][:, self.original_seq_len:] = 1
|
||||
|
||||
# compute the rope embeddings for the input
|
||||
x = torch.cat(x)
|
||||
b, s, n, d = x.size(0), x.size(
|
||||
1), self.num_heads, self.dim // self.num_heads
|
||||
self.pre_compute_freqs = rope_precompute(
|
||||
x.detach().view(b, s, n, d), grid_sizes, self.freqs, start=None)
|
||||
|
||||
x = [u.unsqueeze(0) for u in x]
|
||||
self.pre_compute_freqs = [
|
||||
u.unsqueeze(0) for u in self.pre_compute_freqs
|
||||
]
|
||||
|
||||
x, seq_lens, self.pre_compute_freqs, mask_input = self.inject_motion(
|
||||
x,
|
||||
seq_lens,
|
||||
self.pre_compute_freqs,
|
||||
mask_input,
|
||||
motion_latents,
|
||||
drop_motion_frames=drop_motion_frames,
|
||||
add_last_motion=add_last_motion)
|
||||
|
||||
x = torch.cat(x, dim=0)
|
||||
self.pre_compute_freqs = torch.cat(self.pre_compute_freqs, dim=0)
|
||||
mask_input = torch.cat(mask_input, dim=0)
|
||||
|
||||
x = x + self.trainable_cond_mask(mask_input).to(x.dtype)
|
||||
|
||||
# time embeddings
|
||||
if self.zero_timestep:
|
||||
t = torch.cat([t, torch.zeros([1], dtype=t.dtype, device=t.device)])
|
||||
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
|
||||
|
||||
if self.zero_timestep:
|
||||
e = e[:-1]
|
||||
zero_e0 = e0[-1:]
|
||||
e0 = e0[:-1]
|
||||
token_len = x.shape[1]
|
||||
e0 = torch.cat([
|
||||
e0.unsqueeze(2),
|
||||
zero_e0.unsqueeze(2).repeat(e0.size(0), 1, 1, 1)
|
||||
],
|
||||
dim=2)
|
||||
e0 = [e0, self.original_seq_len]
|
||||
else:
|
||||
e0 = e0.unsqueeze(2).repeat(1, 1, 2, 1)
|
||||
e0 = [e0, 0]
|
||||
|
||||
# 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
|
||||
]))
|
||||
|
||||
# grad ckpt args
|
||||
def create_custom_forward(module, return_dict=None):
|
||||
|
||||
def custom_forward(*inputs, **kwargs):
|
||||
if return_dict is not None:
|
||||
return module(*inputs, **kwargs, return_dict=return_dict)
|
||||
else:
|
||||
return module(*inputs, **kwargs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
if self.use_context_parallel:
|
||||
# sharded tensors for long context attn
|
||||
sp_rank = get_rank()
|
||||
x = torch.chunk(x, get_world_size(), dim=1)
|
||||
sq_size = [u.shape[1] for u in x]
|
||||
sq_start_size = sum(sq_size[:sp_rank])
|
||||
x = x[sp_rank]
|
||||
# Confirm the application range of the time embedding in e0[0] for each sequence:
|
||||
# - For tokens before seg_id: apply e0[0][:, :, 0]
|
||||
# - For tokens after seg_id: apply e0[0][:, :, 1]
|
||||
sp_size = x.shape[1]
|
||||
seg_idx = e0[1] - sq_start_size
|
||||
e0[1] = seg_idx
|
||||
|
||||
self.pre_compute_freqs = torch.chunk(
|
||||
self.pre_compute_freqs, get_world_size(), dim=1)
|
||||
self.pre_compute_freqs = self.pre_compute_freqs[sp_rank]
|
||||
|
||||
# arguments
|
||||
kwargs = dict(
|
||||
e=e0,
|
||||
seq_lens=seq_lens,
|
||||
grid_sizes=grid_sizes,
|
||||
freqs=self.pre_compute_freqs,
|
||||
context=context,
|
||||
context_lens=context_lens)
|
||||
for idx, block in enumerate(self.blocks):
|
||||
x = block(x, **kwargs)
|
||||
x = self.after_transformer_block(idx, x)
|
||||
|
||||
# Context Parallel
|
||||
if self.use_context_parallel:
|
||||
x = gather_forward(x.contiguous(), dim=1)
|
||||
# unpatchify
|
||||
x = x[:, :self.original_seq_len]
|
||||
# head
|
||||
x = self.head(x, e)
|
||||
x = self.unpatchify(x, original_grid_sizes)
|
||||
return [u.float() for u in x]
|
||||
|
||||
def unpatchify(self, x, grid_sizes):
|
||||
"""
|
||||
Reconstruct video tensors from patch embeddings.
|
||||
|
||||
Args:
|
||||
x (List[Tensor]):
|
||||
List of patchified features, each with shape [L, C_out * prod(patch_size)]
|
||||
grid_sizes (Tensor):
|
||||
Original spatial-temporal grid dimensions before patching,
|
||||
shape [B, 3] (3 dimensions correspond to F_patches, H_patches, W_patches)
|
||||
|
||||
Returns:
|
||||
List[Tensor]:
|
||||
Reconstructed video tensors with shape [C_out, F, H / 8, W / 8]
|
||||
"""
|
||||
|
||||
c = self.out_dim
|
||||
out = []
|
||||
for u, v in zip(x, grid_sizes.tolist()):
|
||||
u = u[:math.prod(v)].view(*v, *self.patch_size, c)
|
||||
u = torch.einsum('fhwpqrc->cfphqwr', u)
|
||||
u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
|
||||
out.append(u)
|
||||
return out
|
||||
|
||||
def init_weights(self):
|
||||
r"""
|
||||
Initialize model parameters using Xavier initialization.
|
||||
"""
|
||||
|
||||
# basic init
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.xavier_uniform_(m.weight)
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
# init embeddings
|
||||
nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
|
||||
for m in self.text_embedding.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, std=.02)
|
||||
for m in self.time_embedding.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.normal_(m.weight, std=.02)
|
||||
|
||||
# init output layer
|
||||
nn.init.zeros_(self.head.head.weight)
|
||||
@@ -0,0 +1,794 @@
|
||||
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||||
import math
|
||||
from typing import Any, Dict, List, Literal, Optional, Union
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.cuda.amp as amp
|
||||
import torch.nn as nn
|
||||
from diffusers.loaders import FromOriginalModelMixin, PeftAdapterMixin
|
||||
from diffusers.utils import BaseOutput, is_torch_version
|
||||
from einops import rearrange, repeat
|
||||
|
||||
from ..model import flash_attention
|
||||
from .s2v_utils import rope_precompute
|
||||
|
||||
|
||||
def sinusoidal_embedding_1d(dim, position):
|
||||
# preprocess
|
||||
assert dim % 2 == 0
|
||||
half = dim // 2
|
||||
position = position.type(torch.float64)
|
||||
|
||||
# calculation
|
||||
sinusoid = torch.outer(
|
||||
position, torch.pow(10000, -torch.arange(half).to(position).div(half)))
|
||||
x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
|
||||
return x
|
||||
|
||||
|
||||
@amp.autocast(enabled=False)
|
||||
def rope_params(max_seq_len, dim, theta=10000):
|
||||
assert dim % 2 == 0
|
||||
freqs = torch.outer(
|
||||
torch.arange(max_seq_len),
|
||||
1.0 / torch.pow(theta,
|
||||
torch.arange(0, dim, 2).to(torch.float64).div(dim)))
|
||||
freqs = torch.polar(torch.ones_like(freqs), freqs)
|
||||
return freqs
|
||||
|
||||
|
||||
@amp.autocast(enabled=False)
|
||||
def rope_apply(x, grid_sizes, freqs, start=None):
|
||||
n, c = x.size(2), x.size(3) // 2
|
||||
|
||||
# split freqs
|
||||
if type(freqs) is list:
|
||||
trainable_freqs = freqs[1]
|
||||
freqs = freqs[0]
|
||||
freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
|
||||
|
||||
# loop over samples
|
||||
output = []
|
||||
output = x.clone()
|
||||
seq_bucket = [0]
|
||||
if not type(grid_sizes) is list:
|
||||
grid_sizes = [grid_sizes]
|
||||
for g in grid_sizes:
|
||||
if not type(g) is list:
|
||||
g = [torch.zeros_like(g), g]
|
||||
batch_size = g[0].shape[0]
|
||||
for i in range(batch_size):
|
||||
if start is None:
|
||||
f_o, h_o, w_o = g[0][i]
|
||||
else:
|
||||
f_o, h_o, w_o = start[i]
|
||||
|
||||
f, h, w = g[1][i]
|
||||
t_f, t_h, t_w = g[2][i]
|
||||
seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o
|
||||
seq_len = int(seq_f * seq_h * seq_w)
|
||||
if seq_len > 0:
|
||||
if t_f > 0:
|
||||
factor_f, factor_h, factor_w = (t_f / seq_f).item(), (
|
||||
t_h / seq_h).item(), (t_w / seq_w).item()
|
||||
|
||||
if f_o >= 0:
|
||||
f_sam = np.linspace(f_o.item(), (t_f + f_o).item() - 1,
|
||||
seq_f).astype(int).tolist()
|
||||
else:
|
||||
f_sam = np.linspace(-f_o.item(),
|
||||
(-t_f - f_o).item() + 1,
|
||||
seq_f).astype(int).tolist()
|
||||
h_sam = np.linspace(h_o.item(), (t_h + h_o).item() - 1,
|
||||
seq_h).astype(int).tolist()
|
||||
w_sam = np.linspace(w_o.item(), (t_w + w_o).item() - 1,
|
||||
seq_w).astype(int).tolist()
|
||||
|
||||
assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0
|
||||
freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][
|
||||
f_sam].conj()
|
||||
freqs_0 = freqs_0.view(seq_f, 1, 1, -1)
|
||||
|
||||
freqs_i = torch.cat([
|
||||
freqs_0.expand(seq_f, seq_h, seq_w, -1),
|
||||
freqs[1][h_sam].view(1, seq_h, 1, -1).expand(
|
||||
seq_f, seq_h, seq_w, -1),
|
||||
freqs[2][w_sam].view(1, 1, seq_w, -1).expand(
|
||||
seq_f, seq_h, seq_w, -1),
|
||||
],
|
||||
dim=-1).reshape(seq_len, 1, -1)
|
||||
elif t_f < 0:
|
||||
freqs_i = trainable_freqs.unsqueeze(1)
|
||||
# apply rotary embedding
|
||||
# precompute multipliers
|
||||
x_i = torch.view_as_complex(
|
||||
x[i, seq_bucket[-1]:seq_bucket[-1] + seq_len].to(
|
||||
torch.float64).reshape(seq_len, n, -1, 2))
|
||||
x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
|
||||
output[i, seq_bucket[-1]:seq_bucket[-1] + seq_len] = x_i
|
||||
seq_bucket.append(seq_bucket[-1] + seq_len)
|
||||
return output.float()
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
|
||||
def __init__(self, dim, eps=1e-5):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.eps = eps
|
||||
self.weight = nn.Parameter(torch.ones(dim))
|
||||
|
||||
def forward(self, x):
|
||||
return self._norm(x.float()).type_as(x) * self.weight
|
||||
|
||||
def _norm(self, x):
|
||||
return x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)
|
||||
|
||||
|
||||
class LayerNorm(nn.LayerNorm):
|
||||
|
||||
def __init__(self, dim, eps=1e-6, elementwise_affine=False):
|
||||
super().__init__(dim, elementwise_affine=elementwise_affine, eps=eps)
|
||||
|
||||
def forward(self, x):
|
||||
return super().forward(x.float()).type_as(x)
|
||||
|
||||
|
||||
class SelfAttention(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
dim,
|
||||
num_heads,
|
||||
window_size=(-1, -1),
|
||||
qk_norm=True,
|
||||
eps=1e-6):
|
||||
assert dim % num_heads == 0
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.num_heads = num_heads
|
||||
self.head_dim = dim // num_heads
|
||||
self.window_size = window_size
|
||||
self.qk_norm = qk_norm
|
||||
self.eps = eps
|
||||
|
||||
# layers
|
||||
self.q = nn.Linear(dim, dim)
|
||||
self.k = nn.Linear(dim, dim)
|
||||
self.v = nn.Linear(dim, dim)
|
||||
self.o = nn.Linear(dim, dim)
|
||||
self.norm_q = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
|
||||
self.norm_k = RMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
|
||||
|
||||
def forward(self, x, seq_lens, grid_sizes, freqs):
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
|
||||
# query, key, value function
|
||||
def qkv_fn(x):
|
||||
q = self.norm_q(self.q(x)).view(b, s, n, d)
|
||||
k = self.norm_k(self.k(x)).view(b, s, n, d)
|
||||
v = self.v(x).view(b, s, n, d)
|
||||
return q, k, v
|
||||
|
||||
q, k, v = qkv_fn(x)
|
||||
|
||||
x = flash_attention(
|
||||
q=rope_apply(q, grid_sizes, freqs),
|
||||
k=rope_apply(k, grid_sizes, freqs),
|
||||
v=v,
|
||||
k_lens=seq_lens,
|
||||
window_size=self.window_size)
|
||||
|
||||
# output
|
||||
x = x.flatten(2)
|
||||
x = self.o(x)
|
||||
return x
|
||||
|
||||
|
||||
class SwinSelfAttention(SelfAttention):
|
||||
|
||||
def forward(self, x, seq_lens, grid_sizes, freqs):
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
assert b == 1, 'Only support batch_size 1'
|
||||
|
||||
# query, key, value function
|
||||
def qkv_fn(x):
|
||||
q = self.norm_q(self.q(x)).view(b, s, n, d)
|
||||
k = self.norm_k(self.k(x)).view(b, s, n, d)
|
||||
v = self.v(x).view(b, s, n, d)
|
||||
return q, k, v
|
||||
|
||||
q, k, v = qkv_fn(x)
|
||||
|
||||
q = rope_apply(q, grid_sizes, freqs)
|
||||
k = rope_apply(k, grid_sizes, freqs)
|
||||
T, H, W = grid_sizes[0].tolist()
|
||||
|
||||
q = rearrange(q, 'b (t h w) n d -> (b t) (h w) n d', t=T, h=H, w=W)
|
||||
k = rearrange(k, 'b (t h w) n d -> (b t) (h w) n d', t=T, h=H, w=W)
|
||||
v = rearrange(v, 'b (t h w) n d -> (b t) (h w) n d', t=T, h=H, w=W)
|
||||
|
||||
ref_q = q[-1:]
|
||||
q = q[:-1]
|
||||
|
||||
ref_k = repeat(
|
||||
k[-1:], "1 s n d -> t s n d", t=k.shape[0] - 1) # t hw n d
|
||||
k = k[:-1]
|
||||
k = torch.cat([k[:1], k, k[-1:]])
|
||||
k = torch.cat([k[1:-1], k[2:], k[:-2], ref_k], dim=1) # (bt) (3hw) n d
|
||||
|
||||
ref_v = repeat(v[-1:], "1 s n d -> t s n d", t=v.shape[0] - 1)
|
||||
v = v[:-1]
|
||||
v = torch.cat([v[:1], v, v[-1:]])
|
||||
v = torch.cat([v[1:-1], v[2:], v[:-2], ref_v], dim=1)
|
||||
|
||||
# q: b (t h w) n d
|
||||
# k: b (t h w) n d
|
||||
out = flash_attention(
|
||||
q=q,
|
||||
k=k,
|
||||
v=v,
|
||||
# k_lens=torch.tensor([k.shape[1]] * k.shape[0], device=x.device, dtype=torch.long),
|
||||
window_size=self.window_size)
|
||||
out = torch.cat([out, ref_v[:1]], axis=0)
|
||||
out = rearrange(out, '(b t) (h w) n d -> b (t h w) n d', t=T, h=H, w=W)
|
||||
x = out
|
||||
|
||||
# output
|
||||
x = x.flatten(2)
|
||||
x = self.o(x)
|
||||
return x
|
||||
|
||||
|
||||
#Fix the reference frame RoPE to 1,H,W.
|
||||
#Set the current frame RoPE to 1.
|
||||
#Set the previous frame RoPE to 0.
|
||||
class CasualSelfAttention(SelfAttention):
|
||||
|
||||
def forward(self, x, seq_lens, grid_sizes, freqs):
|
||||
shifting = 3
|
||||
b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
|
||||
assert b == 1, 'Only support batch_size 1'
|
||||
|
||||
# query, key, value function
|
||||
def qkv_fn(x):
|
||||
q = self.norm_q(self.q(x)).view(b, s, n, d)
|
||||
k = self.norm_k(self.k(x)).view(b, s, n, d)
|
||||
v = self.v(x).view(b, s, n, d)
|
||||
return q, k, v
|
||||
|
||||
q, k, v = qkv_fn(x)
|
||||
|
||||
T, H, W = grid_sizes[0].tolist()
|
||||
|
||||
q = rearrange(q, 'b (t h w) n d -> (b t) (h w) n d', t=T, h=H, w=W)
|
||||
k = rearrange(k, 'b (t h w) n d -> (b t) (h w) n d', t=T, h=H, w=W)
|
||||
v = rearrange(v, 'b (t h w) n d -> (b t) (h w) n d', t=T, h=H, w=W)
|
||||
|
||||
ref_q = q[-1:]
|
||||
q = q[:-1]
|
||||
|
||||
grid_sizes = torch.tensor([[1, H, W]] * q.shape[0], dtype=torch.long)
|
||||
start = [[shifting, 0, 0]] * q.shape[0]
|
||||
q = rope_apply(q, grid_sizes, freqs, start=start)
|
||||
|
||||
ref_k = k[-1:]
|
||||
grid_sizes = torch.tensor([[1, H, W]], dtype=torch.long)
|
||||
# start = [[shifting, H, W]]
|
||||
|
||||
start = [[shifting + 10, 0, 0]]
|
||||
ref_k = rope_apply(ref_k, grid_sizes, freqs, start)
|
||||
ref_k = repeat(
|
||||
ref_k, "1 s n d -> t s n d", t=k.shape[0] - 1) # t hw n d
|
||||
|
||||
k = k[:-1]
|
||||
k = torch.cat([*([k[:1]] * shifting), k])
|
||||
cat_k = []
|
||||
for i in range(shifting):
|
||||
cat_k.append(k[i:i - shifting])
|
||||
cat_k.append(k[shifting:])
|
||||
k = torch.cat(cat_k, dim=1) # (bt) (3hw) n d
|
||||
|
||||
grid_sizes = torch.tensor(
|
||||
[[shifting + 1, H, W]] * q.shape[0], dtype=torch.long)
|
||||
k = rope_apply(k, grid_sizes, freqs)
|
||||
k = torch.cat([k, ref_k], dim=1)
|
||||
|
||||
ref_v = repeat(v[-1:], "1 s n d -> t s n d", t=q.shape[0]) # t hw n d
|
||||
v = v[:-1]
|
||||
v = torch.cat([*([v[:1]] * shifting), v])
|
||||
cat_v = []
|
||||
for i in range(shifting):
|
||||
cat_v.append(v[i:i - shifting])
|
||||
cat_v.append(v[shifting:])
|
||||
v = torch.cat(cat_v, dim=1) # (bt) (3hw) n d
|
||||
v = torch.cat([v, ref_v], dim=1)
|
||||
|
||||
# q: b (t h w) n d
|
||||
# k: b (t h w) n d
|
||||
outs = []
|
||||
for i in range(q.shape[0]):
|
||||
out = flash_attention(
|
||||
q=q[i:i + 1],
|
||||
k=k[i:i + 1],
|
||||
v=v[i:i + 1],
|
||||
window_size=self.window_size)
|
||||
outs.append(out)
|
||||
out = torch.cat(outs, dim=0)
|
||||
out = torch.cat([out, ref_v[:1]], axis=0)
|
||||
out = rearrange(out, '(b t) (h w) n d -> b (t h w) n d', t=T, h=H, w=W)
|
||||
x = out
|
||||
|
||||
# output
|
||||
x = x.flatten(2)
|
||||
x = self.o(x)
|
||||
return x
|
||||
|
||||
|
||||
class MotionerAttentionBlock(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
dim,
|
||||
ffn_dim,
|
||||
num_heads,
|
||||
window_size=(-1, -1),
|
||||
qk_norm=True,
|
||||
cross_attn_norm=False,
|
||||
eps=1e-6,
|
||||
self_attn_block="SelfAttention"):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.ffn_dim = ffn_dim
|
||||
self.num_heads = num_heads
|
||||
self.window_size = window_size
|
||||
self.qk_norm = qk_norm
|
||||
self.cross_attn_norm = cross_attn_norm
|
||||
self.eps = eps
|
||||
|
||||
# layers
|
||||
self.norm1 = LayerNorm(dim, eps)
|
||||
if self_attn_block == "SelfAttention":
|
||||
self.self_attn = SelfAttention(dim, num_heads, window_size, qk_norm,
|
||||
eps)
|
||||
elif self_attn_block == "SwinSelfAttention":
|
||||
self.self_attn = SwinSelfAttention(dim, num_heads, window_size,
|
||||
qk_norm, eps)
|
||||
elif self_attn_block == "CasualSelfAttention":
|
||||
self.self_attn = CasualSelfAttention(dim, num_heads, window_size,
|
||||
qk_norm, eps)
|
||||
|
||||
self.norm2 = LayerNorm(dim, eps)
|
||||
self.ffn = nn.Sequential(
|
||||
nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
|
||||
nn.Linear(ffn_dim, dim))
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
seq_lens,
|
||||
grid_sizes,
|
||||
freqs,
|
||||
):
|
||||
# self-attention
|
||||
y = self.self_attn(self.norm1(x).float(), seq_lens, grid_sizes, freqs)
|
||||
x = x + y
|
||||
y = self.ffn(self.norm2(x).float())
|
||||
x = x + y
|
||||
return x
|
||||
|
||||
|
||||
class Head(nn.Module):
|
||||
|
||||
def __init__(self, dim, out_dim, patch_size, eps=1e-6):
|
||||
super().__init__()
|
||||
self.dim = dim
|
||||
self.out_dim = out_dim
|
||||
self.patch_size = patch_size
|
||||
self.eps = eps
|
||||
|
||||
# layers
|
||||
out_dim = math.prod(patch_size) * out_dim
|
||||
self.norm = LayerNorm(dim, eps)
|
||||
self.head = nn.Linear(dim, out_dim)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.head(self.norm(x))
|
||||
return x
|
||||
|
||||
|
||||
class MotionerTransformers(nn.Module, PeftAdapterMixin):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
patch_size=(1, 2, 2),
|
||||
in_dim=16,
|
||||
dim=2048,
|
||||
ffn_dim=8192,
|
||||
freq_dim=256,
|
||||
out_dim=16,
|
||||
num_heads=16,
|
||||
num_layers=32,
|
||||
window_size=(-1, -1),
|
||||
qk_norm=True,
|
||||
cross_attn_norm=False,
|
||||
eps=1e-6,
|
||||
self_attn_block="SelfAttention",
|
||||
motion_token_num=1024,
|
||||
enable_tsm=False,
|
||||
motion_stride=4,
|
||||
expand_ratio=2,
|
||||
trainable_token_pos_emb=False,
|
||||
):
|
||||
super().__init__()
|
||||
self.patch_size = patch_size
|
||||
self.in_dim = in_dim
|
||||
self.dim = dim
|
||||
self.ffn_dim = ffn_dim
|
||||
self.freq_dim = freq_dim
|
||||
self.out_dim = out_dim
|
||||
self.num_heads = num_heads
|
||||
self.num_layers = num_layers
|
||||
self.window_size = window_size
|
||||
self.qk_norm = qk_norm
|
||||
self.cross_attn_norm = cross_attn_norm
|
||||
self.eps = eps
|
||||
|
||||
self.enable_tsm = enable_tsm
|
||||
self.motion_stride = motion_stride
|
||||
self.expand_ratio = expand_ratio
|
||||
self.sample_c = self.patch_size[0]
|
||||
|
||||
# embeddings
|
||||
self.patch_embedding = nn.Conv3d(
|
||||
in_dim, dim, kernel_size=patch_size, stride=patch_size)
|
||||
|
||||
# blocks
|
||||
self.blocks = nn.ModuleList([
|
||||
MotionerAttentionBlock(
|
||||
dim,
|
||||
ffn_dim,
|
||||
num_heads,
|
||||
window_size,
|
||||
qk_norm,
|
||||
cross_attn_norm,
|
||||
eps,
|
||||
self_attn_block=self_attn_block) for _ in range(num_layers)
|
||||
])
|
||||
|
||||
# buffers (don't use register_buffer otherwise dtype will be changed in to())
|
||||
assert (dim % num_heads) == 0 and (dim // num_heads) % 2 == 0
|
||||
d = dim // num_heads
|
||||
self.freqs = torch.cat([
|
||||
rope_params(1024, d - 4 * (d // 6)),
|
||||
rope_params(1024, 2 * (d // 6)),
|
||||
rope_params(1024, 2 * (d // 6))
|
||||
],
|
||||
dim=1)
|
||||
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
self.motion_side_len = int(math.sqrt(motion_token_num))
|
||||
assert self.motion_side_len**2 == motion_token_num
|
||||
self.token = nn.Parameter(
|
||||
torch.zeros(1, motion_token_num, dim).contiguous())
|
||||
|
||||
self.trainable_token_pos_emb = trainable_token_pos_emb
|
||||
if trainable_token_pos_emb:
|
||||
x = torch.zeros([1, motion_token_num, num_heads, d])
|
||||
x[..., ::2] = 1
|
||||
|
||||
gride_sizes = [[
|
||||
torch.tensor([0, 0, 0]).unsqueeze(0).repeat(1, 1),
|
||||
torch.tensor([1, self.motion_side_len,
|
||||
self.motion_side_len]).unsqueeze(0).repeat(1, 1),
|
||||
torch.tensor([1, self.motion_side_len,
|
||||
self.motion_side_len]).unsqueeze(0).repeat(1, 1),
|
||||
]]
|
||||
token_freqs = rope_apply(x, gride_sizes, self.freqs)
|
||||
token_freqs = token_freqs[0, :, 0].reshape(motion_token_num, -1, 2)
|
||||
token_freqs = token_freqs * 0.01
|
||||
self.token_freqs = torch.nn.Parameter(token_freqs)
|
||||
|
||||
def after_patch_embedding(self, x):
|
||||
return x
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x,
|
||||
):
|
||||
"""
|
||||
x: A list of videos each with shape [C, T, H, W].
|
||||
t: [B].
|
||||
context: A list of text embeddings each with shape [L, C].
|
||||
"""
|
||||
# params
|
||||
motion_frames = x[0].shape[1]
|
||||
device = self.patch_embedding.weight.device
|
||||
freqs = self.freqs
|
||||
if freqs.device != device:
|
||||
freqs = freqs.to(device)
|
||||
|
||||
if self.trainable_token_pos_emb:
|
||||
with amp.autocast(dtype=torch.float64):
|
||||
token_freqs = self.token_freqs.to(torch.float64)
|
||||
token_freqs = token_freqs / token_freqs.norm(
|
||||
dim=-1, keepdim=True)
|
||||
freqs = [freqs, torch.view_as_complex(token_freqs)]
|
||||
|
||||
if self.enable_tsm:
|
||||
sample_idx = [
|
||||
sample_indices(
|
||||
u.shape[1],
|
||||
stride=self.motion_stride,
|
||||
expand_ratio=self.expand_ratio,
|
||||
c=self.sample_c) for u in x
|
||||
]
|
||||
x = [
|
||||
torch.flip(torch.flip(u, [1])[:, idx], [1])
|
||||
for idx, u in zip(sample_idx, x)
|
||||
]
|
||||
|
||||
# embeddings
|
||||
x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
|
||||
x = self.after_patch_embedding(x)
|
||||
|
||||
seq_f, seq_h, seq_w = x[0].shape[-3:]
|
||||
batch_size = len(x)
|
||||
if not self.enable_tsm:
|
||||
grid_sizes = torch.stack(
|
||||
[torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
|
||||
grid_sizes = [[
|
||||
torch.zeros_like(grid_sizes), grid_sizes, grid_sizes
|
||||
]]
|
||||
seq_f = 0
|
||||
else:
|
||||
grid_sizes = []
|
||||
for idx in sample_idx[0][::-1][::self.sample_c]:
|
||||
tsm_frame_grid_sizes = [[
|
||||
torch.tensor([idx, 0,
|
||||
0]).unsqueeze(0).repeat(batch_size, 1),
|
||||
torch.tensor([idx + 1, seq_h,
|
||||
seq_w]).unsqueeze(0).repeat(batch_size, 1),
|
||||
torch.tensor([1, seq_h,
|
||||
seq_w]).unsqueeze(0).repeat(batch_size, 1),
|
||||
]]
|
||||
grid_sizes += tsm_frame_grid_sizes
|
||||
seq_f = sample_idx[0][-1] + 1
|
||||
|
||||
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)
|
||||
x = torch.cat([u for u in x])
|
||||
|
||||
batch_size = len(x)
|
||||
|
||||
token_grid_sizes = [[
|
||||
torch.tensor([seq_f, 0, 0]).unsqueeze(0).repeat(batch_size, 1),
|
||||
torch.tensor(
|
||||
[seq_f + 1, self.motion_side_len,
|
||||
self.motion_side_len]).unsqueeze(0).repeat(batch_size, 1),
|
||||
torch.tensor(
|
||||
[1 if not self.trainable_token_pos_emb else -1, seq_h,
|
||||
seq_w]).unsqueeze(0).repeat(batch_size, 1),
|
||||
] # 第三行代表rope emb的想要覆盖到的范围
|
||||
]
|
||||
|
||||
grid_sizes = grid_sizes + token_grid_sizes
|
||||
token_unpatch_grid_sizes = torch.stack([
|
||||
torch.tensor([1, 32, 32], dtype=torch.long)
|
||||
for b in range(batch_size)
|
||||
])
|
||||
token_len = self.token.shape[1]
|
||||
token = self.token.clone().repeat(x.shape[0], 1, 1).contiguous()
|
||||
seq_lens = seq_lens + torch.tensor([t.size(0) for t in token],
|
||||
dtype=torch.long)
|
||||
x = torch.cat([x, token], dim=1)
|
||||
# arguments
|
||||
kwargs = dict(
|
||||
seq_lens=seq_lens,
|
||||
grid_sizes=grid_sizes,
|
||||
freqs=freqs,
|
||||
)
|
||||
|
||||
# grad ckpt args
|
||||
def create_custom_forward(module, return_dict=None):
|
||||
|
||||
def custom_forward(*inputs, **kwargs):
|
||||
if return_dict is not None:
|
||||
return module(*inputs, **kwargs, return_dict=return_dict)
|
||||
else:
|
||||
return module(*inputs, **kwargs)
|
||||
|
||||
return custom_forward
|
||||
|
||||
ckpt_kwargs: Dict[str, Any] = ({
|
||||
"use_reentrant": False
|
||||
} if is_torch_version(">=", "1.11.0") else {})
|
||||
|
||||
for idx, block in enumerate(self.blocks):
|
||||
if self.training and self.gradient_checkpointing:
|
||||
x = torch.utils.checkpoint.checkpoint(
|
||||
create_custom_forward(block),
|
||||
x,
|
||||
**kwargs,
|
||||
**ckpt_kwargs,
|
||||
)
|
||||
else:
|
||||
x = block(x, **kwargs)
|
||||
# head
|
||||
out = x[:, -token_len:]
|
||||
return out
|
||||
|
||||
def unpatchify(self, x, grid_sizes):
|
||||
c = self.out_dim
|
||||
out = []
|
||||
for u, v in zip(x, grid_sizes.tolist()):
|
||||
u = u[:math.prod(v)].view(*v, *self.patch_size, c)
|
||||
u = torch.einsum('fhwpqrc->cfphqwr', u)
|
||||
u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
|
||||
out.append(u)
|
||||
return out
|
||||
|
||||
def init_weights(self):
|
||||
# basic init
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.xavier_uniform_(m.weight)
|
||||
if m.bias is not None:
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
# init embeddings
|
||||
nn.init.xavier_uniform_(self.patch_embedding.weight.flatten(1))
|
||||
|
||||
|
||||
class FramePackMotioner(nn.Module):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
inner_dim=1024,
|
||||
num_heads=16, # Used to indicate the number of heads in the backbone network; unrelated to this module's design
|
||||
zip_frame_buckets=[
|
||||
1, 2, 16
|
||||
], # Three numbers representing the number of frames sampled for patch operations from the nearest to the farthest frames
|
||||
drop_mode="drop", # If not "drop", it will use "padd", meaning padding instead of deletion
|
||||
*args,
|
||||
**kwargs):
|
||||
super().__init__(*args, **kwargs)
|
||||
self.proj = nn.Conv3d(
|
||||
16, inner_dim, kernel_size=(1, 2, 2), stride=(1, 2, 2))
|
||||
self.proj_2x = nn.Conv3d(
|
||||
16, inner_dim, kernel_size=(2, 4, 4), stride=(2, 4, 4))
|
||||
self.proj_4x = nn.Conv3d(
|
||||
16, inner_dim, kernel_size=(4, 8, 8), stride=(4, 8, 8))
|
||||
self.zip_frame_buckets = torch.tensor(
|
||||
zip_frame_buckets, dtype=torch.long)
|
||||
|
||||
self.inner_dim = inner_dim
|
||||
self.num_heads = num_heads
|
||||
|
||||
assert (inner_dim %
|
||||
num_heads) == 0 and (inner_dim // num_heads) % 2 == 0
|
||||
d = inner_dim // num_heads
|
||||
self.freqs = torch.cat([
|
||||
rope_params(1024, d - 4 * (d // 6)),
|
||||
rope_params(1024, 2 * (d // 6)),
|
||||
rope_params(1024, 2 * (d // 6))
|
||||
],
|
||||
dim=1)
|
||||
self.drop_mode = drop_mode
|
||||
|
||||
def forward(self, motion_latents, add_last_motion=2):
|
||||
motion_frames = motion_latents[0].shape[1]
|
||||
mot = []
|
||||
mot_remb = []
|
||||
for m in motion_latents:
|
||||
lat_height, lat_width = m.shape[2], m.shape[3]
|
||||
padd_lat = torch.zeros(16, self.zip_frame_buckets.sum(), lat_height,
|
||||
lat_width).to(
|
||||
device=m.device, dtype=m.dtype)
|
||||
overlap_frame = min(padd_lat.shape[1], m.shape[1])
|
||||
if overlap_frame > 0:
|
||||
padd_lat[:, -overlap_frame:] = m[:, -overlap_frame:]
|
||||
|
||||
if add_last_motion < 2 and self.drop_mode != "drop":
|
||||
zero_end_frame = self.zip_frame_buckets[:self.zip_frame_buckets.
|
||||
__len__() -
|
||||
add_last_motion -
|
||||
1].sum()
|
||||
padd_lat[:, -zero_end_frame:] = 0
|
||||
|
||||
padd_lat = padd_lat.unsqueeze(0)
|
||||
clean_latents_4x, clean_latents_2x, clean_latents_post = padd_lat[:, :, -self.zip_frame_buckets.sum(
|
||||
):, :, :].split(
|
||||
list(self.zip_frame_buckets)[::-1], dim=2) # 16, 2 ,1
|
||||
|
||||
# patchfy
|
||||
clean_latents_post = self.proj(clean_latents_post).flatten(
|
||||
2).transpose(1, 2)
|
||||
clean_latents_2x = self.proj_2x(clean_latents_2x).flatten(
|
||||
2).transpose(1, 2)
|
||||
clean_latents_4x = self.proj_4x(clean_latents_4x).flatten(
|
||||
2).transpose(1, 2)
|
||||
|
||||
if add_last_motion < 2 and self.drop_mode == "drop":
|
||||
clean_latents_post = clean_latents_post[:, :
|
||||
0] if add_last_motion < 2 else clean_latents_post
|
||||
clean_latents_2x = clean_latents_2x[:, :
|
||||
0] if add_last_motion < 1 else clean_latents_2x
|
||||
|
||||
motion_lat = torch.cat(
|
||||
[clean_latents_post, clean_latents_2x, clean_latents_4x], dim=1)
|
||||
|
||||
# rope
|
||||
start_time_id = -(self.zip_frame_buckets[:1].sum())
|
||||
end_time_id = start_time_id + self.zip_frame_buckets[0]
|
||||
grid_sizes = [] if add_last_motion < 2 and self.drop_mode == "drop" else \
|
||||
[
|
||||
[torch.tensor([start_time_id, 0, 0]).unsqueeze(0).repeat(1, 1),
|
||||
torch.tensor([end_time_id, lat_height // 2, lat_width // 2]).unsqueeze(0).repeat(1, 1),
|
||||
torch.tensor([self.zip_frame_buckets[0], lat_height // 2, lat_width // 2]).unsqueeze(0).repeat(1, 1), ]
|
||||
]
|
||||
|
||||
start_time_id = -(self.zip_frame_buckets[:2].sum())
|
||||
end_time_id = start_time_id + self.zip_frame_buckets[1] // 2
|
||||
grid_sizes_2x = [] if add_last_motion < 1 and self.drop_mode == "drop" else \
|
||||
[
|
||||
[torch.tensor([start_time_id, 0, 0]).unsqueeze(0).repeat(1, 1),
|
||||
torch.tensor([end_time_id, lat_height // 4, lat_width // 4]).unsqueeze(0).repeat(1, 1),
|
||||
torch.tensor([self.zip_frame_buckets[1], lat_height // 2, lat_width // 2]).unsqueeze(0).repeat(1, 1), ]
|
||||
]
|
||||
|
||||
start_time_id = -(self.zip_frame_buckets[:3].sum())
|
||||
end_time_id = start_time_id + self.zip_frame_buckets[2] // 4
|
||||
grid_sizes_4x = [[
|
||||
torch.tensor([start_time_id, 0, 0]).unsqueeze(0).repeat(1, 1),
|
||||
torch.tensor([end_time_id, lat_height // 8,
|
||||
lat_width // 8]).unsqueeze(0).repeat(1, 1),
|
||||
torch.tensor([
|
||||
self.zip_frame_buckets[2], lat_height // 2, lat_width // 2
|
||||
]).unsqueeze(0).repeat(1, 1),
|
||||
]]
|
||||
|
||||
grid_sizes = grid_sizes + grid_sizes_2x + grid_sizes_4x
|
||||
|
||||
motion_rope_emb = rope_precompute(
|
||||
motion_lat.detach().view(1, motion_lat.shape[1], self.num_heads,
|
||||
self.inner_dim // self.num_heads),
|
||||
grid_sizes,
|
||||
self.freqs,
|
||||
start=None)
|
||||
|
||||
mot.append(motion_lat)
|
||||
mot_remb.append(motion_rope_emb)
|
||||
return mot, mot_remb
|
||||
|
||||
|
||||
def sample_indices(N, stride, expand_ratio, c):
|
||||
indices = []
|
||||
current_start = 0
|
||||
|
||||
while current_start < N:
|
||||
bucket_width = int(stride * (expand_ratio**(len(indices) / stride)))
|
||||
|
||||
interval = int(bucket_width / stride * c)
|
||||
current_end = min(N, current_start + bucket_width)
|
||||
bucket_samples = []
|
||||
for i in range(current_end - 1, current_start - 1, -interval):
|
||||
for near in range(c):
|
||||
bucket_samples.append(i - near)
|
||||
|
||||
indices += bucket_samples[::-1]
|
||||
current_start += bucket_width
|
||||
|
||||
return indices
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
device = "cuda"
|
||||
model = FramePackMotioner(inner_dim=1024)
|
||||
batch_size = 2
|
||||
num_frame, height, width = (28, 32, 32)
|
||||
single_input = torch.ones([16, num_frame, height, width], device=device)
|
||||
for i in range(num_frame):
|
||||
single_input[:, num_frame - 1 - i] *= i
|
||||
x = [single_input] * batch_size
|
||||
model.forward(x)
|
||||
@@ -0,0 +1,70 @@
|
||||
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def rope_precompute(x, grid_sizes, freqs, start=None):
|
||||
b, s, n, c = x.size(0), x.size(1), x.size(2), x.size(3) // 2
|
||||
|
||||
# split freqs
|
||||
if type(freqs) is list:
|
||||
trainable_freqs = freqs[1]
|
||||
freqs = freqs[0]
|
||||
freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
|
||||
|
||||
# loop over samples
|
||||
output = torch.view_as_complex(x.detach().reshape(b, s, n, -1,
|
||||
2).to(torch.float64))
|
||||
seq_bucket = [0]
|
||||
if not type(grid_sizes) is list:
|
||||
grid_sizes = [grid_sizes]
|
||||
for g in grid_sizes:
|
||||
if not type(g) is list:
|
||||
g = [torch.zeros_like(g), g]
|
||||
batch_size = g[0].shape[0]
|
||||
for i in range(batch_size):
|
||||
if start is None:
|
||||
f_o, h_o, w_o = g[0][i]
|
||||
else:
|
||||
f_o, h_o, w_o = start[i]
|
||||
|
||||
f, h, w = g[1][i]
|
||||
t_f, t_h, t_w = g[2][i]
|
||||
seq_f, seq_h, seq_w = f - f_o, h - h_o, w - w_o
|
||||
seq_len = int(seq_f * seq_h * seq_w)
|
||||
if seq_len > 0:
|
||||
if t_f > 0:
|
||||
factor_f, factor_h, factor_w = (t_f / seq_f).item(), (
|
||||
t_h / seq_h).item(), (t_w / seq_w).item()
|
||||
# Generate a list of seq_f integers starting from f_o and ending at math.ceil(factor_f * seq_f.item() + f_o.item())
|
||||
if f_o >= 0:
|
||||
f_sam = np.linspace(f_o.item(), (t_f + f_o).item() - 1,
|
||||
seq_f).astype(int).tolist()
|
||||
else:
|
||||
f_sam = np.linspace(-f_o.item(),
|
||||
(-t_f - f_o).item() + 1,
|
||||
seq_f).astype(int).tolist()
|
||||
h_sam = np.linspace(h_o.item(), (t_h + h_o).item() - 1,
|
||||
seq_h).astype(int).tolist()
|
||||
w_sam = np.linspace(w_o.item(), (t_w + w_o).item() - 1,
|
||||
seq_w).astype(int).tolist()
|
||||
|
||||
assert f_o * f >= 0 and h_o * h >= 0 and w_o * w >= 0
|
||||
freqs_0 = freqs[0][f_sam] if f_o >= 0 else freqs[0][
|
||||
f_sam].conj()
|
||||
freqs_0 = freqs_0.view(seq_f, 1, 1, -1)
|
||||
|
||||
freqs_i = torch.cat([
|
||||
freqs_0.expand(seq_f, seq_h, seq_w, -1),
|
||||
freqs[1][h_sam].view(1, seq_h, 1, -1).expand(
|
||||
seq_f, seq_h, seq_w, -1),
|
||||
freqs[2][w_sam].view(1, 1, seq_w, -1).expand(
|
||||
seq_f, seq_h, seq_w, -1),
|
||||
],
|
||||
dim=-1).reshape(seq_len, 1, -1)
|
||||
elif t_f < 0:
|
||||
freqs_i = trainable_freqs.unsqueeze(1)
|
||||
# apply rotary embedding
|
||||
output[i, seq_bucket[-1]:seq_bucket[-1] + seq_len] = freqs_i
|
||||
seq_bucket.append(seq_bucket[-1] + seq_len)
|
||||
return output
|
||||
@@ -0,0 +1,673 @@
|
||||
# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
|
||||
import gc
|
||||
import logging
|
||||
import math
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
import types
|
||||
from contextlib import contextmanager
|
||||
from copy import deepcopy
|
||||
from functools import partial
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.cuda.amp as amp
|
||||
import torch.distributed as dist
|
||||
import torchvision.transforms.functional as TF
|
||||
from decord import VideoReader
|
||||
from PIL import Image
|
||||
from safetensors import safe_open
|
||||
from torchvision import transforms
|
||||
from tqdm import tqdm
|
||||
|
||||
from .distributed.fsdp import shard_model
|
||||
from .distributed.sequence_parallel import sp_attn_forward, sp_dit_forward
|
||||
from .distributed.util import get_world_size
|
||||
from .modules.s2v.audio_encoder import AudioEncoder
|
||||
from .modules.s2v.model_s2v import WanModel_S2V, sp_attn_forward_s2v
|
||||
from .modules.t5 import T5EncoderModel
|
||||
from .modules.vae2_1 import Wan2_1_VAE
|
||||
from .utils.fm_solvers import (
|
||||
FlowDPMSolverMultistepScheduler,
|
||||
get_sampling_sigmas,
|
||||
retrieve_timesteps,
|
||||
)
|
||||
from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
|
||||
|
||||
|
||||
def load_safetensors(path):
|
||||
tensors = {}
|
||||
with safe_open(path, framework="pt", device="cpu") as f:
|
||||
for key in f.keys():
|
||||
tensors[key] = f.get_tensor(key)
|
||||
return tensors
|
||||
|
||||
|
||||
class WanS2V:
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
config,
|
||||
checkpoint_dir,
|
||||
device_id=0,
|
||||
rank=0,
|
||||
t5_fsdp=False,
|
||||
dit_fsdp=False,
|
||||
use_sp=False,
|
||||
t5_cpu=False,
|
||||
init_on_cpu=True,
|
||||
convert_model_dtype=False,
|
||||
):
|
||||
r"""
|
||||
Initializes the image-to-video generation model components.
|
||||
|
||||
Args:
|
||||
config (EasyDict):
|
||||
Object containing model parameters initialized from config.py
|
||||
checkpoint_dir (`str`):
|
||||
Path to directory containing model checkpoints
|
||||
device_id (`int`, *optional*, defaults to 0):
|
||||
Id of target GPU device
|
||||
rank (`int`, *optional*, defaults to 0):
|
||||
Process rank for distributed training
|
||||
t5_fsdp (`bool`, *optional*, defaults to False):
|
||||
Enable FSDP sharding for T5 model
|
||||
dit_fsdp (`bool`, *optional*, defaults to False):
|
||||
Enable FSDP sharding for DiT model
|
||||
use_sp (`bool`, *optional*, defaults to False):
|
||||
Enable distribution strategy of sequence parallel.
|
||||
t5_cpu (`bool`, *optional*, defaults to False):
|
||||
Whether to place T5 model on CPU. Only works without t5_fsdp.
|
||||
init_on_cpu (`bool`, *optional*, defaults to True):
|
||||
Enable initializing Transformer Model on CPU. Only works without FSDP or USP.
|
||||
convert_model_dtype (`bool`, *optional*, defaults to False):
|
||||
Convert DiT model parameters dtype to 'config.param_dtype'.
|
||||
Only works without FSDP.
|
||||
"""
|
||||
self.device = torch.device(f"cuda:{device_id}")
|
||||
self.config = config
|
||||
self.rank = rank
|
||||
self.t5_cpu = t5_cpu
|
||||
self.init_on_cpu = init_on_cpu
|
||||
|
||||
self.num_train_timesteps = config.num_train_timesteps
|
||||
self.param_dtype = config.param_dtype
|
||||
|
||||
if t5_fsdp or dit_fsdp or use_sp:
|
||||
self.init_on_cpu = False
|
||||
|
||||
shard_fn = partial(shard_model, device_id=device_id)
|
||||
self.text_encoder = T5EncoderModel(
|
||||
text_len=config.text_len,
|
||||
dtype=config.t5_dtype,
|
||||
device=torch.device('cpu'),
|
||||
checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),
|
||||
tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),
|
||||
shard_fn=shard_fn if t5_fsdp else None,
|
||||
)
|
||||
|
||||
self.vae = Wan2_1_VAE(
|
||||
vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
|
||||
device=self.device)
|
||||
|
||||
logging.info(f"Creating WanModel from {checkpoint_dir}")
|
||||
if not dit_fsdp:
|
||||
self.noise_model = WanModel_S2V.from_pretrained(
|
||||
checkpoint_dir,
|
||||
torch_dtype=self.param_dtype,
|
||||
device_map=self.device)
|
||||
else:
|
||||
self.noise_model = WanModel_S2V.from_pretrained(
|
||||
checkpoint_dir, torch_dtype=self.param_dtype)
|
||||
|
||||
self.noise_model = self._configure_model(
|
||||
model=self.noise_model,
|
||||
use_sp=use_sp,
|
||||
dit_fsdp=dit_fsdp,
|
||||
shard_fn=shard_fn,
|
||||
convert_model_dtype=convert_model_dtype)
|
||||
|
||||
self.audio_encoder = AudioEncoder(
|
||||
model_id=os.path.join(checkpoint_dir,
|
||||
"wav2vec2-large-xlsr-53-english"))
|
||||
|
||||
if use_sp:
|
||||
self.sp_size = get_world_size()
|
||||
else:
|
||||
self.sp_size = 1
|
||||
|
||||
self.sample_neg_prompt = config.sample_neg_prompt
|
||||
self.motion_frames = config.transformer.motion_frames
|
||||
self.drop_first_motion = config.drop_first_motion
|
||||
self.fps = config.sample_fps
|
||||
self.audio_sample_m = 0
|
||||
|
||||
def _configure_model(self, model, use_sp, dit_fsdp, shard_fn,
|
||||
convert_model_dtype):
|
||||
"""
|
||||
Configures a model object. This includes setting evaluation modes,
|
||||
applying distributed parallel strategy, and handling device placement.
|
||||
|
||||
Args:
|
||||
model (torch.nn.Module):
|
||||
The model instance to configure.
|
||||
use_sp (`bool`):
|
||||
Enable distribution strategy of sequence parallel.
|
||||
dit_fsdp (`bool`):
|
||||
Enable FSDP sharding for DiT model.
|
||||
shard_fn (callable):
|
||||
The function to apply FSDP sharding.
|
||||
convert_model_dtype (`bool`):
|
||||
Convert DiT model parameters dtype to 'config.param_dtype'.
|
||||
Only works without FSDP.
|
||||
|
||||
Returns:
|
||||
torch.nn.Module:
|
||||
The configured model.
|
||||
"""
|
||||
model.eval().requires_grad_(False)
|
||||
if use_sp:
|
||||
for block in model.blocks:
|
||||
block.self_attn.forward = types.MethodType(
|
||||
sp_attn_forward_s2v, block.self_attn)
|
||||
model.use_context_parallel = True
|
||||
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
|
||||
if dit_fsdp:
|
||||
model = shard_fn(model)
|
||||
else:
|
||||
if convert_model_dtype:
|
||||
model.to(self.param_dtype)
|
||||
if not self.init_on_cpu:
|
||||
model.to(self.device)
|
||||
|
||||
return model
|
||||
|
||||
def get_size_less_than_area(self,
|
||||
height,
|
||||
width,
|
||||
target_area=1024 * 704,
|
||||
divisor=64):
|
||||
if height * width <= target_area:
|
||||
# If the original image area is already less than or equal to the target,
|
||||
# no resizing is needed—just padding. Still need to ensure that the padded area doesn't exceed the target.
|
||||
max_upper_area = target_area
|
||||
min_scale = 0.1
|
||||
max_scale = 1.0
|
||||
else:
|
||||
# Resize to fit within the target area and then pad to multiples of `divisor`
|
||||
max_upper_area = target_area # Maximum allowed total pixel count after padding
|
||||
d = divisor - 1
|
||||
b = d * (height + width)
|
||||
a = height * width
|
||||
c = d**2 - max_upper_area
|
||||
|
||||
# Calculate scale boundaries using quadratic equation
|
||||
min_scale = (-b + math.sqrt(b**2 - 2 * a * c)) / (
|
||||
2 * a) # Scale when maximum padding is applied
|
||||
max_scale = math.sqrt(max_upper_area /
|
||||
(height * width)) # Scale without any padding
|
||||
|
||||
# We want to choose the largest possible scale such that the final padded area does not exceed max_upper_area
|
||||
# Use binary search-like iteration to find this scale
|
||||
find_it = False
|
||||
for i in range(100):
|
||||
scale = max_scale - (max_scale - min_scale) * i / 100
|
||||
new_height, new_width = int(height * scale), int(width * scale)
|
||||
|
||||
# Pad to make dimensions divisible by 64
|
||||
pad_height = (64 - new_height % 64) % 64
|
||||
pad_width = (64 - new_width % 64) % 64
|
||||
pad_top = pad_height // 2
|
||||
pad_bottom = pad_height - pad_top
|
||||
pad_left = pad_width // 2
|
||||
pad_right = pad_width - pad_left
|
||||
|
||||
padded_height, padded_width = new_height + pad_height, new_width + pad_width
|
||||
|
||||
if padded_height * padded_width <= max_upper_area:
|
||||
find_it = True
|
||||
break
|
||||
|
||||
if find_it:
|
||||
return padded_height, padded_width
|
||||
else:
|
||||
# Fallback: calculate target dimensions based on aspect ratio and divisor alignment
|
||||
aspect_ratio = width / height
|
||||
target_width = int(
|
||||
(target_area * aspect_ratio)**0.5 // divisor * divisor)
|
||||
target_height = int(
|
||||
(target_area / aspect_ratio)**0.5 // divisor * divisor)
|
||||
|
||||
# Ensure the result is not larger than the original resolution
|
||||
if target_width >= width or target_height >= height:
|
||||
target_width = int(width // divisor * divisor)
|
||||
target_height = int(height // divisor * divisor)
|
||||
|
||||
return target_height, target_width
|
||||
|
||||
def prepare_default_cond_input(self,
|
||||
map_shape=[3, 12, 64, 64],
|
||||
motion_frames=5,
|
||||
lat_motion_frames=2,
|
||||
enable_mano=False,
|
||||
enable_kp=False,
|
||||
enable_pose=False):
|
||||
default_value = [1.0, -1.0, -1.0]
|
||||
cond_enable = [enable_mano, enable_kp, enable_pose]
|
||||
cond = []
|
||||
for d, c in zip(default_value, cond_enable):
|
||||
if c:
|
||||
map_value = torch.ones(
|
||||
map_shape, dtype=self.param_dtype, device=self.device) * d
|
||||
cond_lat = torch.cat([
|
||||
map_value[:, :, 0:1].repeat(1, 1, motion_frames, 1, 1),
|
||||
map_value
|
||||
],
|
||||
dim=2)
|
||||
cond_lat = torch.stack(
|
||||
self.vae.encode(cond_lat.to(
|
||||
self.param_dtype)))[:, :, lat_motion_frames:].to(
|
||||
self.param_dtype)
|
||||
|
||||
cond.append(cond_lat)
|
||||
if len(cond) >= 1:
|
||||
cond = torch.cat(cond, dim=1)
|
||||
else:
|
||||
cond = None
|
||||
return cond
|
||||
|
||||
def encode_audio(self, audio_path, infer_frames):
|
||||
z = self.audio_encoder.extract_audio_feat(
|
||||
audio_path, return_all_layers=True)
|
||||
audio_embed_bucket, num_repeat = self.audio_encoder.get_audio_embed_bucket_fps(
|
||||
z, fps=self.fps, batch_frames=infer_frames, m=self.audio_sample_m)
|
||||
audio_embed_bucket = audio_embed_bucket.to(self.device,
|
||||
self.param_dtype)
|
||||
audio_embed_bucket = audio_embed_bucket.unsqueeze(0)
|
||||
if len(audio_embed_bucket.shape) == 3:
|
||||
audio_embed_bucket = audio_embed_bucket.permute(0, 2, 1)
|
||||
elif len(audio_embed_bucket.shape) == 4:
|
||||
audio_embed_bucket = audio_embed_bucket.permute(0, 2, 3, 1)
|
||||
return audio_embed_bucket, num_repeat
|
||||
|
||||
def read_last_n_frames(self,
|
||||
video_path,
|
||||
n_frames,
|
||||
target_fps=16,
|
||||
reverse=False):
|
||||
"""
|
||||
Read the last `n_frames` from a video at the specified frame rate.
|
||||
|
||||
Parameters:
|
||||
video_path (str): Path to the video file.
|
||||
n_frames (int): Number of frames to read.
|
||||
target_fps (int, optional): Target sampling frame rate. Defaults to 16.
|
||||
reverse (bool, optional): Whether to read frames in reverse order.
|
||||
If True, reads the first `n_frames` instead of the last ones.
|
||||
|
||||
Returns:
|
||||
np.ndarray: A NumPy array of shape [n_frames, H, W, 3], representing the sampled video frames.
|
||||
"""
|
||||
vr = VideoReader(video_path)
|
||||
original_fps = vr.get_avg_fps()
|
||||
total_frames = len(vr)
|
||||
|
||||
interval = max(1, round(original_fps / target_fps))
|
||||
|
||||
required_span = (n_frames - 1) * interval
|
||||
|
||||
start_frame = max(0, total_frames - required_span -
|
||||
1) if not reverse else 0
|
||||
|
||||
sampled_indices = []
|
||||
for i in range(n_frames):
|
||||
indice = start_frame + i * interval
|
||||
if indice >= total_frames:
|
||||
break
|
||||
else:
|
||||
sampled_indices.append(indice)
|
||||
|
||||
return vr.get_batch(sampled_indices).asnumpy()
|
||||
|
||||
def load_pose_cond(self, pose_video, num_repeat, infer_frames, size):
|
||||
HEIGHT, WIDTH = size
|
||||
if not pose_video is None:
|
||||
pose_seq = self.read_last_n_frames(
|
||||
pose_video,
|
||||
n_frames=infer_frames * num_repeat,
|
||||
target_fps=self.fps,
|
||||
reverse=True)
|
||||
|
||||
resize_opreat = transforms.Resize(min(HEIGHT, WIDTH))
|
||||
crop_opreat = transforms.CenterCrop((HEIGHT, WIDTH))
|
||||
tensor_trans = transforms.ToTensor()
|
||||
|
||||
cond_tensor = torch.from_numpy(pose_seq)
|
||||
cond_tensor = cond_tensor.permute(0, 3, 1, 2) / 255.0 * 2 - 1.0
|
||||
cond_tensor = crop_opreat(resize_opreat(cond_tensor)).permute(
|
||||
1, 0, 2, 3).unsqueeze(0)
|
||||
|
||||
padding_frame_num = num_repeat * infer_frames - cond_tensor.shape[2]
|
||||
cond_tensor = torch.cat([
|
||||
cond_tensor,
|
||||
- torch.ones([1, 3, padding_frame_num, HEIGHT, WIDTH])
|
||||
],
|
||||
dim=2)
|
||||
|
||||
cond_tensors = torch.chunk(cond_tensor, num_repeat, dim=2)
|
||||
else:
|
||||
cond_tensors = [-torch.ones([1, 3, infer_frames, HEIGHT, WIDTH])]
|
||||
|
||||
COND = []
|
||||
for r in range(len(cond_tensors)):
|
||||
cond = cond_tensors[r]
|
||||
cond = torch.cat([cond[:, :, 0:1].repeat(1, 1, 1, 1, 1), cond],
|
||||
dim=2)
|
||||
cond_lat = torch.stack(
|
||||
self.vae.encode(
|
||||
cond.to(dtype=self.param_dtype,
|
||||
device=self.device)))[:, :,
|
||||
1:].cpu() # for mem save
|
||||
COND.append(cond_lat)
|
||||
return COND
|
||||
|
||||
def get_gen_size(self, size, max_area, ref_image_path, pre_video_path):
|
||||
if not size is None:
|
||||
HEIGHT, WIDTH = size
|
||||
else:
|
||||
if pre_video_path:
|
||||
ref_image = self.read_last_n_frames(
|
||||
pre_video_path, n_frames=1)[0]
|
||||
else:
|
||||
ref_image = np.array(Image.open(ref_image_path).convert('RGB'))
|
||||
HEIGHT, WIDTH = ref_image.shape[:2]
|
||||
HEIGHT, WIDTH = self.get_size_less_than_area(
|
||||
HEIGHT, WIDTH, target_area=max_area)
|
||||
return (HEIGHT, WIDTH)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
input_prompt,
|
||||
ref_image_path,
|
||||
audio_path,
|
||||
num_repeat=1,
|
||||
pose_video=None,
|
||||
max_area=720 * 1280,
|
||||
infer_frames=80,
|
||||
shift=5.0,
|
||||
sample_solver='unipc',
|
||||
sampling_steps=40,
|
||||
guide_scale=5.0,
|
||||
n_prompt="",
|
||||
seed=-1,
|
||||
offload_model=True,
|
||||
init_first_frame=False,
|
||||
):
|
||||
r"""
|
||||
Generates video frames from input image and text prompt using diffusion process.
|
||||
|
||||
Args:
|
||||
input_prompt (`str`):
|
||||
Text prompt for content generation.
|
||||
ref_image_path ('str'):
|
||||
Input image path
|
||||
audio_path ('str'):
|
||||
Audio for video driven
|
||||
num_repeat ('int'):
|
||||
Number of clips to generate; will be automatically adjusted based on the audio length
|
||||
pose_video ('str'):
|
||||
If provided, uses a sequence of poses to drive the generated video
|
||||
max_area (`int`, *optional*, defaults to 720*1280):
|
||||
Maximum pixel area for latent space calculation. Controls video resolution scaling
|
||||
infer_frames (`int`, *optional*, defaults to 80):
|
||||
How many frames to generate per clips. The number should be 4n
|
||||
shift (`float`, *optional*, defaults to 5.0):
|
||||
Noise schedule shift parameter. Affects temporal dynamics
|
||||
[NOTE]: If you want to generate a 480p video, it is recommended to set the shift value to 3.0.
|
||||
sample_solver (`str`, *optional*, defaults to 'unipc'):
|
||||
Solver used to sample the video.
|
||||
sampling_steps (`int`, *optional*, defaults to 40):
|
||||
Number of diffusion sampling steps. Higher values improve quality but slow generation
|
||||
guide_scale (`float` or tuple[`float`], *optional*, defaults 5.0):
|
||||
Classifier-free guidance scale. Controls prompt adherence vs. creativity.
|
||||
If tuple, the first guide_scale will be used for low noise model and
|
||||
the second guide_scale will be used for high noise model.
|
||||
n_prompt (`str`, *optional*, defaults to ""):
|
||||
Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt`
|
||||
seed (`int`, *optional*, defaults to -1):
|
||||
Random seed for noise generation. If -1, use random seed
|
||||
offload_model (`bool`, *optional*, defaults to True):
|
||||
If True, offloads models to CPU during generation to save VRAM
|
||||
init_first_frame (`bool`, *optional*, defaults to False):
|
||||
Whether to use the reference image as the first frame (i.e., standard image-to-video generation)
|
||||
|
||||
Returns:
|
||||
torch.Tensor:
|
||||
Generated video frames tensor. Dimensions: (C, N H, W) where:
|
||||
- C: Color channels (3 for RGB)
|
||||
- N: Number of frames (81)
|
||||
- H: Frame height (from max_area)
|
||||
- W: Frame width from max_area)
|
||||
"""
|
||||
# preprocess
|
||||
size = self.get_gen_size(
|
||||
size=None,
|
||||
max_area=max_area,
|
||||
ref_image_path=ref_image_path,
|
||||
pre_video_path=None)
|
||||
HEIGHT, WIDTH = size
|
||||
channel = 3
|
||||
|
||||
resize_opreat = transforms.Resize(min(HEIGHT, WIDTH))
|
||||
crop_opreat = transforms.CenterCrop((HEIGHT, WIDTH))
|
||||
tensor_trans = transforms.ToTensor()
|
||||
|
||||
ref_image = None
|
||||
motion_latents = None
|
||||
|
||||
if ref_image is None:
|
||||
ref_image = np.array(Image.open(ref_image_path).convert('RGB'))
|
||||
if motion_latents is None:
|
||||
motion_latents = torch.zeros(
|
||||
[1, channel, self.motion_frames, HEIGHT, WIDTH],
|
||||
dtype=self.param_dtype,
|
||||
device=self.device)
|
||||
|
||||
# extract audio emb
|
||||
audio_emb, nr = self.encode_audio(audio_path, infer_frames=infer_frames)
|
||||
if num_repeat is None or num_repeat > nr:
|
||||
num_repeat = nr
|
||||
|
||||
lat_motion_frames = (self.motion_frames + 3) // 4
|
||||
model_pic = crop_opreat(resize_opreat(Image.fromarray(ref_image)))
|
||||
|
||||
ref_pixel_values = tensor_trans(model_pic)
|
||||
ref_pixel_values = ref_pixel_values.unsqueeze(1).unsqueeze(
|
||||
0) * 2 - 1.0 # b c 1 h w
|
||||
ref_pixel_values = ref_pixel_values.to(
|
||||
dtype=self.vae.dtype, device=self.vae.device)
|
||||
ref_latents = torch.stack(self.vae.encode(ref_pixel_values))
|
||||
|
||||
# encode the motion latents
|
||||
videos_last_frames = motion_latents.detach()
|
||||
drop_first_motion = self.drop_first_motion
|
||||
if init_first_frame:
|
||||
drop_first_motion = False
|
||||
motion_latents[:, :, -6:] = ref_pixel_values
|
||||
motion_latents = torch.stack(self.vae.encode(motion_latents))
|
||||
|
||||
# get pose cond input if need
|
||||
COND = self.load_pose_cond(
|
||||
pose_video=pose_video,
|
||||
num_repeat=num_repeat,
|
||||
infer_frames=infer_frames,
|
||||
size=size)
|
||||
|
||||
seed = seed if seed >= 0 else random.randint(0, sys.maxsize)
|
||||
|
||||
if n_prompt == "":
|
||||
n_prompt = self.sample_neg_prompt
|
||||
|
||||
# preprocess
|
||||
if not self.t5_cpu:
|
||||
self.text_encoder.model.to(self.device)
|
||||
context = self.text_encoder([input_prompt], self.device)
|
||||
context_null = self.text_encoder([n_prompt], self.device)
|
||||
if offload_model:
|
||||
self.text_encoder.model.cpu()
|
||||
else:
|
||||
context = self.text_encoder([input_prompt], torch.device('cpu'))
|
||||
context_null = self.text_encoder([n_prompt], torch.device('cpu'))
|
||||
context = [t.to(self.device) for t in context]
|
||||
context_null = [t.to(self.device) for t in context_null]
|
||||
|
||||
out = []
|
||||
# evaluation mode
|
||||
with (
|
||||
torch.amp.autocast('cuda', dtype=self.param_dtype),
|
||||
torch.no_grad(),
|
||||
):
|
||||
for r in range(num_repeat):
|
||||
seed_g = torch.Generator(device=self.device)
|
||||
seed_g.manual_seed(seed + r)
|
||||
|
||||
lat_target_frames = (infer_frames + 3 + self.motion_frames
|
||||
) // 4 - lat_motion_frames
|
||||
target_shape = [lat_target_frames, HEIGHT // 8, WIDTH // 8]
|
||||
noise = [
|
||||
torch.randn(
|
||||
16,
|
||||
target_shape[0],
|
||||
target_shape[1],
|
||||
target_shape[2],
|
||||
dtype=self.param_dtype,
|
||||
device=self.device,
|
||||
generator=seed_g)
|
||||
]
|
||||
max_seq_len = np.prod(target_shape) // 4
|
||||
|
||||
if sample_solver == 'unipc':
|
||||
sample_scheduler = FlowUniPCMultistepScheduler(
|
||||
num_train_timesteps=self.num_train_timesteps,
|
||||
shift=1,
|
||||
use_dynamic_shifting=False)
|
||||
sample_scheduler.set_timesteps(
|
||||
sampling_steps, device=self.device, shift=shift)
|
||||
timesteps = sample_scheduler.timesteps
|
||||
elif sample_solver == 'dpm++':
|
||||
sample_scheduler = FlowDPMSolverMultistepScheduler(
|
||||
num_train_timesteps=self.num_train_timesteps,
|
||||
shift=1,
|
||||
use_dynamic_shifting=False)
|
||||
sampling_sigmas = get_sampling_sigmas(sampling_steps, shift)
|
||||
timesteps, _ = retrieve_timesteps(
|
||||
sample_scheduler,
|
||||
device=self.device,
|
||||
sigmas=sampling_sigmas)
|
||||
else:
|
||||
raise NotImplementedError("Unsupported solver.")
|
||||
|
||||
latents = deepcopy(noise)
|
||||
with torch.no_grad():
|
||||
left_idx = r * infer_frames
|
||||
right_idx = r * infer_frames + infer_frames
|
||||
cond_latents = COND[r] if pose_video else COND[0] * 0
|
||||
cond_latents = cond_latents.to(
|
||||
dtype=self.param_dtype, device=self.device)
|
||||
audio_input = audio_emb[..., left_idx:right_idx]
|
||||
input_motion_latents = motion_latents.clone()
|
||||
|
||||
arg_c = {
|
||||
'context': context[0:1],
|
||||
'seq_len': max_seq_len,
|
||||
'cond_states': cond_latents,
|
||||
"motion_latents": input_motion_latents,
|
||||
'ref_latents': ref_latents,
|
||||
"audio_input": audio_input,
|
||||
"motion_frames": [self.motion_frames, lat_motion_frames],
|
||||
"drop_motion_frames": drop_first_motion and r == 0,
|
||||
}
|
||||
if guide_scale > 1:
|
||||
arg_null = {
|
||||
'context': context_null[0:1],
|
||||
'seq_len': max_seq_len,
|
||||
'cond_states': cond_latents,
|
||||
"motion_latents": input_motion_latents,
|
||||
'ref_latents': ref_latents,
|
||||
"audio_input": 0.0 * audio_input,
|
||||
"motion_frames": [
|
||||
self.motion_frames, lat_motion_frames
|
||||
],
|
||||
"drop_motion_frames": drop_first_motion and r == 0,
|
||||
}
|
||||
if offload_model or self.init_on_cpu:
|
||||
self.noise_model.to(self.device)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
for i, t in enumerate(tqdm(timesteps)):
|
||||
latent_model_input = latents[0:1]
|
||||
timestep = [t]
|
||||
|
||||
timestep = torch.stack(timestep).to(self.device)
|
||||
|
||||
noise_pred_cond = self.noise_model(
|
||||
latent_model_input, t=timestep, **arg_c)
|
||||
|
||||
if guide_scale > 1:
|
||||
noise_pred_uncond = self.noise_model(
|
||||
latent_model_input, t=timestep, **arg_null)
|
||||
noise_pred = [
|
||||
u + guide_scale * (c - u)
|
||||
for c, u in zip(noise_pred_cond, noise_pred_uncond)
|
||||
]
|
||||
else:
|
||||
noise_pred = noise_pred_cond
|
||||
|
||||
temp_x0 = sample_scheduler.step(
|
||||
noise_pred[0].unsqueeze(0),
|
||||
t,
|
||||
latents[0].unsqueeze(0),
|
||||
return_dict=False,
|
||||
generator=seed_g)[0]
|
||||
latents[0] = temp_x0.squeeze(0)
|
||||
|
||||
if offload_model:
|
||||
self.noise_model.cpu()
|
||||
torch.cuda.synchronize()
|
||||
torch.cuda.empty_cache()
|
||||
latents = torch.stack(latents)
|
||||
if not (drop_first_motion and r == 0):
|
||||
decode_latents = torch.cat([motion_latents, latents], dim=2)
|
||||
else:
|
||||
decode_latents = torch.cat([ref_latents, latents], dim=2)
|
||||
image = torch.stack(self.vae.decode(decode_latents))
|
||||
image = image[:, :, -(infer_frames):]
|
||||
if (drop_first_motion and r == 0):
|
||||
image = image[:, :, 3:]
|
||||
|
||||
overlap_frames_num = min(self.motion_frames, image.shape[2])
|
||||
videos_last_frames = torch.cat([
|
||||
videos_last_frames[:, :, overlap_frames_num:],
|
||||
image[:, :, -overlap_frames_num:]
|
||||
],
|
||||
dim=2)
|
||||
videos_last_frames = videos_last_frames.to(
|
||||
dtype=motion_latents.dtype, device=motion_latents.device)
|
||||
motion_latents = torch.stack(
|
||||
self.vae.encode(videos_last_frames))
|
||||
out.append(image.cpu())
|
||||
|
||||
videos = torch.cat(out, dim=2)
|
||||
del noise, latents
|
||||
del sample_scheduler
|
||||
if offload_model:
|
||||
gc.collect()
|
||||
torch.cuda.synchronize()
|
||||
if dist.is_initialized():
|
||||
dist.barrier()
|
||||
|
||||
return videos[0] if self.rank == 0 else None
|
||||
@@ -4,6 +4,8 @@ import binascii
|
||||
import logging
|
||||
import os
|
||||
import os.path as osp
|
||||
import shutil
|
||||
import subprocess
|
||||
|
||||
import imageio
|
||||
import torch
|
||||
@@ -21,6 +23,70 @@ def rand_name(length=8, suffix=''):
|
||||
return name
|
||||
|
||||
|
||||
def merge_video_audio(video_path: str, audio_path: str):
|
||||
"""
|
||||
Merge the video and audio into a new video, with the duration set to the shorter of the two,
|
||||
and overwrite the original video file.
|
||||
|
||||
Parameters:
|
||||
video_path (str): Path to the original video file
|
||||
audio_path (str): Path to the audio file
|
||||
"""
|
||||
# set logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
# check
|
||||
if not os.path.exists(video_path):
|
||||
raise FileNotFoundError(f"video file {video_path} does not exist")
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"audio file {audio_path} does not exist")
|
||||
|
||||
base, ext = os.path.splitext(video_path)
|
||||
temp_output = f"{base}_temp{ext}"
|
||||
|
||||
try:
|
||||
# create ffmpeg command
|
||||
command = [
|
||||
'ffmpeg',
|
||||
'-y', # overwrite
|
||||
'-i',
|
||||
video_path,
|
||||
'-i',
|
||||
audio_path,
|
||||
'-c:v',
|
||||
'copy', # copy video stream
|
||||
'-c:a',
|
||||
'aac', # use AAC audio encoder
|
||||
'-b:a',
|
||||
'192k', # set audio bitrate (optional)
|
||||
'-map',
|
||||
'0:v:0', # select the first video stream
|
||||
'-map',
|
||||
'1:a:0', # select the first audio stream
|
||||
'-shortest', # choose the shortest duration
|
||||
temp_output
|
||||
]
|
||||
|
||||
# execute the command
|
||||
logging.info("Start merging video and audio...")
|
||||
result = subprocess.run(
|
||||
command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
|
||||
|
||||
# check result
|
||||
if result.returncode != 0:
|
||||
error_msg = f"FFmpeg execute failed: {result.stderr}"
|
||||
logging.error(error_msg)
|
||||
raise RuntimeError(error_msg)
|
||||
|
||||
shutil.move(temp_output, video_path)
|
||||
logging.info(f"Merge completed, saved to {video_path}")
|
||||
|
||||
except Exception as e:
|
||||
if os.path.exists(temp_output):
|
||||
os.remove(temp_output)
|
||||
logging.error(f"merge_video_audio failed with error: {e}")
|
||||
|
||||
|
||||
def save_video(tensor,
|
||||
save_file=None,
|
||||
fps=30,
|
||||
|
||||
Reference in New Issue
Block a user