diff --git a/.gitattributes b/.gitattributes
index a6344aac8c09253b3b630fb776ae94478aa0275b..527dd55c3f8a5edf0a333c05af2ca7c3a6edc708 100644
--- a/.gitattributes
+++ b/.gitattributes
@@ -33,3 +33,7 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
*.zip filter=lfs diff=lfs merge=lfs -text
*.zst filter=lfs diff=lfs merge=lfs -text
*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.wav filter=lfs diff=lfs merge=lfs -text
+*.mp4 filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
diff --git a/.gradio/certificate.pem b/.gradio/certificate.pem
new file mode 100644
index 0000000000000000000000000000000000000000..b85c8037f6b60976b2546fdbae88312c5246d9a3
--- /dev/null
+++ b/.gradio/certificate.pem
@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----
diff --git a/LICENSE.txt b/LICENSE.txt
new file mode 100644
index 0000000000000000000000000000000000000000..89282bc0d4adcd69235c4e3680bba83b52f7b364
--- /dev/null
+++ b/LICENSE.txt
@@ -0,0 +1,21 @@
+ MIT License
+
+ Copyright (c) Shuyuan Tu.
+
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ copies of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+
+ The above copyright notice and this permission notice shall be included in all
+ copies or substantial portions of the Software.
+
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+ SOFTWARE
\ No newline at end of file
diff --git a/accelerate_config/accelerate_config_machine_14B_multiple.yaml b/accelerate_config/accelerate_config_machine_14B_multiple.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..a7366cc4e41196d84d49aaff29ea23a053dbce23
--- /dev/null
+++ b/accelerate_config/accelerate_config_machine_14B_multiple.yaml
@@ -0,0 +1,19 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+deepspeed_config:
+ deepspeed_config_file: path/StableAvatar/deepspeed_config/zero_stage2_config.json
+ deepspeed_multinode_launcher: standard
+ zero3_init_flag: False
+distributed_type: DEEPSPEED
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+main_training_function: main
+dynamo_backend: 'no'
+num_machines: 8
+num_processes: 64
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
\ No newline at end of file
diff --git a/accelerate_config/accelerate_config_machine_1B_multiple.yaml b/accelerate_config/accelerate_config_machine_1B_multiple.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..b6f60269b003f9d991bc4c90e310d13def19ae53
--- /dev/null
+++ b/accelerate_config/accelerate_config_machine_1B_multiple.yaml
@@ -0,0 +1,15 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+enable_cpu_affinity: false
+main_training_function: main
+dynamo_backend: 'no'
+num_machines: 8
+num_processes: 64
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
\ No newline at end of file
diff --git a/app.py b/app.py
new file mode 100644
index 0000000000000000000000000000000000000000..fb74205363f40d0a010bc49f08b16cfeddfc489d
--- /dev/null
+++ b/app.py
@@ -0,0 +1,718 @@
+import torch
+import psutil
+import argparse
+import gradio as gr
+import os
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import load_image
+from transformers import AutoTokenizer, Wav2Vec2Model, Wav2Vec2Processor
+from omegaconf import OmegaConf
+from wan.models.cache_utils import get_teacache_coefficients
+from wan.models.wan_fantasy_transformer3d_1B import WanTransformer3DFantasyModel
+from wan.models.wan_text_encoder import WanT5EncoderModel
+from wan.models.wan_vae import AutoencoderKLWan
+from wan.models.wan_image_encoder import CLIPModel
+from wan.pipeline.wan_inference_long_pipeline import WanI2VTalkingInferenceLongPipeline
+from wan.utils.fp8_optimization import replace_parameters_by_name, convert_weight_dtype_wrapper, convert_model_weight_to_float8
+from wan.utils.utils import get_image_to_video_latent, save_videos_grid
+import numpy as np
+import librosa
+import datetime
+import random
+import math
+import subprocess
+from moviepy.editor import VideoFileClip
+from huggingface_hub import snapshot_download
+import shutil
+try:
+ from audio_separator.separator import Separator
+except:
+ print("Unable to use vocal separation feature. Please install audio-separator[gpu].")
+
+
+if torch.cuda.is_available():
+ device = "cuda"
+ if torch.cuda.get_device_capability()[0] >= 8:
+ dtype = torch.bfloat16
+ else:
+ dtype = torch.float16
+else:
+ device = "cpu"
+ dtype = torch.float32
+
+
+def filter_kwargs(cls, kwargs):
+ import inspect
+ sig = inspect.signature(cls.__init__)
+ valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
+ filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
+ return filtered_kwargs
+
+
+def load_transformer_model(model_version):
+ """
+ 根据选择的模型版本加载对应的transformer模型
+
+ Args:
+ model_version (str): 模型版本，"square" 或 "rec_vec"
+
+ Returns:
+ WanTransformer3DFantasyModel: 加载的transformer模型
+ """
+ global transformer3d
+
+ if model_version == "square":
+ transformer_path = os.path.join(repo_root, "StableAvatar-1.3B", "transformer3d-square.pt")
+ elif model_version == "rec_vec":
+ transformer_path = os.path.join(repo_root, "StableAvatar-1.3B", "transformer3d-rec-vec.pt")
+ else:
+ # 默认使用square版本
+ transformer_path = os.path.join(repo_root, "StableAvatar-1.3B", "transformer3d-square.pt")
+
+ print(f"正在加载模型: {transformer_path}")
+
+ if os.path.exists(transformer_path):
+ state_dict = torch.load(transformer_path, map_location="cpu")
+ state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
+ m, u = transformer3d.load_state_dict(state_dict, strict=False)
+ print(f"模型加载成功: {transformer_path}")
+ print(f"Missing keys: {len(m)}; Unexpected keys: {len(u)}")
+ return transformer3d
+ else:
+ print(f"错误：模型文件不存在: {transformer_path}")
+ return None
+
+
+REPO_ID = "FrancisRing/StableAvatar"
+repo_root = snapshot_download(
+ repo_id=REPO_ID,
+ allow_patterns=[
+ "StableAvatar-1.3B/*",
+ "Wan2.1-Fun-V1.1-1.3B-InP/*",
+ "wav2vec2-base-960h/*",
+ "assets/**",
+ "Kim_Vocal_2.onnx",
+ ],
+)
+pretrained_model_name_or_path = os.path.join(repo_root, "Wan2.1-Fun-V1.1-1.3B-InP")
+pretrained_wav2vec_path = os.path.join(repo_root, "wav2vec2-base-960h")
+
+
+# 人声分离 onnx
+audio_separator_model_file = os.path.join(repo_root, "Kim_Vocal_2.onnx")
+
+# model_path = "/datadrive/stableavatar/checkpoints"
+# pretrained_model_name_or_path = f"{model_path}/Wan2.1-Fun-V1.1-1.3B-InP"
+# pretrained_wav2vec_path = f"{model_path}/wav2vec2-base-960h"
+# transformer_path = f"{model_path}/StableAvatar-1.3B/transformer3d-square.pt"
+config = OmegaConf.load("deepspeed_config/wan2.1/wan_civitai.yaml")
+sampler_name = "Flow"
+clip_sample_n_frames = 81
+tokenizer = AutoTokenizer.from_pretrained(os.path.join(pretrained_model_name_or_path, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')), )
+text_encoder = WanT5EncoderModel.from_pretrained(
+ os.path.join(pretrained_model_name_or_path, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
+ additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
+ low_cpu_mem_usage=True,
+ torch_dtype=dtype,
+)
+text_encoder = text_encoder.eval()
+vae = AutoencoderKLWan.from_pretrained(
+ os.path.join(pretrained_model_name_or_path, config['vae_kwargs'].get('vae_subpath', 'vae')),
+ additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
+)
+wav2vec_processor = Wav2Vec2Processor.from_pretrained(pretrained_wav2vec_path)
+wav2vec = Wav2Vec2Model.from_pretrained(pretrained_wav2vec_path).to("cpu")
+clip_image_encoder = CLIPModel.from_pretrained(os.path.join(pretrained_model_name_or_path, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')), )
+clip_image_encoder = clip_image_encoder.eval()
+transformer3d = WanTransformer3DFantasyModel.from_pretrained(
+ os.path.join(pretrained_model_name_or_path, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
+ transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
+ low_cpu_mem_usage=False,
+ torch_dtype=dtype,
+)
+
+# 默认加载square版本模型
+load_transformer_model("square")
+Choosen_Scheduler = scheduler_dict = {
+ "Flow": FlowMatchEulerDiscreteScheduler,
+}[sampler_name]
+scheduler = Choosen_Scheduler(
+ **filter_kwargs(Choosen_Scheduler, OmegaConf.to_container(config['scheduler_kwargs']))
+)
+pipeline = WanI2VTalkingInferenceLongPipeline(
+ tokenizer=tokenizer,
+ text_encoder=text_encoder,
+ vae=vae,
+ transformer=transformer3d,
+ clip_image_encoder=clip_image_encoder,
+ scheduler=scheduler,
+ wav2vec_processor=wav2vec_processor,
+ wav2vec=wav2vec,
+)
+
+
+def generate(
+ GPU_memory_mode,
+ teacache_threshold,
+ num_skip_start_steps,
+ image_path,
+ audio_path,
+ prompt,
+ negative_prompt,
+ width,
+ height,
+ guidance_scale,
+ num_inference_steps,
+ text_guide_scale,
+ audio_guide_scale,
+ motion_frame,
+ fps,
+ overlap_window_length,
+ seed_param,
+ overlapping_weight_scheme,
+ progress=gr.Progress(track_tqdm=True),
+):
+ global pipeline, transformer3d
+ timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+ if seed_param<0:
+ seed = random.randint(0, np.iinfo(np.int32).max)
+ else:
+ seed = seed_param
+
+ if GPU_memory_mode == "sequential_cpu_offload":
+ replace_parameters_by_name(transformer3d, ["modulation", ], device=device)
+ transformer3d.freqs = transformer3d.freqs.to(device=device)
+ pipeline.enable_sequential_cpu_offload(device=device)
+ elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+ convert_model_weight_to_float8(transformer3d, exclude_module_name=["modulation", ])
+ convert_weight_dtype_wrapper(transformer3d, dtype)
+ pipeline.enable_model_cpu_offload(device=device)
+ elif GPU_memory_mode == "model_cpu_offload":
+ pipeline.enable_model_cpu_offload(device=device)
+ else:
+ pipeline.to(device=device)
+
+ if teacache_threshold > 0:
+ coefficients = get_teacache_coefficients(pretrained_model_name_or_path)
+ pipeline.transformer.enable_teacache(
+ coefficients,
+ num_inference_steps,
+ teacache_threshold,
+ num_skip_start_steps=num_skip_start_steps,
+ )
+
+ with torch.no_grad():
+ video_length = int((clip_sample_n_frames - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if clip_sample_n_frames != 1 else 1
+ input_video, input_video_mask, clip_image = get_image_to_video_latent(image_path, None, video_length=video_length, sample_size=[height, width])
+ sr = 16000
+ vocal_input, sample_rate = librosa.load(audio_path, sr=sr)
+ sample = pipeline(
+ prompt,
+ num_frames=video_length,
+ negative_prompt=negative_prompt,
+ width=width,
+ height=height,
+ guidance_scale=guidance_scale,
+ generator=torch.Generator().manual_seed(seed),
+ num_inference_steps=num_inference_steps,
+ video=input_video,
+ mask_video=input_video_mask,
+ clip_image=clip_image,
+ text_guide_scale=text_guide_scale,
+ audio_guide_scale=audio_guide_scale,
+ vocal_input_values=vocal_input,
+ motion_frame=motion_frame,
+ fps=fps,
+ sr=sr,
+ cond_file_path=image_path,
+ overlap_window_length=overlap_window_length,
+ seed=seed,
+ overlapping_weight_scheme=overlapping_weight_scheme,
+ ).videos
+ os.makedirs("outputs", exist_ok=True)
+ video_path = os.path.join("outputs", f"{timestamp}.mp4")
+ save_videos_grid(sample, video_path, fps=fps)
+ output_video_with_audio = os.path.join("outputs", f"{timestamp}_audio.mp4")
+ subprocess.run([
+ "ffmpeg", "-y", "-loglevel", "quiet", "-i", video_path, "-i", audio_path,
+ "-c:v", "copy", "-c:a", "aac", "-strict", "experimental",
+ output_video_with_audio
+ ], check=True)
+
+ return output_video_with_audio, seed, f"Generated outputs/{timestamp}.mp4 / 已生成outputs/{timestamp}.mp4"
+
+
+def exchange_width_height(width, height):
+ return height, width, "✅ Width and Height Swapped / 宽高交换完毕"
+
+
+def adjust_width_height(image):
+ image = load_image(image)
+ width, height = image.size
+ original_area = width * height
+ default_area = 512*512
+ ratio = math.sqrt(original_area / default_area)
+ width = width / ratio // 16 * 16
+ height = height / ratio // 16 * 16
+ return int(width), int(height), "✅ Adjusted Size Based on Image / 根据图片调整宽高"
+
+
+def audio_extractor(video_path):
+ timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+ os.makedirs("outputs", exist_ok=True) # 确保目录存在
+ out_wav = os.path.abspath(os.path.join("outputs", f"{timestamp}.wav"))
+ video = VideoFileClip(video_path)
+ audio = video.audio
+ audio.write_audiofile(out_wav, codec="pcm_s16le")
+ return out_wav, f"Generated {out_wav} / 已生成 {out_wav}", out_wav # ← 第3个返回给 gr.File
+
+def vocal_separation(audio_path):
+ timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+ os.makedirs("outputs", exist_ok=True)
+ # audio_separator_model_file = "checkpoints/Kim_Vocal_2.onnx"
+ audio_separator = Separator(
+ output_dir=os.path.abspath(os.path.join("outputs", timestamp)),
+ output_single_stem="vocals",
+ model_file_dir=os.path.dirname(audio_separator_model_file),
+ )
+ audio_separator.load_model(os.path.basename(audio_separator_model_file))
+ assert audio_separator.model_instance is not None, "Fail to load audio separate model."
+ outputs = audio_separator.separate(audio_path)
+ vocal_audio_file = os.path.join(audio_separator.output_dir, outputs[0])
+ destination_file = os.path.abspath(os.path.join("outputs", f"{timestamp}.wav"))
+ shutil.copy(vocal_audio_file, destination_file)
+ os.remove(vocal_audio_file)
+ return destination_file, f"Generated {destination_file} / 已生成 {destination_file}", destination_file
+
+
+def update_language(language):
+ if language == "English":
+ return {
+ GPU_memory_mode: gr.Dropdown(label="GPU Memory Mode", info="Normal uses 25G VRAM, model_cpu_offload uses 13G VRAM"),
+ teacache_threshold: gr.Slider(label="TeaCache Threshold", info="Recommended 0.1, 0 disables TeaCache acceleration"),
+ num_skip_start_steps: gr.Slider(label="Skip Start Steps", info="Recommended 5"),
+ model_version: gr.Dropdown(label="Model Version", choices=["square", "rec_vec"], value="square"),
+ image_path: gr.Image(label="Upload Image"),
+ audio_path: gr.Audio(label="Upload Audio"),
+ prompt: gr.Textbox(label="Prompt"),
+ negative_prompt: gr.Textbox(label="Negative Prompt"),
+ generate_button: gr.Button("🎬 Start Generation"),
+ width: gr.Slider(label="Width"),
+ height: gr.Slider(label="Height"),
+ exchange_button: gr.Button("🔄 Swap Width/Height"),
+ adjust_button: gr.Button("Adjust Size Based on Image"),
+ guidance_scale: gr.Slider(label="Guidance Scale"),
+ num_inference_steps: gr.Slider(label="Sampling Steps (Recommended 50)"),
+ text_guide_scale: gr.Slider(label="Text Guidance Scale"),
+ audio_guide_scale: gr.Slider(label="Audio Guidance Scale"),
+ motion_frame: gr.Slider(label="Motion Frame"),
+ fps: gr.Slider(label="FPS"),
+ overlap_window_length: gr.Slider(label="Overlap Window Length"),
+ seed_param: gr.Number(label="Seed (positive integer, -1 for random)"),
+ overlapping_weight_scheme: gr.Dropdown(label="Overlapping Weight Scheme", choices=["uniform", "log"], value="uniform"),
+ info: gr.Textbox(label="Status"),
+ video_output: gr.Video(label="Generated Result"),
+ seed_output: gr.Textbox(label="Seed"),
+ video_path: gr.Video(label="Upload Video"),
+ extractor_button: gr.Button("🎬 Start Extraction"),
+ info2: gr.Textbox(label="Status"),
+ audio_output: gr.Audio(label="Generated Result"),
+ audio_path3: gr.Audio(label="Upload Audio"),
+ separation_button: gr.Button("🎬 Start Separation"),
+ info3: gr.Textbox(label="Status"),
+ audio_output3: gr.Audio(label="Generated Result"),
+ example_title: gr.Markdown(value="### Select the following example cases for testing:"),
+ example1_label: gr.Markdown(value="**Example 1**"),
+ example2_label: gr.Markdown(value="**Example 2**"),
+ example3_label: gr.Markdown(value="**Example 3**"),
+ example4_label: gr.Markdown(value="**Example 4**"),
+ example5_label: gr.Markdown(value="**Example 5**"),
+ example1_btn: gr.Button("🚀 Use Example 1", variant="secondary"),
+ example2_btn: gr.Button("🚀 Use Example 2", variant="secondary"),
+ example3_btn: gr.Button("🚀 Use Example 3", variant="secondary"),
+ example4_btn: gr.Button("🚀 Use Example 4", variant="secondary"),
+ example5_btn: gr.Button("🚀 Use Example 5", variant="secondary"),
+ parameter_settings_title: gr.Accordion(label="Parameter Settings", open=True),
+ example_cases_title: gr.Accordion(label="Example Cases", open=True),
+ stableavatar_title: gr.TabItem(label="StableAvatar"),
+ audio_extraction_title: gr.TabItem(label="Audio Extraction"),
+ vocal_separation_title: gr.TabItem(label="Vocal Separation")
+ }
+ else:
+ return {
+ GPU_memory_mode: gr.Dropdown(label="显存模式", info="Normal占用25G显存，model_cpu_offload占用13G显存"),
+ teacache_threshold: gr.Slider(label="teacache threshold", info="推荐参数0.1，0为禁用teacache加速"),
+ num_skip_start_steps: gr.Slider(label="跳过开始步数", info="推荐参数5"),
+ model_version: gr.Dropdown(label="模型版本", choices=["square", "rec_vec"], value="square"),
+ image_path: gr.Image(label="上传图片"),
+ audio_path: gr.Audio(label="上传音频"),
+ prompt: gr.Textbox(label="提示词"),
+ negative_prompt: gr.Textbox(label="负面提示词"),
+ generate_button: gr.Button("🎬 开始生成"),
+ width: gr.Slider(label="宽度"),
+ height: gr.Slider(label="高度"),
+ exchange_button: gr.Button("🔄 交换宽高"),
+ adjust_button: gr.Button("根据图片调整宽高"),
+ guidance_scale: gr.Slider(label="guidance scale"),
+ num_inference_steps: gr.Slider(label="采样步数（推荐50步）", minimum=1, maximum=100, step=1, value=50),
+ text_guide_scale: gr.Slider(label="text guidance scale"),
+ audio_guide_scale: gr.Slider(label="audio guidance scale"),
+ motion_frame: gr.Slider(label="motion frame"),
+ fps: gr.Slider(label="帧率"),
+ overlap_window_length: gr.Slider(label="overlap window length"),
+ seed_param: gr.Number(label="种子，请输入正整数，-1为随机"),
+ overlapping_weight_scheme: gr.Dropdown(label="Overlapping Weight Scheme", choices=["uniform", "log"], value="uniform"),
+ info: gr.Textbox(label="提示信息"),
+ video_output: gr.Video(label="生成结果"),
+ seed_output: gr.Textbox(label="种子"),
+ video_path: gr.Video(label="上传视频"),
+ extractor_button: gr.Button("🎬 开始提取"),
+ info2: gr.Textbox(label="提示信息"),
+ audio_output: gr.Audio(label="生成结果"),
+ audio_path3: gr.Audio(label="上传音频"),
+ separation_button: gr.Button("🎬 开始分离"),
+ info3: gr.Textbox(label="提示信息"),
+ audio_output3: gr.Audio(label="生成结果"),
+ example_title: gr.Markdown(value="### 选择以下示例案例进行测试："),
+ example1_label: gr.Markdown(value="**示例 1**"),
+ example2_label: gr.Markdown(value="**示例 2**"),
+ example3_label: gr.Markdown(value="**示例 3**"),
+ example4_label: gr.Markdown(value="**示例 4**"),
+ example5_label: gr.Markdown(value="**示例 5**"),
+ example1_btn: gr.Button("🚀 使用示例 1", variant="secondary"),
+ example2_btn: gr.Button("🚀 使用示例 2", variant="secondary"),
+ example3_btn: gr.Button("🚀 使用示例 3", variant="secondary"),
+ example4_btn: gr.Button("🚀 使用示例 4", variant="secondary"),
+ example5_btn: gr.Button("🚀 使用示例 5", variant="secondary"),
+ parameter_settings_title: gr.Accordion(label="参数设置", open=True),
+ example_cases_title: gr.Accordion(label="示例案例", open=True),
+ stableavatar_title: gr.TabItem(label="StableAvatar"),
+ audio_extraction_title: gr.TabItem(label="音频提取"),
+ vocal_separation_title: gr.TabItem(label="人声分离")
+ }
+
+BANNER_HTML = """
+
+
+
+
+ STABLEAVATAR
+
+
+
+"""
+
+BANNER_CSS = """
+.hero{display:flex;align-items:center;gap:18px;padding:18px;border-radius:14px;
+ background:#111;color:#fff;margin-bottom:12px}
+.brand-text{font-weight:800;letter-spacing:2px}
+.brand img{height:46px}
+.titles h1{font-size:28px;margin:0 0 6px 0}
+.badges{display:flex;gap:10px;flex-wrap:wrap}
+.badge img{height:22px}
+.divider{border:0;border-top:1px solid rgba(255,255,255,0.18);margin:6px 0 18px}
+"""
+
+
+# with gr.Blocks(theme=gr.themes.Base()) as demo:
+# gr.Markdown("""
+#
+#
StableAvatar
+#
+# """)
+with gr.Blocks(theme=gr.themes.Base(), css=BANNER_CSS) as demo:
+ gr.HTML(BANNER_HTML)
+
+ language_radio = gr.Radio(
+ choices=["English", "中文"],
+ value="English",
+ label="Language / 语言"
+ )
+
+ with gr.Accordion("Model Settings / 模型设置", open=False):
+ with gr.Row():
+ GPU_memory_mode = gr.Dropdown(
+ label = "显存模式",
+ info = "Normal占用25G显存，model_cpu_offload占用13G显存",
+ choices = ["Normal", "model_cpu_offload", "model_cpu_offloadand_qfloat8", "sequential_cpu_offload"],
+ value = "model_cpu_offload"
+ )
+ teacache_threshold = gr.Slider(label="teacache threshold", info = "推荐参数0.1，0为禁用teacache加速", minimum=0, maximum=1, step=0.01, value=0)
+ num_skip_start_steps = gr.Slider(label="跳过开始步数", info = "推荐参数5", minimum=0, maximum=100, step=1, value=5)
+ with gr.Row():
+ model_version = gr.Dropdown(
+ label = "模型版本",
+ choices = ["square","rec_vec"],
+ value = "square"
+ )
+
+ stableavatar_title = gr.TabItem(label="StableAvatar")
+ with stableavatar_title:
+ with gr.Row():
+ with gr.Column():
+ with gr.Row():
+ image_path = gr.Image(label="上传图片", type="filepath", height=280)
+ audio_path = gr.Audio(label="上传音频", type="filepath")
+ prompt = gr.Textbox(label="提示词", value="")
+ negative_prompt = gr.Textbox(label="负面提示词", value="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走")
+ generate_button = gr.Button("🎬 开始生成", variant='primary')
+ parameter_settings_title = gr.Accordion(label="参数设置", open=True)
+ with parameter_settings_title:
+ with gr.Row():
+ width = gr.Slider(label="宽度", minimum=256, maximum=2048, step=16, value=512)
+ height = gr.Slider(label="高度", minimum=256, maximum=2048, step=16, value=512)
+ with gr.Row():
+ exchange_button = gr.Button("🔄 交换宽高")
+ adjust_button = gr.Button("根据图片调整宽高")
+ with gr.Row():
+ guidance_scale = gr.Slider(label="guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=6.0)
+ num_inference_steps = gr.Slider(label="采样步数（推荐50步）", minimum=1, maximum=100, step=1, value=50)
+ with gr.Row():
+ text_guide_scale = gr.Slider(label="text guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=3.0)
+ audio_guide_scale = gr.Slider(label="audio guidance scale", minimum=1.0, maximum=10.0, step=0.1, value=5.0)
+ with gr.Row():
+ motion_frame = gr.Slider(label="motion frame", minimum=1, maximum=50, step=1, value=25)
+ fps = gr.Slider(label="帧率", minimum=1, maximum=60, step=1, value=25)
+ with gr.Row():
+ overlap_window_length = gr.Slider(label="overlap window length", minimum=1, maximum=20, step=1, value=10)
+ seed_param = gr.Number(label="种子，请输入正整数，-1为随机", value=42)
+ with gr.Row():
+ overlapping_weight_scheme = gr.Dropdown(label="Overlapping Weight Scheme", choices=["uniform", "log"], value="uniform")
+ with gr.Column():
+ info = gr.Textbox(label="提示信息", interactive=False)
+ video_output = gr.Video(label="生成结果", interactive=False)
+ seed_output = gr.Textbox(label="种子")
+
+ # 示例案例部分移到StableAvatar标签页内部
+ example_cases_title = gr.Accordion(label="示例案例", open=True)
+ with example_cases_title:
+ example_title = gr.Markdown(value="### 选择以下示例案例进行测试：")
+ with gr.Row():
+ with gr.Column():
+ example1_label = gr.Markdown(value="**示例 1**")
+ example1_image = gr.Image(value="example_case/case-1/reference.png", label="", interactive=False, height=120, show_label=False)
+ example1_audio = gr.Audio(value="example_case/case-1/audio.wav", label="", interactive=False, show_label=False)
+ example1_btn = gr.Button("🚀 使用示例 1", variant="secondary", size="sm")
+
+ with gr.Column():
+ example2_label = gr.Markdown(value="**示例 2**")
+ example2_image = gr.Image(value="example_case/case-2/reference.png", label="", interactive=False, height=120, show_label=False)
+ example2_audio = gr.Audio(value="example_case/case-2/audio.wav", label="", interactive=False, show_label=False)
+ example2_btn = gr.Button("🚀 使用示例 2", variant="secondary", size="sm")
+
+ with gr.Column():
+ example3_label = gr.Markdown(value="**示例 3**")
+ example3_image = gr.Image(value="example_case/case-6/reference.png", label="", interactive=False, height=120, show_label=False)
+ example3_audio = gr.Audio(value="example_case/case-6/audio.wav", label="", interactive=False, show_label=False)
+ example3_btn = gr.Button("🚀 使用示例 3", variant="secondary", size="sm")
+
+ with gr.Column():
+ example4_label = gr.Markdown(value="**示例 4**")
+ example4_image = gr.Image(value="example_case/case-45/reference.png", label="", interactive=False, height=120, show_label=False)
+ example4_audio = gr.Audio(value="example_case/case-45/audio.wav", label="", interactive=False, show_label=False)
+ example4_btn = gr.Button("🚀 使用示例 4", variant="secondary", size="sm")
+
+ with gr.Column():
+ example5_label = gr.Markdown(value="**示例 5**")
+ example5_image = gr.Image(value="example_case/case-3/reference.jpg", label="", interactive=False, height=120, show_label=False)
+ example5_audio = gr.Audio(value="example_case/case-3/audio.wav", label="", interactive=False, show_label=False)
+ example5_btn = gr.Button("🚀 使用示例 5", variant="secondary", size="sm")
+
+ audio_extraction_title = gr.TabItem(label="音频提取")
+ with audio_extraction_title:
+ with gr.Row():
+ with gr.Column():
+ video_path = gr.Video(label="上传视频", height=500)
+ extractor_button = gr.Button("🎬 开始提取", variant='primary')
+ with gr.Column():
+ info2 = gr.Textbox(label="提示信息", interactive=False)
+ audio_output = gr.Audio(label="生成结果", interactive=False)
+ audio_file = gr.File(label="download audio file")
+
+ vocal_separation_title = gr.TabItem(label="人声分离")
+ with vocal_separation_title:
+ with gr.Row():
+ with gr.Column():
+ audio_path3 = gr.Audio(label="上传音频", type="filepath")
+ separation_button = gr.Button("🎬 开始分离", variant='primary')
+ with gr.Column():
+ info3 = gr.Textbox(label="提示信息", interactive=False)
+ audio_output3 = gr.Audio(label="生成结果", interactive=False)
+ audio_file3 = gr.File(label="download audio file")
+
+ # 示例案例部分移到末尾
+ # example_cases_title = gr.Accordion(label="示例案例", open=True)
+ # with example_cases_title:
+ # example_title = gr.Markdown(value="### 选择以下示例案例进行测试：")
+ # with gr.Row():
+ # with gr.Column():
+ # example1_label = gr.Markdown(value="**示例 1**")
+ # example1_image = gr.Image(value="example_case/case-1/reference.png", label="", interactive=False, height=120, show_label=False)
+ # example1_audio = gr.Audio(value="example_case/case-1/audio.wav", label="", interactive=False, show_label=False)
+ # example1_btn = gr.Button("🚀 使用示例 1", variant="secondary", size="sm")
+
+ # with gr.Column():
+ # example2_label = gr.Markdown(value="**示例 2**")
+ # example2_image = gr.Image(value="example_case/case-2/reference.png", label="", interactive=False, height=120, show_label=False)
+ # example2_audio = gr.Audio(value="example_case/case-2/audio.wav", label="", interactive=False, show_label=False)
+ # example2_btn = gr.Button("🚀 使用示例 2", variant="secondary", size="sm")
+
+ # with gr.Column():
+ # example3_label = gr.Markdown(value="**示例 3**")
+ # example3_image = gr.Image(value="example_case/case-6/reference.png", label="", interactive=False, height=120, show_label=False)
+ # example3_audio = gr.Audio(value="example_case/case-6/audio.wav", label="", interactive=False, show_label=False)
+ # example3_btn = gr.Button("🚀 使用示例 3", variant="secondary", size="sm")
+
+ # with gr.Column():
+ # example4_label = gr.Markdown(value="**示例 4**")
+ # example4_image = gr.Image(value="example_case/case-45/reference.png", label="", interactive=False, height=120, show_label=False)
+ # example4_audio = gr.Audio(value="example_case/case-45/audio.wav", label="", interactive=False, show_label=False)
+ # example4_btn = gr.Button("🚀 使用示例 4", variant="secondary", size="sm")
+
+ # with gr.Column():
+ # example5_label = gr.Markdown(value="**示例 5**")
+ # example5_image = gr.Image(value="example_case/case-3/reference.jpg", label="", interactive=False, height=120, show_label=False)
+ # example5_audio = gr.Audio(value="example_case/case-3/audio.wav", label="", interactive=False, show_label=False)
+ # example5_btn = gr.Button("🚀 使用示例 5", variant="secondary", size="sm")
+
+ all_components = [GPU_memory_mode, teacache_threshold, num_skip_start_steps, model_version, image_path, audio_path, prompt, negative_prompt, generate_button, width, height, exchange_button, adjust_button, guidance_scale, num_inference_steps, text_guide_scale, audio_guide_scale, motion_frame, fps, overlap_window_length, seed_param, overlapping_weight_scheme, info, video_output, seed_output, video_path, extractor_button, info2, audio_output, audio_path3, separation_button, info3, audio_output3, example_title, example1_label, example2_label, example3_label, example4_label, example1_btn, example2_btn, example3_btn, example4_btn, example5_label, example5_btn, parameter_settings_title, example_cases_title, stableavatar_title, audio_extraction_title, vocal_separation_title]
+
+ language_radio.change(
+ fn=update_language,
+ inputs=[language_radio],
+ outputs=all_components
+ )
+
+ # 添加模型版本选择的事件处理
+ def on_model_version_change(model_version):
+ """当模型版本改变时，重新加载对应的模型"""
+ result = load_transformer_model(model_version)
+ if result is not None:
+ return f"✅ 模型已切换到 {model_version} 版本"
+ else:
+ return f"❌ 模型切换失败，请检查文件是否存在"
+
+ model_version.change(
+ fn=on_model_version_change,
+ inputs=[model_version],
+ outputs=[info]
+ )
+
+ demo.load(fn=update_language, inputs=[language_radio], outputs=all_components)
+ # 添加示例案例按钮的事件处理
+ def load_example1():
+ try:
+ with open("example_case/case-1/prompt.txt", "r", encoding="utf-8") as f:
+ prompt_text = f.read().strip()
+ except:
+ prompt_text = ""
+ return "example_case/case-1/reference.png", "example_case/case-1/audio.wav", prompt_text
+
+ def load_example2():
+ try:
+ with open("example_case/case-2/prompt.txt", "r", encoding="utf-8") as f:
+ prompt_text = f.read().strip()
+ except:
+ prompt_text = ""
+ return "example_case/case-2/reference.png", "example_case/case-2/audio.wav", prompt_text
+
+ def load_example3():
+ try:
+ with open("example_case/case-6/prompt.txt", "r", encoding="utf-8") as f:
+ prompt_text = f.read().strip()
+ except:
+ prompt_text = ""
+ return "example_case/case-6/reference.png", "example_case/case-6/audio.wav", prompt_text
+
+ def load_example4():
+ try:
+ with open("example_case/case-45/prompt.txt", "r", encoding="utf-8") as f:
+ prompt_text = f.read().strip()
+ except:
+ prompt_text = ""
+ return "example_case/case-45/reference.png", "example_case/case-45/audio.wav", prompt_text
+
+ def load_example5():
+ try:
+ with open("example_case/case-3/prompt.txt", "r", encoding="utf-8") as f:
+ prompt_text = f.read().strip()
+ except:
+ prompt_text = ""
+ return "example_case/case-3/reference.jpg", "example_case/case-3/audio.wav", prompt_text
+
+ example1_btn.click(fn=load_example1, outputs=[image_path, audio_path, prompt])
+ example2_btn.click(fn=load_example2, outputs=[image_path, audio_path, prompt])
+ example3_btn.click(fn=load_example3, outputs=[image_path, audio_path, prompt])
+ example4_btn.click(fn=load_example4, outputs=[image_path, audio_path, prompt])
+ example5_btn.click(fn=load_example5, outputs=[image_path, audio_path, prompt])
+ gr.on(
+ triggers=[generate_button.click, prompt.submit, negative_prompt.submit],
+ fn = generate,
+ inputs = [
+ GPU_memory_mode,
+ teacache_threshold,
+ num_skip_start_steps,
+ image_path,
+ audio_path,
+ prompt,
+ negative_prompt,
+ width,
+ height,
+ guidance_scale,
+ num_inference_steps,
+ text_guide_scale,
+ audio_guide_scale,
+ motion_frame,
+ fps,
+ overlap_window_length,
+ seed_param,
+ overlapping_weight_scheme,
+ ],
+ outputs = [video_output, seed_output, info]
+ )
+ exchange_button.click(
+ fn=exchange_width_height,
+ inputs=[width, height],
+ outputs=[width, height, info]
+ )
+ adjust_button.click(
+ fn=adjust_width_height,
+ inputs=[image_path],
+ outputs=[width, height, info]
+ )
+ extractor_button.click(
+ fn=audio_extractor,
+ inputs=[video_path],
+ outputs=[audio_output, info2, audio_file]
+ )
+ separation_button.click(
+ fn=vocal_separation,
+ inputs=[audio_path3],
+ outputs=[audio_output3, info3, audio_file3]
+ )
+
+
+if __name__ == "__main__":
+ demo.launch(
+ server_name="0.0.0.0",
+ server_port=int(os.getenv("PORT", 7860)),
+ share=False,
+ inbrowser=False,
+ )
diff --git a/audio_extractor.py b/audio_extractor.py
new file mode 100644
index 0000000000000000000000000000000000000000..1ba206145dbd2ba6158931af0c86fe503f330462
--- /dev/null
+++ b/audio_extractor.py
@@ -0,0 +1,14 @@
+import os
+from moviepy.editor import VideoFileClip
+import argparse
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--video_path", type=str)
+ parser.add_argument("--saved_audio_path", type=str)
+ args = parser.parse_args()
+ video_path = args.video_path
+ saved_audio_path = args.saved_audio_path
+ video = VideoFileClip(video_path)
+ audio = video.audio
+ audio.write_audiofile(saved_audio_path, codec='pcm_s16le')
\ No newline at end of file
diff --git a/deepspeed_config/wan2.1/wan_civitai.yaml b/deepspeed_config/wan2.1/wan_civitai.yaml
new file mode 100644
index 0000000000000000000000000000000000000000..3a2747bdc9e86a6812c3ca09d2a8653aaea6cb54
--- /dev/null
+++ b/deepspeed_config/wan2.1/wan_civitai.yaml
@@ -0,0 +1,39 @@
+format: civitai
+pipeline: Wan
+transformer_additional_kwargs:
+ transformer_subpath: ./
+ dict_mapping:
+ in_dim: in_channels
+ dim: hidden_size
+
+vae_kwargs:
+ vae_subpath: Wan2.1_VAE.pth
+ temporal_compression_ratio: 4
+ spatial_compression_ratio: 8
+
+text_encoder_kwargs:
+ text_encoder_subpath: models_t5_umt5-xxl-enc-bf16.pth
+ tokenizer_subpath: google/umt5-xxl
+ text_length: 512
+ vocab: 256384
+ dim: 4096
+ dim_attn: 4096
+ dim_ffn: 10240
+ num_heads: 64
+ num_layers: 24
+ num_buckets: 32
+ shared_pos: False
+ dropout: 0.0
+
+scheduler_kwargs:
+ scheduler_subpath: null
+ num_train_timesteps: 1000
+ shift: 5.0
+ use_dynamic_shifting: false
+ base_shift: 0.5
+ max_shift: 1.15
+ base_image_seq_len: 256
+ max_image_seq_len: 4096
+
+image_encoder_kwargs:
+ image_encoder_subpath: models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth
\ No newline at end of file
diff --git a/deepspeed_config/zero2_offload_cpu.json b/deepspeed_config/zero2_offload_cpu.json
new file mode 100644
index 0000000000000000000000000000000000000000..a0aff0eaad65a00fe12bc7642569c0458a5549c3
--- /dev/null
+++ b/deepspeed_config/zero2_offload_cpu.json
@@ -0,0 +1,35 @@
+{
+ "bf16": {
+ "enabled": "auto"
+ },
+ "train_micro_batch_size_per_gpu": "auto",
+ "train_batch_size": "auto",
+ "gradient_accumulation_steps": "auto",
+ "zero_optimization": {
+ "stage": 2,
+ "overlap_comm": true,
+ "contiguous_gradients": true,
+ "sub_group_size": 1e9,
+ "offload_optimizer": {
+ "device": "cpu",
+ "pin_memory": true
+ }
+ },
+
+ "optimizer": {
+ "type": "AdamW",
+ "params": {
+ "lr": 5e-5,
+ "betas": [0.9, 0.95],
+ "weight_decay": 0.01
+ }
+ },
+ "scheduler": {
+ "type": "WarmupDecayLR",
+ "params": {
+ "warmup_min_lr": 1e-6,
+ "warmup_max_lr": 5e-5,
+ "total_num_steps": 10000
+ }
+ }
+}
\ No newline at end of file
diff --git a/deepspeed_config/zero_stage2_config.json b/deepspeed_config/zero_stage2_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..3099d255426c78ade836befbd086148093d9049c
--- /dev/null
+++ b/deepspeed_config/zero_stage2_config.json
@@ -0,0 +1,35 @@
+{
+ "bf16": {
+ "enabled": true
+ },
+ "train_micro_batch_size_per_gpu": 1,
+ "train_batch_size": 64,
+ "gradient_clipping": 1.0,
+ "gradient_accumulation_steps": 1,
+ "dump_state": true,
+ "zero_optimization": {
+ "stage": 2,
+ "allgather_partitions": true,
+ "allgather_bucket_size": 2e8,
+ "overlap_comm": true,
+ "reduce_scatter": true,
+ "reduce_bucket_size": 1e8,
+ "contiguous_gradients": true
+ },
+ "optimizer": {
+ "type": "AdamW",
+ "params": {
+ "lr": 1e-4,
+ "betas": [0.9, 0.999],
+ "weight_decay": 3e-2
+ }
+ },
+ "scheduler": {
+ "type": "WarmupLR",
+ "params": {
+ "warmup_min_lr": 1e-7,
+ "warmup_max_lr": 1e-4,
+ "warmup_num_steps": 100
+ }
+ }
+}
diff --git a/deepspeed_config/zero_stage3_config.json b/deepspeed_config/zero_stage3_config.json
new file mode 100644
index 0000000000000000000000000000000000000000..07c4d8194bc135f1345369481a9f6bca33320f92
--- /dev/null
+++ b/deepspeed_config/zero_stage3_config.json
@@ -0,0 +1,46 @@
+{
+ "bf16": {
+ "enabled": true
+ },
+ "optimizer": {
+ "type": "AdamW",
+ "params": {
+ "lr": 2e-5,
+ "betas": [0.9, 0.999],
+ "weight_decay": 3e-2
+ }
+ },
+ "scheduler": {
+ "type": "WarmupLR",
+ "params": {
+ "warmup_min_lr": 1e-7,
+ "warmup_max_lr": 2e-5,
+ "warmup_num_steps": 6400
+ }
+ },
+ "train_micro_batch_size_per_gpu": 1,
+ "gradient_accumulation_steps": 1,
+ "train_batch_size": 64,
+ "gradient_clipping": 1.0,
+ "steps_per_print": 2000,
+ "wall_clock_breakdown": false,
+ "zero_optimization": {
+ "stage": 3,
+ "overlap_comm": true,
+ "contiguous_gradients": true,
+ "reduce_bucket_size": 5e8,
+ "sub_group_size": 1e9,
+ "stage3_max_live_parameters": 1e9,
+ "stage3_max_reuse_distance": 1e9,
+ "stage3_gather_16bit_weights_on_model_save": "auto",
+ "offload_optimizer": {
+ "device": "cpu",
+ "pin_memory": true
+ },
+ "offload_param": {
+ "device": "cpu",
+ "pin_memory": true
+ }
+ }
+}
+
diff --git a/example_case/case-1/audio.wav b/example_case/case-1/audio.wav
new file mode 100644
index 0000000000000000000000000000000000000000..6c0a7fa7c6e812cd0a6348aa4691fda7f644a1a0
--- /dev/null
+++ b/example_case/case-1/audio.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d12a8745971f1472c1ac5b3e3e5349163be7555b187ef3ad3cc4718393174458
+size 17645370
diff --git a/example_case/case-1/prompt.txt b/example_case/case-1/prompt.txt
new file mode 100644
index 0000000000000000000000000000000000000000..6f57790a97a591db1d4312b150266ac454279b62
--- /dev/null
+++ b/example_case/case-1/prompt.txt
@@ -0,0 +1 @@
+Front-facing head-and-shoulders close-up of a middle-aged woman with short light brown hair, pearl earrings, and a blue blazer under soft studio lighting – She delivers a clear, confident speech with precise lip movements, steady gaze toward the camera, subtle eyebrow emphasis, slight nods, and occasional blinks while maintaining composed posture – Blurred civic architecture in the background resembling a government building, shallow depth of field, static camera.
\ No newline at end of file
diff --git a/example_case/case-1/reference.png b/example_case/case-1/reference.png
new file mode 100644
index 0000000000000000000000000000000000000000..4a425fcad9addecf363a0ca2531cbeb3cfce449f
--- /dev/null
+++ b/example_case/case-1/reference.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:238117e216e32488130fa3b1ef4147339e7b2f157dcf7e62af46f3be34d19a89
+size 657395
diff --git a/example_case/case-2/audio.wav b/example_case/case-2/audio.wav
new file mode 100644
index 0000000000000000000000000000000000000000..fea194ae60b4e3fffb91fb972798b4032df7bcb9
--- /dev/null
+++ b/example_case/case-2/audio.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eb004cdfc7ba33e44c4128c555f43bbfdd88049a8937b1c5585db56efd59da15
+size 2568018
diff --git a/example_case/case-2/prompt.txt b/example_case/case-2/prompt.txt
new file mode 100644
index 0000000000000000000000000000000000000000..fa0d7b19845579f13d835d6a25ab2ae6705cced3
--- /dev/null
+++ b/example_case/case-2/prompt.txt
@@ -0,0 +1 @@
+Front-facing head-and-shoulders close-up of a middle-aged man with a shaved head, thin-rim glasses, and a striped shirt under soft warm lighting – He speaks clearly and thoughtfully with precise lip-sync, subtle eyebrow movement, slight nods, and occasional blinks while maintaining a steady posture – Indoor studio with blurred shutters and two warm pendant lights, shallow depth of field, and a static camera.
\ No newline at end of file
diff --git a/example_case/case-2/reference.png b/example_case/case-2/reference.png
new file mode 100644
index 0000000000000000000000000000000000000000..eaec45494a53cafc3867e85c17e131e4f2ccdbc5
--- /dev/null
+++ b/example_case/case-2/reference.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0282b2e0a1401db84ae97d06adadea1d58b9c61b2e632bcf2a377d55eacdaecf
+size 704823
diff --git a/example_case/case-3/audio.wav b/example_case/case-3/audio.wav
new file mode 100644
index 0000000000000000000000000000000000000000..c11a5ae64729679b68e564c9faedc4f222701b0e
--- /dev/null
+++ b/example_case/case-3/audio.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d2162d8ca2e9ff692c132683c9f197cc73c84a6bb6cd8a3ed5aeefbc4711ad87
+size 168014
diff --git a/example_case/case-3/prompt.txt b/example_case/case-3/prompt.txt
new file mode 100644
index 0000000000000000000000000000000000000000..8facbd5943da880cad4ce49ca44621e90f224790
--- /dev/null
+++ b/example_case/case-3/prompt.txt
@@ -0,0 +1 @@
+Front-facing head-and-shoulders close-up of an adult woman with wavy dark brown hair and silver hoop earrings under soft warm lighting – She sings “there once was a ship that put to sea, the name of the ship was the Billy” with precise lip-sync, steady tempo, subtle head sway, gentle eyebrow lifts, and occasional blinks while maintaining a composed posture – Indoor studio with a softly blurred background and warm bokeh, shallow depth of field, static camera.
\ No newline at end of file
diff --git a/example_case/case-3/reference.jpg b/example_case/case-3/reference.jpg
new file mode 100644
index 0000000000000000000000000000000000000000..476498c443b1d410bb6e22e01769bab4b78b99f2
--- /dev/null
+++ b/example_case/case-3/reference.jpg
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5875d5c60ed74322fb878604c8e8548aefe4a45bdbbdb998069e398824080b74
+size 26854
diff --git a/example_case/case-45/audio.wav b/example_case/case-45/audio.wav
new file mode 100644
index 0000000000000000000000000000000000000000..713613ef9957724927ed5947db2dd7a900c5515c
--- /dev/null
+++ b/example_case/case-45/audio.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cd5db4a3d4a970a51729aff7001ac34ffb13c486c62c7e44759023403121db66
+size 3076494
diff --git a/example_case/case-45/prompt.txt b/example_case/case-45/prompt.txt
new file mode 100644
index 0000000000000000000000000000000000000000..895cb30abf1e375efe0c5b03d75f241dff97ae5a
--- /dev/null
+++ b/example_case/case-45/prompt.txt
@@ -0,0 +1 @@
+Front-facing medium close-up of a young woman with long silver hair, elf-like ears, a cozy oversized light blue scarf, and a white outfit under soft daylight – She sings a sweet, lighthearted melody with precise lip-sync, a gentle smile, relaxed breathing, subtle head sway, and natural blinks while maintaining a warm and calm demeanor – Cozy indoor room with soft light, bed and curtain in the background, shallow depth of field, static camera.
\ No newline at end of file
diff --git a/example_case/case-45/reference.png b/example_case/case-45/reference.png
new file mode 100644
index 0000000000000000000000000000000000000000..cc4a344478cbaac131c2bd5a3cdf59e54a520274
--- /dev/null
+++ b/example_case/case-45/reference.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:859b53178227630da809ec6be49894c6b26da340aae1532c90167a7736125e59
+size 49369
diff --git a/example_case/case-6/audio.wav b/example_case/case-6/audio.wav
new file mode 100644
index 0000000000000000000000000000000000000000..e32ba9a35e943e7dd7aa0a8dd3c9f434b9b6c8d5
--- /dev/null
+++ b/example_case/case-6/audio.wav
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c4a1460a58d3a7662cb17494e99edc590bbab2254c80fbbd8d5c0ce327645c39
+size 5821278
diff --git a/example_case/case-6/prompt.txt b/example_case/case-6/prompt.txt
new file mode 100644
index 0000000000000000000000000000000000000000..aa72627b3e0d15d011c2163d1a4b3681f52427f0
--- /dev/null
+++ b/example_case/case-6/prompt.txt
@@ -0,0 +1 @@
+Front-facing medium close-up of a young woman with shoulder-length dark hair, wearing a white top and small hoop earrings, a studio microphone visible in the lower left under soft daylight – She sings smoothly with precise lip-sync, relaxed breathing, gentle head sway, subtle eyebrow emphasis, and natural blinks while maintaining a calm posture – Minimal indoor setting with a light gray wall and decorative molding, diagonal light and soft shadows, shallow depth of field, static camera.
\ No newline at end of file
diff --git a/example_case/case-6/reference.png b/example_case/case-6/reference.png
new file mode 100644
index 0000000000000000000000000000000000000000..698318da02b2728da9be613ed134c58a0377999e
--- /dev/null
+++ b/example_case/case-6/reference.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:43717c247b9c19950216b409632c0daae8b6568ed1c58cf5abf364b3469a4569
+size 1024823
diff --git a/extract_audio_segment.py b/extract_audio_segment.py
new file mode 100644
index 0000000000000000000000000000000000000000..73d938c78d002fe3d964259a56696672797ccb19
--- /dev/null
+++ b/extract_audio_segment.py
@@ -0,0 +1,146 @@
+#!/usr/bin/env python3
+"""
+音频文件转换和片段提取工具
+将MP3文件转换为WAV格式，并提取指定时间段的音频片段
+"""
+
+import os
+import subprocess
+from pathlib import Path
+
+def convert_mp3_to_wav_and_extract(input_file, start_time, end_time, output_dir=None):
+ """
+ 将MP3文件转换为WAV格式，并提取指定时间段的音频片段
+
+ Args:
+ input_file (str): 输入的MP3文件路径
+ start_time (float): 开始时间（秒）
+ end_time (float): 结束时间（秒）
+ output_dir (str): 输出目录，如果为None则使用输入文件所在目录
+
+ Returns:
+ bool: 操作是否成功
+ """
+ try:
+ # 检查输入文件是否存在
+ if not os.path.exists(input_file):
+ print(f"❌ 错误：输入文件不存在: {input_file}")
+ return False
+
+ # 设置输出目录
+ if output_dir is None:
+ output_dir = os.path.dirname(input_file)
+
+ # 确保输出目录存在
+ Path(output_dir).mkdir(parents=True, exist_ok=True)
+
+ # 生成输出文件名
+ input_name = Path(input_file).stem
+ output_wav = os.path.join(output_dir, f"{input_name}.wav")
+ output_segment = os.path.join(output_dir, f"{input_name}_segment_{start_time}s_to_{end_time}s.wav")
+
+ print(f"🎵 开始处理音频文件: {input_file}")
+ print(f"📁 输出目录: {output_dir}")
+
+ # 步骤1：将MP3转换为WAV格式
+ print(f"\n🔄 步骤1: 将MP3转换为WAV格式")
+ convert_cmd = [
+ 'ffmpeg',
+ '-y', # 覆盖输出文件
+ '-i', input_file, # 输入文件
+ '-ar', '16000', # 采样率16kHz
+ '-ac', '1', # 单声道
+ '-c:a', 'pcm_s16le', # 16位PCM编码
+ output_wav # 输出文件
+ ]
+
+ print(f"执行命令: {' '.join(convert_cmd)}")
+ result = subprocess.run(convert_cmd, capture_output=True, text=True)
+
+ if result.returncode != 0:
+ print(f"❌ MP3转WAV失败: {result.stderr}")
+ return False
+
+ print(f"✅ MP3转WAV成功: {output_wav}")
+
+ # 步骤2：提取音频片段
+ print(f"\n🔄 步骤2: 提取音频片段 ({start_time}s - {end_time}s)")
+ duration = end_time - start_time
+
+ extract_cmd = [
+ 'ffmpeg',
+ '-y', # 覆盖输出文件
+ '-i', output_wav, # 输入WAV文件
+ '-ss', str(start_time), # 开始时间
+ '-t', str(duration), # 持续时间
+ '-c', 'copy', # 直接复制，不重新编码
+ output_segment # 输出片段文件
+ ]
+
+ print(f"执行命令: {' '.join(extract_cmd)}")
+ result = subprocess.run(extract_cmd, capture_output=True, text=True)
+
+ if result.returncode != 0:
+ print(f"❌ 音频片段提取失败: {result.stderr}")
+ return False
+
+ print(f"✅ 音频片段提取成功: {output_segment}")
+
+ # 显示文件信息
+ print(f"\n📊 文件信息:")
+ print(f"原始MP3文件: {input_file}")
+ print(f"转换后的WAV文件: {output_wav}")
+ print(f"提取的音频片段: {output_segment}")
+ print(f"片段时长: {duration:.1f}秒")
+
+ # 检查输出文件大小
+ if os.path.exists(output_wav):
+ wav_size = os.path.getsize(output_wav) / 1024 # KB
+ print(f"WAV文件大小: {wav_size:.1f} KB")
+
+ if os.path.exists(output_segment):
+ segment_size = os.path.getsize(output_segment) / 1024 # KB
+ print(f"片段文件大小: {segment_size:.1f} KB")
+
+ return True
+
+ except Exception as e:
+ print(f"❌ 处理过程中出现错误: {str(e)}")
+ return False
+
+def main():
+ """主函数"""
+ print("🎵 音频文件转换和片段提取工具")
+ print("=" * 50)
+
+ # 检查ffmpeg是否安装
+ try:
+ subprocess.run(['ffmpeg', '-version'], capture_output=True, check=True)
+ print("✅ 检测到ffmpeg")
+ except (subprocess.CalledProcessError, FileNotFoundError):
+ print("❌ 错误：未找到ffmpeg，请先安装ffmpeg")
+ print("Ubuntu/Debian: sudo apt install ffmpeg")
+ print("CentOS/RHEL: sudo yum install ffmpeg")
+ print("macOS: brew install ffmpeg")
+ return
+
+ # 设置文件路径和时间参数
+ input_file = "/home/t2vg-a100-G4-42/v-shuyuantu/StableAvatar/example_case/case-3/ssvid.net--Wellerman-Female-Cover-LYRICS-Sea-Shanty.mp3"
+ start_time = 1.9 # 开始时间（秒）
+ end_time = 7.1 # 结束时间（秒）
+
+ print(f"📁 输入文件: {input_file}")
+ print(f"⏰ 提取时间段: {start_time}s - {end_time}s")
+ print(f"⏱️ 片段时长: {end_time - start_time:.1f}秒")
+
+ # 执行转换和提取
+ success = convert_mp3_to_wav_and_extract(input_file, start_time, end_time)
+
+ if success:
+ print(f"\n🎉 所有操作完成！")
+ print(f"输出文件保存在: {os.path.dirname(input_file)}")
+ else:
+ print(f"\n❌ 操作失败，请检查错误信息")
+
+if __name__ == "__main__":
+ main()
diff --git a/lip_mask_extractor.py b/lip_mask_extractor.py
new file mode 100644
index 0000000000000000000000000000000000000000..a35306cf022a374248d29e934021826fbcfd1abb
--- /dev/null
+++ b/lip_mask_extractor.py
@@ -0,0 +1,70 @@
+import argparse
+import os
+
+import cv2
+import mediapipe as mp
+import numpy as np
+
+if __name__ == "__main__":
+
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--folder_root", type=str)
+ parser.add_argument("--start", type=int, help="Specify the value of start")
+ parser.add_argument("--end", type=int, help="Specify the value of end")
+ args = parser.parse_args()
+
+ folder_root = args.folder_root
+ start = args.start
+ end = args.end
+
+ mp_face_mesh = mp.solutions.face_mesh
+ face_mesh = mp_face_mesh.FaceMesh(static_image_mode=True, max_num_faces=10)
+
+ upper_lip_idx = [61, 185, 40, 39, 37, 0, 267, 269, 270, 409, 291]
+ lower_lip_idx = [61, 146, 91, 181, 84, 17, 314, 405, 321, 375, 291]
+
+ for idx in range(start, end):
+ subfolder = str(idx).zfill(5)
+ subfolder_path = os.path.join(folder_root, subfolder)
+ images_folder = os.path.join(subfolder_path, "images")
+ if os.path.exists(images_folder):
+ face_masks_folder = os.path.join(subfolder_path, "lip_masks")
+ os.makedirs(face_masks_folder, exist_ok=True)
+ for root, dirs, files in os.walk(images_folder):
+ for file in files:
+ if file.endswith('.png'):
+ file_name = os.path.splitext(file)[0]
+ image_name = file_name + '.png'
+ image_legal_path = os.path.join(images_folder, image_name)
+ if os.path.exists(os.path.join(face_masks_folder, file_name + '.png')):
+ existed_path = os.path.join(face_masks_folder, file_name + '.png')
+ print(f"{existed_path} already exists!")
+ continue
+
+ face_save_path = os.path.join(face_masks_folder, file_name + '.png')
+
+ image = cv2.imread(image_legal_path)
+ h, w, _ = image.shape
+ rgb_image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+ results = face_mesh.process(rgb_image)
+ mask = np.zeros((h, w), dtype=np.uint8)
+
+ if results.multi_face_landmarks:
+ for face_landmarks in results.multi_face_landmarks:
+ upper_points = np.array([
+ [int(face_landmarks.landmark[i].x * w), int(face_landmarks.landmark[i].y * h)]
+ for i in upper_lip_idx
+ ], dtype=np.int32)
+ lower_points = np.array([
+ [int(face_landmarks.landmark[i].x * w), int(face_landmarks.landmark[i].y * h)]
+ for i in lower_lip_idx
+ ], dtype=np.int32)
+ cv2.fillPoly(mask, [upper_points], 255)
+ cv2.fillPoly(mask, [lower_points], 255)
+ else:
+ print(f"No face detected in {image_legal_path}. Saving empty mask.")
+ cv2.imwrite(face_save_path, mask)
+ print(f"Lip mask saved to {face_save_path}")
+ else:
+ print(f"{images_folder} does not exist")
+ continue
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000000000000000000000000000000000000..cfc80836efe6867c4df347e6146c5159954474df
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,170 @@
+absl-py==2.3.1
+accelerate==1.10.0
+aiofiles==24.1.0
+aiohappyeyeballs==2.6.1
+aiohttp==3.12.15
+aiosignal==1.4.0
+albucore==0.0.24
+albumentations==2.0.8
+annotated-types==0.7.0
+antlr4-python3-runtime==4.9.3
+anyio==4.10.0
+attrs==25.3.0
+audio-separator==0.36.1
+audioread==3.0.1
+av==15.0.0
+beartype==0.18.5
+beautifulsoup4==4.13.4
+Brotli==1.1.0
+certifi==2025.8.3
+cffi==1.17.1
+charset-normalizer==3.4.3
+click==8.2.1
+coloredlogs==15.0.1
+cryptography==45.0.6
+Cython==3.1.3
+dashscope==1.24.1
+datasets==4.0.0
+decorator==4.4.2
+decord==0.6.0
+diffq==0.2.4
+diffusers==0.30.1
+dill==0.3.8
+easydict==1.13
+einops==0.8.1
+fastapi==0.116.1
+ffmpy==0.6.1
+filelock==3.13.1
+flatbuffers==25.2.10
+frozenlist==1.7.0
+fsspec==2024.6.1
+ftfy==6.3.1
+gradio==5.42.0
+gradio_client==1.11.1
+groovy==0.1.2
+grpcio==1.74.0
+h11==0.16.0
+hf-xet==1.1.7
+httpcore==1.0.9
+httpx==0.28.1
+huggingface-hub==0.34.4
+humanfriendly==10.0
+idna==3.10
+imageio==2.37.0
+imageio-ffmpeg==0.6.0
+importlib_metadata==8.7.0
+Jinja2==3.1.4
+joblib==1.5.1
+julius==0.2.7
+lazy_loader==0.4
+librosa==0.11.0
+llvmlite==0.44.0
+Markdown==3.8.2
+markdown-it-py==4.0.0
+MarkupSafe==2.1.5
+mdurl==0.1.2
+ml_collections==1.1.0
+ml_dtypes==0.5.3
+moviepy==1.0.3
+mpmath==1.3.0
+msgpack==1.1.1
+multidict==6.6.4
+multiprocess==0.70.16
+networkx==3.3
+ninja==1.13.0
+numba==0.61.2
+numpy==2.2.6
+nvidia-cublas-cu12==12.4.5.8
+nvidia-cuda-cupti-cu12==12.4.127
+nvidia-cuda-nvrtc-cu12==12.4.127
+nvidia-cuda-runtime-cu12==12.4.127
+nvidia-cudnn-cu12==9.1.0.70
+nvidia-cufft-cu12==11.2.1.3
+nvidia-curand-cu12==10.3.5.147
+nvidia-cusolver-cu12==11.6.1.9
+nvidia-cusparse-cu12==12.3.1.170
+nvidia-cusparselt-cu12==0.6.2
+nvidia-nccl-cu12==2.21.5
+nvidia-nvjitlink-cu12==12.4.127
+nvidia-nvtx-cu12==12.4.127
+omegaconf==2.3.0
+onnx-weekly==1.19.0.dev20250726
+onnx2torch-py313==1.6.0
+onnxruntime-gpu==1.22.0
+opencv-python==4.11.0.86
+opencv-python-headless==4.11.0.86
+orjson==3.11.2
+packaging==25.0
+pandas==2.3.1
+pillow==11.0.0
+platformdirs==4.3.8
+pooch==1.8.2
+proglog==0.1.12
+propcache==0.3.2
+protobuf==6.31.1
+psutil==7.0.0
+pyarrow==21.0.0
+pycparser==2.22
+pydantic==2.11.7
+pydantic_core==2.33.2
+pydub==0.25.1
+Pygments==2.19.2
+python-dateutil==2.9.0.post0
+python-dotenv==1.1.1
+python-multipart==0.0.20
+pytz==2025.2
+PyYAML==6.0.2
+regex==2025.7.34
+requests==2.32.4
+resampy==0.4.3
+rich==14.1.0
+rotary-embedding-torch==0.6.5
+ruff==0.12.8
+safehttpx==0.1.6
+safetensors==0.6.2
+samplerate==0.1.0
+scikit-image==0.25.2
+scikit-learn==1.7.1
+scipy==1.16.1
+semantic-version==2.10.0
+sentencepiece==0.2.1
+shellingham==1.5.4
+simsimd==6.5.0
+six==1.17.0
+sniffio==1.3.1
+soundfile==0.13.1
+soupsieve==2.7
+soxr==0.5.0.post1
+starlette==0.47.2
+stringzilla==3.12.6
+sympy==1.13.1
+tensorboard==2.20.0
+tensorboard-data-server==0.7.2
+threadpoolctl==3.6.0
+tifffile==2025.6.11
+timm==1.0.19
+tokenizers==0.21.4
+tomesd==0.1.3
+tomlkit==0.13.3
+torch==2.6.0+cu124
+torchaudio==2.6.0+cu124
+torchdiffeq==0.2.5
+torchsde==0.2.6
+torchvision==0.21.0+cu124
+tqdm==4.67.1
+trampoline==0.1.2
+transformers==4.51.3
+triton==3.2.0
+typer==0.16.0
+typing-inspection==0.4.1
+typing_extensions==4.12.2
+tzdata==2025.2
+urllib3==2.5.0
+uvicorn==0.35.0
+wcwidth==0.2.13
+websocket-client==1.8.0
+websockets==15.0.1
+Werkzeug==3.1.3
+xxhash==3.5.0
+yarl==1.20.1
+zipp==3.23.0
diff --git a/vocal_seperator.py b/vocal_seperator.py
new file mode 100644
index 0000000000000000000000000000000000000000..3ecdb84400457a4a08946b4a72600096c2bd5e6b
--- /dev/null
+++ b/vocal_seperator.py
@@ -0,0 +1,31 @@
+import argparse
+import os
+import shutil
+import subprocess
+from audio_separator.separator import Separator
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--audio_file_path", type=str)
+ parser.add_argument("--saved_vocal_path", type=str)
+ parser.add_argument("--audio_separator_model_file", type=str)
+ args = parser.parse_args()
+ audio_file_path = args.audio_file_path
+ audio_separator_model_file = args.audio_separator_model_file
+ saved_vocal_path = args.saved_vocal_path
+ cache_dir = os.path.join(os.path.dirname(audio_file_path), "vocals")
+ os.makedirs(cache_dir, exist_ok=True)
+ audio_separator = Separator(
+ output_dir=cache_dir,
+ output_single_stem="vocals",
+ model_file_dir=os.path.dirname(audio_separator_model_file),
+ )
+ audio_separator.load_model(os.path.basename(audio_separator_model_file))
+ assert audio_separator.model_instance is not None, "Fail to load audio separate model."
+ outputs = audio_separator.separate(audio_file_path)
+ subfolder_path = os.path.dirname(audio_file_path)
+ vocal_audio_file = os.path.join(audio_separator.output_dir, outputs[0])
+ destination_file = os.path.join(subfolder_path, "vocal.wav")
+ shutil.copy(vocal_audio_file, destination_file)
+ os.remove(vocal_audio_file)
diff --git a/wan/__init__.py b/wan/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..62b57c67879519f8fea9b97f71e335affed761e1
--- /dev/null
+++ b/wan/__init__.py
@@ -0,0 +1,3 @@
+# from . import configs, distributed, modules
+# from .image2video import WanI2V
+# from .text2video import WanT2V
diff --git a/wan/__pycache__/__init__.cpython-311.pyc b/wan/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..28206bde8d0ecc3e972a36406abffba1f010f8df
Binary files /dev/null and b/wan/__pycache__/__init__.cpython-311.pyc differ
diff --git a/wan/configs/__init__.py b/wan/configs/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..c72d2d01be834882d659701fc0dc67beb152383f
--- /dev/null
+++ b/wan/configs/__init__.py
@@ -0,0 +1,42 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import copy
+import os
+
+os.environ['TOKENIZERS_PARALLELISM'] = 'false'
+
+from .wan_i2v_14B import i2v_14B
+from .wan_t2v_1_3B import t2v_1_3B
+from .wan_t2v_14B import t2v_14B
+
+# the config of t2i_14B is the same as t2v_14B
+t2i_14B = copy.deepcopy(t2v_14B)
+t2i_14B.__name__ = 'Config: Wan T2I 14B'
+
+WAN_CONFIGS = {
+ 't2v-14B': t2v_14B,
+ 't2v-1.3B': t2v_1_3B,
+ 'i2v-14B': i2v_14B,
+ 't2i-14B': t2i_14B,
+}
+
+SIZE_CONFIGS = {
+ '720*1280': (720, 1280),
+ '1280*720': (1280, 720),
+ '480*832': (480, 832),
+ '832*480': (832, 480),
+ '1024*1024': (1024, 1024),
+}
+
+MAX_AREA_CONFIGS = {
+ '720*1280': 720 * 1280,
+ '1280*720': 1280 * 720,
+ '480*832': 480 * 832,
+ '832*480': 832 * 480,
+}
+
+SUPPORTED_SIZES = {
+ 't2v-14B': ('720*1280', '1280*720', '480*832', '832*480'),
+ 't2v-1.3B': ('480*832', '832*480'),
+ 'i2v-14B': ('720*1280', '1280*720', '480*832', '832*480'),
+ 't2i-14B': tuple(SIZE_CONFIGS.keys()),
+}
diff --git a/wan/configs/shared_config.py b/wan/configs/shared_config.py
new file mode 100644
index 0000000000000000000000000000000000000000..04a9f454218fc1ce958b628e71ad5738222e2aa4
--- /dev/null
+++ b/wan/configs/shared_config.py
@@ -0,0 +1,19 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+from easydict import EasyDict
+
+#------------------------ Wan shared config ------------------------#
+wan_shared_cfg = EasyDict()
+
+# t5
+wan_shared_cfg.t5_model = 'umt5_xxl'
+wan_shared_cfg.t5_dtype = torch.bfloat16
+wan_shared_cfg.text_len = 512
+
+# transformer
+wan_shared_cfg.param_dtype = torch.bfloat16
+
+# inference
+wan_shared_cfg.num_train_timesteps = 1000
+wan_shared_cfg.sample_fps = 16
+wan_shared_cfg.sample_neg_prompt = '色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走'
diff --git a/wan/configs/wan_i2v_14B.py b/wan/configs/wan_i2v_14B.py
new file mode 100644
index 0000000000000000000000000000000000000000..12e8e205bffb343a6e27d2828fb573db1d6349f8
--- /dev/null
+++ b/wan/configs/wan_i2v_14B.py
@@ -0,0 +1,35 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+from easydict import EasyDict
+
+from .shared_config import wan_shared_cfg
+
+#------------------------ Wan I2V 14B ------------------------#
+
+i2v_14B = EasyDict(__name__='Config: Wan I2V 14B')
+i2v_14B.update(wan_shared_cfg)
+
+i2v_14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+i2v_14B.t5_tokenizer = 'google/umt5-xxl'
+
+# clip
+i2v_14B.clip_model = 'clip_xlm_roberta_vit_h_14'
+i2v_14B.clip_dtype = torch.float16
+i2v_14B.clip_checkpoint = 'models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth'
+i2v_14B.clip_tokenizer = 'xlm-roberta-large'
+
+# vae
+i2v_14B.vae_checkpoint = 'Wan2.1_VAE.pth'
+i2v_14B.vae_stride = (4, 8, 8)
+
+# transformer
+i2v_14B.patch_size = (1, 2, 2)
+i2v_14B.dim = 5120
+i2v_14B.ffn_dim = 13824
+i2v_14B.freq_dim = 256
+i2v_14B.num_heads = 40
+i2v_14B.num_layers = 40
+i2v_14B.window_size = (-1, -1)
+i2v_14B.qk_norm = True
+i2v_14B.cross_attn_norm = True
+i2v_14B.eps = 1e-6
diff --git a/wan/configs/wan_t2v_14B.py b/wan/configs/wan_t2v_14B.py
new file mode 100644
index 0000000000000000000000000000000000000000..9d0ee69dea796bfd6eccdedf4ec04835086227a6
--- /dev/null
+++ b/wan/configs/wan_t2v_14B.py
@@ -0,0 +1,29 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from easydict import EasyDict
+
+from .shared_config import wan_shared_cfg
+
+#------------------------ Wan T2V 14B ------------------------#
+
+t2v_14B = EasyDict(__name__='Config: Wan T2V 14B')
+t2v_14B.update(wan_shared_cfg)
+
+# t5
+t2v_14B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+t2v_14B.t5_tokenizer = 'google/umt5-xxl'
+
+# vae
+t2v_14B.vae_checkpoint = 'Wan2.1_VAE.pth'
+t2v_14B.vae_stride = (4, 8, 8)
+
+# transformer
+t2v_14B.patch_size = (1, 2, 2)
+t2v_14B.dim = 5120
+t2v_14B.ffn_dim = 13824
+t2v_14B.freq_dim = 256
+t2v_14B.num_heads = 40
+t2v_14B.num_layers = 40
+t2v_14B.window_size = (-1, -1)
+t2v_14B.qk_norm = True
+t2v_14B.cross_attn_norm = True
+t2v_14B.eps = 1e-6
diff --git a/wan/configs/wan_t2v_1_3B.py b/wan/configs/wan_t2v_1_3B.py
new file mode 100644
index 0000000000000000000000000000000000000000..ea9502b0df685b5d22f9091cc8cdf5c6a7880c4b
--- /dev/null
+++ b/wan/configs/wan_t2v_1_3B.py
@@ -0,0 +1,29 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from easydict import EasyDict
+
+from .shared_config import wan_shared_cfg
+
+#------------------------ Wan T2V 1.3B ------------------------#
+
+t2v_1_3B = EasyDict(__name__='Config: Wan T2V 1.3B')
+t2v_1_3B.update(wan_shared_cfg)
+
+# t5
+t2v_1_3B.t5_checkpoint = 'models_t5_umt5-xxl-enc-bf16.pth'
+t2v_1_3B.t5_tokenizer = 'google/umt5-xxl'
+
+# vae
+t2v_1_3B.vae_checkpoint = 'Wan2.1_VAE.pth'
+t2v_1_3B.vae_stride = (4, 8, 8)
+
+# transformer
+t2v_1_3B.patch_size = (1, 2, 2)
+t2v_1_3B.dim = 1536
+t2v_1_3B.ffn_dim = 8960
+t2v_1_3B.freq_dim = 256
+t2v_1_3B.num_heads = 12
+t2v_1_3B.num_layers = 30
+t2v_1_3B.window_size = (-1, -1)
+t2v_1_3B.qk_norm = True
+t2v_1_3B.cross_attn_norm = True
+t2v_1_3B.eps = 1e-6
diff --git a/wan/dataset/talking_video_dataset_fantasy.py b/wan/dataset/talking_video_dataset_fantasy.py
new file mode 100644
index 0000000000000000000000000000000000000000..b23b7969ef9c210e9808e4c952a91e14a06cd64f
--- /dev/null
+++ b/wan/dataset/talking_video_dataset_fantasy.py
@@ -0,0 +1,328 @@
+import math
+import os
+import random
+import warnings
+import librosa
+import numpy as np
+import torch
+from PIL import Image
+import cv2
+from einops import rearrange
+import torchvision.transforms.functional as TF
+from torch.utils.data.dataset import Dataset
+import torch.nn.functional as F
+
+
+def get_random_mask(shape, image_start_only=False):
+ f, c, h, w = shape
+ mask = torch.zeros((f, 1, h, w), dtype=torch.uint8)
+
+ if not image_start_only:
+ if f != 1:
+ mask_index = np.random.choice([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], p=[0.05, 0.2, 0.2, 0.2, 0.05, 0.05, 0.05, 0.1, 0.05, 0.05])
+ else:
+ mask_index = np.random.choice([0, 1], p = [0.2, 0.8])
+ if mask_index == 0:
+ center_x = torch.randint(0, w, (1,)).item()
+ center_y = torch.randint(0, h, (1,)).item()
+ block_size_x = torch.randint(w // 4, w // 4 * 3, (1,)).item() # 方块的宽度范围
+ block_size_y = torch.randint(h // 4, h // 4 * 3, (1,)).item() # 方块的高度范围
+
+ start_x = max(center_x - block_size_x // 2, 0)
+ end_x = min(center_x + block_size_x // 2, w)
+ start_y = max(center_y - block_size_y // 2, 0)
+ end_y = min(center_y + block_size_y // 2, h)
+ mask[:, :, start_y:end_y, start_x:end_x] = 1
+ elif mask_index == 1:
+ mask[:, :, :, :] = 1
+ elif mask_index == 2:
+ mask_frame_index = np.random.randint(1, 5)
+ mask[mask_frame_index:, :, :, :] = 1
+ elif mask_index == 3:
+ mask_frame_index = np.random.randint(1, 5)
+ mask[mask_frame_index:-mask_frame_index, :, :, :] = 1
+ elif mask_index == 4:
+ center_x = torch.randint(0, w, (1,)).item()
+ center_y = torch.randint(0, h, (1,)).item()
+ block_size_x = torch.randint(w // 4, w // 4 * 3, (1,)).item() # 方块的宽度范围
+ block_size_y = torch.randint(h // 4, h // 4 * 3, (1,)).item() # 方块的高度范围
+
+ start_x = max(center_x - block_size_x // 2, 0)
+ end_x = min(center_x + block_size_x // 2, w)
+ start_y = max(center_y - block_size_y // 2, 0)
+ end_y = min(center_y + block_size_y // 2, h)
+
+ mask_frame_before = np.random.randint(0, f // 2)
+ mask_frame_after = np.random.randint(f // 2, f)
+ mask[mask_frame_before:mask_frame_after, :, start_y:end_y, start_x:end_x] = 1
+ elif mask_index == 5:
+ mask = torch.randint(0, 2, (f, 1, h, w), dtype=torch.uint8)
+ elif mask_index == 6:
+ num_frames_to_mask = random.randint(1, max(f // 2, 1))
+ frames_to_mask = random.sample(range(f), num_frames_to_mask)
+
+ for i in frames_to_mask:
+ block_height = random.randint(1, h // 4)
+ block_width = random.randint(1, w // 4)
+ top_left_y = random.randint(0, h - block_height)
+ top_left_x = random.randint(0, w - block_width)
+ mask[i, 0, top_left_y:top_left_y + block_height, top_left_x:top_left_x + block_width] = 1
+ elif mask_index == 7:
+ center_x = torch.randint(0, w, (1,)).item()
+ center_y = torch.randint(0, h, (1,)).item()
+ a = torch.randint(min(w, h) // 8, min(w, h) // 4, (1,)).item() # 长半轴
+ b = torch.randint(min(h, w) // 8, min(h, w) // 4, (1,)).item() # 短半轴
+
+ for i in range(h):
+ for j in range(w):
+ if ((i - center_y) ** 2) / (b ** 2) + ((j - center_x) ** 2) / (a ** 2) < 1:
+ mask[:, :, i, j] = 1
+ elif mask_index == 8:
+ center_x = torch.randint(0, w, (1,)).item()
+ center_y = torch.randint(0, h, (1,)).item()
+ radius = torch.randint(min(h, w) // 8, min(h, w) // 4, (1,)).item()
+ for i in range(h):
+ for j in range(w):
+ if (i - center_y) ** 2 + (j - center_x) ** 2 < radius ** 2:
+ mask[:, :, i, j] = 1
+ elif mask_index == 9:
+ for idx in range(f):
+ if np.random.rand() > 0.5:
+ mask[idx, :, :, :] = 1
+ else:
+ raise ValueError(f"The mask_index {mask_index} is not define")
+ else:
+ if f != 1:
+ mask[1:, :, :, :] = 1
+ else:
+ mask[:, :, :, :] = 1
+ return mask
+
+
+class LargeScaleTalkingFantasyVideos(Dataset):
+ def __init__(self, txt_path, width, height, n_sample_frames, sample_frame_rate, only_last_features=False, vocal_encoder=None, audio_encoder=None, vocal_sample_rate=16000, audio_sample_rate=24000, enable_inpaint=True, audio_margin=2, vae_stride=None, patch_size=None, wav2vec_processor=None, wav2vec=None):
+ self.txt_path = txt_path
+ self.width = width
+ self.height = height
+ self.n_sample_frames = n_sample_frames
+ self.sample_frame_rate = sample_frame_rate
+ self.only_last_features = only_last_features
+ self.vocal_encoder = vocal_encoder
+ self.audio_encoder = audio_encoder
+ self.vocal_sample_rate = vocal_sample_rate
+ self.audio_sample_rate = audio_sample_rate
+ self.enable_inpaint = enable_inpaint
+ self.wav2vec_processor = wav2vec_processor
+ self.audio_margin = audio_margin
+ self.vae_stride = vae_stride
+ self.patch_size = patch_size
+ self.max_area = height * width
+ self.aspect_ratio = height / width
+ self.video_files = self._read_txt_file_images()
+
+ self.lat_h = round(
+ np.sqrt(self.max_area * self.aspect_ratio) // self.vae_stride[1] //
+ self.patch_size[1] * self.patch_size[1])
+ self.lat_w = round(
+ np.sqrt(self.max_area / self.aspect_ratio) // self.vae_stride[2] //
+ self.patch_size[2] * self.patch_size[2])
+
+ def _read_txt_file_images(self):
+ with open(self.txt_path, 'r') as file:
+ lines = file.readlines()
+ video_files = []
+ for line in lines:
+ video_file = line.strip()
+ video_files.append(video_file)
+ return video_files
+
+ def __len__(self):
+ return len(self.video_files)
+
+ def frame_count(self, frames_path):
+ files = os.listdir(frames_path)
+ png_files = [file for file in files if file.endswith('.png') or file.endswith('.jpg')]
+ png_files_count = len(png_files)
+ return png_files_count
+
+ def find_frames_list(self, frames_path):
+ files = os.listdir(frames_path)
+ image_files = [file for file in files if file.endswith('.png') or file.endswith('.jpg')]
+ if image_files[0].startswith('frame_'):
+ image_files.sort(key=lambda x: int(x.split('_')[1].split('.')[0]))
+ else:
+ image_files.sort(key=lambda x: int(x.split('.')[0]))
+ return image_files
+
+ def __getitem__(self, idx):
+ warnings.filterwarnings('ignore', category=DeprecationWarning)
+ warnings.filterwarnings('ignore', category=FutureWarning)
+
+ video_path = os.path.join(self.video_files[idx], "sub_clip.mp4")
+ cap = cv2.VideoCapture(video_path)
+ fps = cap.get(cv2.CAP_PROP_FPS)
+ try:
+ is_0_fps = 2 / fps
+ except Exception as e:
+ print(f"The fps of {video_path} is 0 !!!")
+ vocal_audio_path = os.path.join(self.video_files[idx], "audio.wav")
+ vocal_duration = librosa.get_duration(filename=vocal_audio_path)
+ frames_path = os.path.join(self.video_files[idx], "images")
+ total_frame_number = self.frame_count(frames_path)
+ fps = total_frame_number / vocal_duration
+ print(f"The calculated fps of {video_path} is {fps} !!!")
+ # idx = random.randint(0, len(self.video_files) - 1)
+ # video_path = os.path.join(self.video_files[idx], "sub_clip.mp4")
+ # cap = cv2.VideoCapture(video_path)
+ # fps = cap.get(cv2.CAP_PROP_FPS)
+
+ frames_path = os.path.join(self.video_files[idx], "images")
+
+ face_masks_path = os.path.join(self.video_files[idx], "face_masks")
+ lip_masks_path = os.path.join(self.video_files[idx], "lip_masks")
+ raw_audio_path = os.path.join(self.video_files[idx], "audio.wav")
+ # vocal_audio_path = os.path.join(self.video_files[idx], "vocal.wav")
+ vocal_audio_path = os.path.join(self.video_files[idx], "audio.wav")
+ video_length = self.frame_count(frames_path)
+ frames_list = self.find_frames_list(frames_path)
+
+ clip_length = min(video_length, (self.n_sample_frames - 1) * self.sample_frame_rate + 1)
+
+ start_idx = random.randint(0, video_length - clip_length)
+ batch_index = np.linspace(
+ start_idx, start_idx + clip_length - 1, self.n_sample_frames, dtype=int
+ ).tolist()
+ all_indices = list(range(0, video_length))
+ reference_frame_idx = random.choice(all_indices)
+
+ tgt_pil_image_list = []
+ tgt_face_masks_list = []
+ tgt_lip_masks_list = []
+
+ # reference_frame_path = os.path.join(frames_path, frames_list[reference_frame_idx])
+ reference_frame_path = os.path.join(frames_path, frames_list[start_idx])
+ reference_pil_image = Image.open(reference_frame_path).convert('RGB')
+ reference_pil_image = reference_pil_image.resize((self.width, self.height))
+ reference_pil_image = torch.from_numpy(np.array(reference_pil_image)).float()
+ reference_pil_image = reference_pil_image / 127.5 - 1
+
+ for index in batch_index:
+ tgt_img_path = os.path.join(frames_path, frames_list[index])
+ # file_name = os.path.splitext(os.path.basename(tgt_img_path))[0]
+ file_name = os.path.basename(tgt_img_path)
+ face_mask_path = os.path.join(face_masks_path, file_name)
+ lip_mask_path = os.path.join(lip_masks_path, file_name)
+ try:
+ tgt_img_pil = Image.open(tgt_img_path).convert('RGB')
+ except Exception as e:
+ print(f"Fail loading the image: {tgt_img_path}")
+
+ try:
+ tgt_lip_mask = Image.open(lip_mask_path)
+ # tgt_lip_mask = Image.open(lip_mask_path).convert('RGB')
+ tgt_lip_mask = tgt_lip_mask.resize((self.width, self.height))
+ tgt_lip_mask = torch.from_numpy(np.array(tgt_lip_mask)).float()
+ # tgt_lip_mask = tgt_lip_mask / 127.5 - 1
+ tgt_lip_mask = tgt_lip_mask / 255
+ except Exception as e:
+ print(f"Fail loading the lip masks: {lip_mask_path}")
+ tgt_lip_mask = torch.ones(self.height, self.width)
+ # tgt_lip_mask = torch.ones(self.height, self.width, 3)
+ tgt_lip_masks_list.append(tgt_lip_mask)
+
+ try:
+ tgt_face_mask = Image.open(face_mask_path)
+ # tgt_face_mask = Image.open(face_mask_path).convert('RGB')
+ tgt_face_mask = tgt_face_mask.resize((self.width, self.height))
+ tgt_face_mask = torch.from_numpy(np.array(tgt_face_mask)).float()
+ tgt_face_mask = tgt_face_mask / 255
+ # tgt_face_mask = tgt_face_mask / 127.5 - 1
+ except Exception as e:
+ print(f"Fail loading the face masks: {face_mask_path}")
+ tgt_face_mask = torch.ones(self.height, self.width)
+ # tgt_face_mask = torch.ones(self.height, self.width, 3)
+ tgt_face_masks_list.append(tgt_face_mask)
+
+ tgt_img_pil = tgt_img_pil.resize((self.width, self.height))
+ tgt_img_tensor = torch.from_numpy(np.array(tgt_img_pil)).float()
+ tgt_img_normalized = tgt_img_tensor / 127.5 - 1
+ tgt_pil_image_list.append(tgt_img_normalized)
+
+ sr = 16000
+ vocal_input, sample_rate = librosa.load(vocal_audio_path, sr=sr)
+ vocal_duration = librosa.get_duration(filename=vocal_audio_path)
+ start_time = batch_index[0] / fps
+ end_time = (clip_length / fps) + start_time
+ start_sample = int(start_time * sr)
+ end_sample = int(end_time * sr)
+ try:
+ vocal_segment = vocal_input[start_sample:end_sample]
+ except:
+ print(f"The current vocal segment is too short: {vocal_audio_path}, [{batch_index[0]}, {batch_index[-1]}], fps={fps}, clip_length={clip_length}, vocal_duration={vocal_duration}], [{start_time}, {end_time}]")
+ vocal_segment = vocal_input[start_sample:]
+ vocal_input_values = self.wav2vec_processor(
+ vocal_segment, sampling_rate=sample_rate, return_tensors="pt"
+ ).input_values
+
+
+ tgt_pil_image_list = torch.stack(tgt_pil_image_list, dim=0)
+ tgt_pil_image_list = rearrange(tgt_pil_image_list, "f h w c -> f c h w")
+ reference_pil_image = rearrange(reference_pil_image, "h w c -> c h w")
+
+ tgt_face_masks_list = torch.stack(tgt_face_masks_list, dim=0)
+ tgt_face_masks_list = torch.unsqueeze(tgt_face_masks_list, dim=-1)
+ tgt_face_masks_list = rearrange(tgt_face_masks_list, "f h w c -> c f h w")
+ tgt_lip_masks_list = torch.stack(tgt_lip_masks_list, dim=0)
+ tgt_lip_masks_list = torch.unsqueeze(tgt_lip_masks_list, dim=-1)
+ tgt_lip_masks_list = rearrange(tgt_lip_masks_list, "f h w c -> c f h w")
+
+
+ clip_pixel_values = reference_pil_image.permute(1, 2, 0).contiguous()
+ clip_pixel_values = (clip_pixel_values * 0.5 + 0.5) * 255
+
+ cos_similarities = []
+ stride = 8
+ for i in range(0, tgt_pil_image_list.size()[0] - stride, stride):
+ frame1 = tgt_pil_image_list[i]
+ frame2 = tgt_pil_image_list[i + stride]
+ frame1_flat = frame1.contiguous().view(-1)
+ frame2_flat = frame2.contiguous().view(-1)
+ cos_sim = F.cosine_similarity(frame1_flat, frame2_flat, dim=0)
+ cos_sim = (cos_sim + 1) / 2
+ cos_similarities.append(cos_sim.item())
+ overall_cos_sim = F.cosine_similarity(tgt_pil_image_list[0].contiguous().view(-1), tgt_pil_image_list[-1].contiguous().view(-1), dim=0)
+ overall_cos_sim = (overall_cos_sim + 1) / 2
+ cos_similarities.append(overall_cos_sim.item())
+ motion_id = (1.0 - sum(cos_similarities) / len(cos_similarities)) * 100
+
+
+ if "singing" in self.video_files[idx]:
+ text_prompt = "The protagonist is singing"
+ elif "speech" in self.video_files[idx]:
+ text_prompt = "The protagonist is talking"
+ elif "dancing" in self.video_files[idx]:
+ text_prompt = "The protagonist is simultaneously dancing and singing"
+ else:
+ text_prompt = ""
+ print(1 / 0)
+
+ sample = dict(
+ pixel_values=tgt_pil_image_list,
+ reference_image=reference_pil_image,
+ clip_pixel_values=clip_pixel_values,
+ tgt_face_masks=tgt_face_masks_list,
+ vocal_input_values=vocal_input_values,
+ text_prompt=text_prompt,
+ motion_id=motion_id,
+ tgt_lip_masks=tgt_lip_masks_list,
+ audio_path=raw_audio_path,
+ )
+
+ if self.enable_inpaint:
+ pixel_value_masks = get_random_mask(tgt_pil_image_list.size(), image_start_only=True)
+ masked_pixel_values = tgt_pil_image_list * (1-pixel_value_masks)
+ sample["masked_pixel_values"] = masked_pixel_values
+ sample["pixel_value_masks"] = pixel_value_masks
+
+
+ return sample
\ No newline at end of file
diff --git a/wan/dist/__init__.py b/wan/dist/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..8da6edfd6da71f09f47c2bdcbb9a09e47fd19711
--- /dev/null
+++ b/wan/dist/__init__.py
@@ -0,0 +1,40 @@
+import torch
+import torch.distributed as dist
+
+try:
+ import xfuser
+ from xfuser.core.distributed import (get_sequence_parallel_rank,
+ get_sequence_parallel_world_size,
+ get_sp_group, get_world_group,
+ init_distributed_environment,
+ initialize_model_parallel)
+ from xfuser.core.long_ctx_attention import xFuserLongContextAttention
+except Exception as ex:
+ get_sequence_parallel_world_size = None
+ get_sequence_parallel_rank = None
+ xFuserLongContextAttention = None
+ get_sp_group = None
+ get_world_group = None
+ init_distributed_environment = None
+ initialize_model_parallel = None
+
+def set_multi_gpus_devices(ulysses_degree, ring_degree):
+ if ulysses_degree > 1 or ring_degree > 1:
+ if get_sp_group is None:
+ raise RuntimeError("xfuser is not installed.")
+ dist.init_process_group("nccl")
+ print('parallel inference enabled: ulysses_degree=%d ring_degree=%d rank=%d world_size=%d' % (
+ ulysses_degree, ring_degree, dist.get_rank(),
+ dist.get_world_size()))
+ assert dist.get_world_size() == ring_degree * ulysses_degree, \
+ "number of GPUs(%d) should be equal to ring_degree * ulysses_degree." % dist.get_world_size()
+ init_distributed_environment(rank=dist.get_rank(), world_size=dist.get_world_size())
+ initialize_model_parallel(sequence_parallel_degree=dist.get_world_size(),
+ ring_degree=ring_degree,
+ ulysses_degree=ulysses_degree)
+ # device = torch.device("cuda:%d" % dist.get_rank())
+ device = torch.device(f"cuda:{get_world_group().local_rank}")
+ print('rank=%d device=%s' % (get_world_group().rank, str(device)))
+ else:
+ device = "cuda"
+ return device
\ No newline at end of file
diff --git a/wan/dist/__pycache__/__init__.cpython-311.pyc b/wan/dist/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..cbce5614de7fd2e11dc7e5983d571c06be50f9ae
Binary files /dev/null and b/wan/dist/__pycache__/__init__.cpython-311.pyc differ
diff --git a/wan/dist/__pycache__/wan_xfuser.cpython-311.pyc b/wan/dist/__pycache__/wan_xfuser.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..6399923a5cb5d038a848ae11fe7e3d484db867e4
Binary files /dev/null and b/wan/dist/__pycache__/wan_xfuser.cpython-311.pyc differ
diff --git a/wan/dist/wan_xfuser.py b/wan/dist/wan_xfuser.py
new file mode 100644
index 0000000000000000000000000000000000000000..9e02bf822325e91a57171f6c3f24d4df26588742
--- /dev/null
+++ b/wan/dist/wan_xfuser.py
@@ -0,0 +1,115 @@
+import torch
+import torch.amp as amp
+
+try:
+ import xfuser
+ from xfuser.core.distributed import (get_sequence_parallel_rank,
+ get_sequence_parallel_world_size,
+ get_sp_group,
+ init_distributed_environment,
+ initialize_model_parallel)
+ from xfuser.core.long_ctx_attention import xFuserLongContextAttention
+except Exception as ex:
+ get_sequence_parallel_world_size = None
+ get_sequence_parallel_rank = None
+ xFuserLongContextAttention = None
+ get_sp_group = None
+ init_distributed_environment = None
+ initialize_model_parallel = None
+
+def pad_freqs(original_tensor, target_len):
+ seq_len, s1, s2 = original_tensor.shape
+ pad_size = target_len - seq_len
+ padding_tensor = torch.ones(
+ pad_size,
+ s1,
+ s2,
+ dtype=original_tensor.dtype,
+ device=original_tensor.device)
+ padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0)
+ return padded_tensor
+
+@amp.autocast('cuda', enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+ """
+ x: [B, L, N, C].
+ grid_sizes: [B, 3].
+ freqs: [M, C // 2].
+ """
+ s, n, c = x.size(1), x.size(2), x.size(3) // 2
+ # split freqs
+ freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+ # loop over samples
+ output = []
+ for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+ seq_len = f * h * w
+
+ # precompute multipliers
+ x_i = torch.view_as_complex(x[i, :s].to(torch.float32).reshape(
+ s, n, -1, 2))
+ freqs_i = torch.cat([
+ freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+ freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+ freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+ ],
+ dim=-1).reshape(seq_len, 1, -1)
+
+ # apply rotary embedding
+ sp_size = get_sequence_parallel_world_size()
+ sp_rank = get_sequence_parallel_rank()
+ freqs_i = pad_freqs(freqs_i, s * sp_size)
+ s_per_rank = s
+ freqs_i_rank = freqs_i[(sp_rank * s_per_rank):((sp_rank + 1) *
+ s_per_rank), :, :]
+ x_i = torch.view_as_real(x_i * freqs_i_rank).flatten(2)
+ x_i = torch.cat([x_i, x[i, s:]])
+
+ # append to collection
+ output.append(x_i)
+ return torch.stack(output)
+
+def usp_attn_forward(self,
+ x,
+ seq_lens,
+ grid_sizes,
+ freqs,
+ dtype=torch.bfloat16):
+ b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+ half_dtypes = (torch.float16, torch.bfloat16)
+
+ def half(x):
+ return x if x.dtype in half_dtypes else x.to(dtype)
+
+ # query, key, value function
+ def qkv_fn(x):
+ q = self.norm_q(self.q(x)).view(b, s, n, d)
+ k = self.norm_k(self.k(x)).view(b, s, n, d)
+ v = self.v(x).view(b, s, n, d)
+ return q, k, v
+
+ q, k, v = qkv_fn(x)
+ q = rope_apply(q, grid_sizes, freqs)
+ k = rope_apply(k, grid_sizes, freqs)
+
+ # TODO: We should use unpaded q,k,v for attention.
+ # k_lens = seq_lens // get_sequence_parallel_world_size()
+ # if k_lens is not None:
+ # q = torch.cat([u[:l] for u, l in zip(q, k_lens)]).unsqueeze(0)
+ # k = torch.cat([u[:l] for u, l in zip(k, k_lens)]).unsqueeze(0)
+ # v = torch.cat([u[:l] for u, l in zip(v, k_lens)]).unsqueeze(0)
+
+ x = xFuserLongContextAttention()(
+ None,
+ query=half(q),
+ key=half(k),
+ value=half(v),
+ window_size=self.window_size)
+
+ # TODO: padding after attention.
+ # x = torch.cat([x, x.new_zeros(b, s - x.size(1), n, d)], dim=1)
+
+ # output
+ x = x.flatten(2)
+ x = self.o(x)
+ return x
\ No newline at end of file
diff --git a/wan/distributed/__init__.py b/wan/distributed/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/wan/distributed/__pycache__/__init__.cpython-311.pyc b/wan/distributed/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c5e295cca5d7988a3af7f2ecbfc721666c321df2
Binary files /dev/null and b/wan/distributed/__pycache__/__init__.cpython-311.pyc differ
diff --git a/wan/distributed/__pycache__/fsdp.cpython-311.pyc b/wan/distributed/__pycache__/fsdp.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..76122f94c08f75d8ae7430d4cffb3608ea630777
Binary files /dev/null and b/wan/distributed/__pycache__/fsdp.cpython-311.pyc differ
diff --git a/wan/distributed/fsdp.py b/wan/distributed/fsdp.py
new file mode 100644
index 0000000000000000000000000000000000000000..18ba2f3eb77b5df0dafdb3d967b66b999092ac83
--- /dev/null
+++ b/wan/distributed/fsdp.py
@@ -0,0 +1,41 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import gc
+from functools import partial
+
+import torch
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.distributed.fsdp import MixedPrecision, ShardingStrategy
+from torch.distributed.fsdp.wrap import lambda_auto_wrap_policy
+from torch.distributed.utils import _free_storage
+
+def shard_model(
+ model,
+ device_id,
+ param_dtype=torch.bfloat16,
+ reduce_dtype=torch.float32,
+ buffer_dtype=torch.float32,
+ process_group=None,
+ sharding_strategy=ShardingStrategy.FULL_SHARD,
+ sync_module_states=True,
+):
+ model = FSDP(
+ module=model,
+ process_group=process_group,
+ sharding_strategy=sharding_strategy,
+ auto_wrap_policy=partial(
+ lambda_auto_wrap_policy, lambda_fn=lambda m: m in model.blocks),
+ mixed_precision=MixedPrecision(
+ param_dtype=param_dtype,
+ reduce_dtype=reduce_dtype,
+ buffer_dtype=buffer_dtype),
+ device_id=device_id,
+ sync_module_states=sync_module_states)
+ return model
+
+def free_model(model):
+ for m in model.modules():
+ if isinstance(m, FSDP):
+ _free_storage(m._handle.flat_param.data)
+ del model
+ gc.collect()
+ torch.cuda.empty_cache()
diff --git a/wan/distributed/xdit_context_parallel.py b/wan/distributed/xdit_context_parallel.py
new file mode 100644
index 0000000000000000000000000000000000000000..01936cee9c31ce0af57af21af1310d69303390e0
--- /dev/null
+++ b/wan/distributed/xdit_context_parallel.py
@@ -0,0 +1,192 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+import torch.cuda.amp as amp
+from xfuser.core.distributed import (get_sequence_parallel_rank,
+ get_sequence_parallel_world_size,
+ get_sp_group)
+from xfuser.core.long_ctx_attention import xFuserLongContextAttention
+
+from ..modules.model import sinusoidal_embedding_1d
+
+
+def pad_freqs(original_tensor, target_len):
+ seq_len, s1, s2 = original_tensor.shape
+ pad_size = target_len - seq_len
+ padding_tensor = torch.ones(
+ pad_size,
+ s1,
+ s2,
+ dtype=original_tensor.dtype,
+ device=original_tensor.device)
+ padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0)
+ return padded_tensor
+
+
+@amp.autocast(enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+ """
+ x: [B, L, N, C].
+ grid_sizes: [B, 3].
+ freqs: [M, C // 2].
+ """
+ s, n, c = x.size(1), x.size(2), x.size(3) // 2
+ # split freqs
+ freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+ # loop over samples
+ output = []
+ for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+ seq_len = f * h * w
+
+ # precompute multipliers
+ x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(
+ s, n, -1, 2))
+ freqs_i = torch.cat([
+ freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+ freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+ freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+ ],
+ dim=-1).reshape(seq_len, 1, -1)
+
+ # apply rotary embedding
+ sp_size = get_sequence_parallel_world_size()
+ sp_rank = get_sequence_parallel_rank()
+ freqs_i = pad_freqs(freqs_i, s * sp_size)
+ s_per_rank = s
+ freqs_i_rank = freqs_i[(sp_rank * s_per_rank):((sp_rank + 1) *
+ s_per_rank), :, :]
+ x_i = torch.view_as_real(x_i * freqs_i_rank).flatten(2)
+ x_i = torch.cat([x_i, x[i, s:]])
+
+ # append to collection
+ output.append(x_i)
+ return torch.stack(output).float()
+
+
+def usp_dit_forward(
+ self,
+ x,
+ t,
+ context,
+ seq_len,
+ clip_fea=None,
+ y=None,
+):
+ """
+ x: A list of videos each with shape [C, T, H, W].
+ t: [B].
+ context: A list of text embeddings each with shape [L, C].
+ """
+ if self.model_type == 'i2v':
+ assert clip_fea is not None and y is not None
+ # params
+ device = self.patch_embedding.weight.device
+ if self.freqs.device != device:
+ self.freqs = self.freqs.to(device)
+
+ if y is not None:
+ x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+ # embeddings
+ x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+ grid_sizes = torch.stack(
+ [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
+ x = [u.flatten(2).transpose(1, 2) for u in x]
+ seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
+ assert seq_lens.max() <= seq_len
+ x = torch.cat([
+ torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)
+ for u in x
+ ])
+
+ # time embeddings
+ with amp.autocast(dtype=torch.float32):
+ e = self.time_embedding(
+ sinusoidal_embedding_1d(self.freq_dim, t).float())
+ e0 = self.time_projection(e).unflatten(1, (6, self.dim))
+ assert e.dtype == torch.float32 and e0.dtype == torch.float32
+
+ # context
+ context_lens = None
+ context = self.text_embedding(
+ torch.stack([
+ torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
+ for u in context
+ ]))
+
+ if clip_fea is not None:
+ context_clip = self.img_emb(clip_fea) # bs x 257 x dim
+ context = torch.concat([context_clip, context], dim=1)
+
+ # arguments
+ kwargs = dict(
+ e=e0,
+ seq_lens=seq_lens,
+ grid_sizes=grid_sizes,
+ freqs=self.freqs,
+ context=context,
+ context_lens=context_lens)
+
+ # Context Parallel
+ x = torch.chunk(
+ x, get_sequence_parallel_world_size(),
+ dim=1)[get_sequence_parallel_rank()]
+
+ for block in self.blocks:
+ x = block(x, **kwargs)
+
+ # head
+ x = self.head(x, e)
+
+ # Context Parallel
+ x = get_sp_group().all_gather(x, dim=1)
+
+ # unpatchify
+ x = self.unpatchify(x, grid_sizes)
+ return [u.float() for u in x]
+
+
+def usp_attn_forward(self,
+ x,
+ seq_lens,
+ grid_sizes,
+ freqs,
+ dtype=torch.bfloat16):
+ b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
+ half_dtypes = (torch.float16, torch.bfloat16)
+
+ def half(x):
+ return x if x.dtype in half_dtypes else x.to(dtype)
+
+ # query, key, value function
+ def qkv_fn(x):
+ q = self.norm_q(self.q(x)).view(b, s, n, d)
+ k = self.norm_k(self.k(x)).view(b, s, n, d)
+ v = self.v(x).view(b, s, n, d)
+ return q, k, v
+
+ q, k, v = qkv_fn(x)
+ q = rope_apply(q, grid_sizes, freqs)
+ k = rope_apply(k, grid_sizes, freqs)
+
+ # TODO: We should use unpaded q,k,v for attention.
+ # k_lens = seq_lens // get_sequence_parallel_world_size()
+ # if k_lens is not None:
+ # q = torch.cat([u[:l] for u, l in zip(q, k_lens)]).unsqueeze(0)
+ # k = torch.cat([u[:l] for u, l in zip(k, k_lens)]).unsqueeze(0)
+ # v = torch.cat([u[:l] for u, l in zip(v, k_lens)]).unsqueeze(0)
+
+ x = xFuserLongContextAttention()(
+ None,
+ query=half(q),
+ key=half(k),
+ value=half(v),
+ window_size=self.window_size)
+
+ # TODO: padding after attention.
+ # x = torch.cat([x, x.new_zeros(b, s - x.size(1), n, d)], dim=1)
+
+ # output
+ x = x.flatten(2)
+ x = self.o(x)
+ return x
diff --git a/wan/image2video.py b/wan/image2video.py
new file mode 100644
index 0000000000000000000000000000000000000000..b375fb96fb3f58353d845a3ecd7f863bce7c2c70
--- /dev/null
+++ b/wan/image2video.py
@@ -0,0 +1,334 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import gc
+import logging
+import math
+import os
+import random
+import sys
+import types
+from contextlib import contextmanager
+from functools import partial
+
+import numpy as np
+import torch
+import torch.cuda.amp as amp
+import torch.distributed as dist
+import torchvision.transforms.functional as TF
+from tqdm import tqdm
+
+from .distributed.fsdp import shard_model
+from .modules.clip import CLIPModel
+from .modules.model import WanModel
+from .modules.t5 import T5EncoderModel
+from .modules.vae import WanVAE
+from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler,
+ get_sampling_sigmas, retrieve_timesteps)
+from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+
+
+class WanI2V:
+
+ def __init__(
+ self,
+ config,
+ checkpoint_dir,
+ device_id=0,
+ rank=0,
+ t5_fsdp=False,
+ dit_fsdp=False,
+ use_usp=False,
+ t5_cpu=False,
+ init_on_cpu=True,
+ ):
+ r"""
+ Initializes the image-to-video generation model components.
+
+ Args:
+ config (EasyDict):
+ Object containing model parameters initialized from config.py
+ checkpoint_dir (`str`):
+ Path to directory containing model checkpoints
+ device_id (`int`, *optional*, defaults to 0):
+ Id of target GPU device
+ rank (`int`, *optional*, defaults to 0):
+ Process rank for distributed training
+ t5_fsdp (`bool`, *optional*, defaults to False):
+ Enable FSDP sharding for T5 model
+ dit_fsdp (`bool`, *optional*, defaults to False):
+ Enable FSDP sharding for DiT model
+ use_usp (`bool`, *optional*, defaults to False):
+ Enable distribution strategy of USP.
+ t5_cpu (`bool`, *optional*, defaults to False):
+ Whether to place T5 model on CPU. Only works without t5_fsdp.
+ init_on_cpu (`bool`, *optional*, defaults to True):
+ Enable initializing Transformer Model on CPU. Only works without FSDP or USP.
+ """
+ self.device = torch.device(f"cuda:{device_id}")
+ self.config = config
+ self.rank = rank
+ self.use_usp = use_usp
+ self.t5_cpu = t5_cpu
+
+ self.num_train_timesteps = config.num_train_timesteps
+ self.param_dtype = config.param_dtype
+
+ shard_fn = partial(shard_model, device_id=device_id)
+ self.text_encoder = T5EncoderModel(
+ text_len=config.text_len,
+ dtype=config.t5_dtype,
+ device=torch.device('cpu'),
+ checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),
+ tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),
+ shard_fn=shard_fn if t5_fsdp else None,
+ )
+
+ self.vae_stride = config.vae_stride
+ self.patch_size = config.patch_size
+ self.vae = WanVAE(
+ vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
+ device=self.device)
+
+ self.clip = CLIPModel(
+ dtype=config.clip_dtype,
+ device=self.device,
+ checkpoint_path=os.path.join(checkpoint_dir,
+ config.clip_checkpoint),
+ tokenizer_path=os.path.join(checkpoint_dir, config.clip_tokenizer))
+
+ logging.info(f"Creating WanModel from {checkpoint_dir}")
+ self.model = WanModel.from_pretrained(checkpoint_dir)
+ self.model.eval().requires_grad_(False)
+
+ if t5_fsdp or dit_fsdp or use_usp:
+ init_on_cpu = False
+
+ if use_usp:
+ from xfuser.core.distributed import \
+ get_sequence_parallel_world_size
+
+ from .distributed.xdit_context_parallel import (usp_attn_forward,
+ usp_dit_forward)
+ for block in self.model.blocks:
+ block.self_attn.forward = types.MethodType(usp_attn_forward, block.self_attn)
+ self.model.forward = types.MethodType(usp_dit_forward, self.model)
+ self.sp_size = get_sequence_parallel_world_size()
+ else:
+ self.sp_size = 1
+
+ if dist.is_initialized():
+ dist.barrier()
+ if dit_fsdp:
+ self.model = shard_fn(self.model)
+ else:
+ if not init_on_cpu:
+ self.model.to(self.device)
+
+ self.sample_neg_prompt = config.sample_neg_prompt
+
+ def generate(self,
+ input_prompt,
+ img,
+ max_area=720 * 1280,
+ frame_num=81,
+ shift=5.0,
+ sample_solver='unipc',
+ sampling_steps=40,
+ guide_scale=5.0,
+ n_prompt="",
+ seed=-1,
+ offload_model=True):
+ r"""
+ Generates video frames from input image and text prompt using diffusion process.
+
+ Args:
+ input_prompt (`str`):
+ Text prompt for content generation.
+ img (PIL.Image.Image):
+ Input image tensor. Shape: [3, H, W]
+ max_area (`int`, *optional*, defaults to 720*1280):
+ Maximum pixel area for latent space calculation. Controls video resolution scaling
+ frame_num (`int`, *optional*, defaults to 81):
+ How many frames to sample from a video. The number should be 4n+1
+ shift (`float`, *optional*, defaults to 5.0):
+ Noise schedule shift parameter. Affects temporal dynamics
+ [NOTE]: If you want to generate a 480p video, it is recommended to set the shift value to 3.0.
+ sample_solver (`str`, *optional*, defaults to 'unipc'):
+ Solver used to sample the video.
+ sampling_steps (`int`, *optional*, defaults to 40):
+ Number of diffusion sampling steps. Higher values improve quality but slow generation
+ guide_scale (`float`, *optional*, defaults 5.0):
+ Classifier-free guidance scale. Controls prompt adherence vs. creativity
+ n_prompt (`str`, *optional*, defaults to ""):
+ Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt`
+ seed (`int`, *optional*, defaults to -1):
+ Random seed for noise generation. If -1, use random seed
+ offload_model (`bool`, *optional*, defaults to True):
+ If True, offloads models to CPU during generation to save VRAM
+
+ Returns:
+ torch.Tensor:
+ Generated video frames tensor. Dimensions: (C, N H, W) where:
+ - C: Color channels (3 for RGB)
+ - N: Number of frames (81)
+ - H: Frame height (from max_area)
+ - W: Frame width from max_area)
+ """
+ img = TF.to_tensor(img).sub_(0.5).div_(0.5).to(self.device)
+
+ F = frame_num
+ h, w = img.shape[1:]
+ aspect_ratio = h / w
+ lat_h = round(
+ np.sqrt(max_area * aspect_ratio) // self.vae_stride[1] //
+ self.patch_size[1] * self.patch_size[1])
+ lat_w = round(
+ np.sqrt(max_area / aspect_ratio) // self.vae_stride[2] //
+ self.patch_size[2] * self.patch_size[2])
+ h = lat_h * self.vae_stride[1]
+ w = lat_w * self.vae_stride[2]
+
+ max_seq_len = ((F - 1) // self.vae_stride[0] + 1) * lat_h * lat_w // (self.patch_size[1] * self.patch_size[2])
+ max_seq_len = int(math.ceil(max_seq_len / self.sp_size)) * self.sp_size
+
+ seed = seed if seed >= 0 else random.randint(0, sys.maxsize)
+ seed_g = torch.Generator(device=self.device)
+ seed_g.manual_seed(seed)
+ noise = torch.randn(
+ 16,
+ 21,
+ lat_h,
+ lat_w,
+ dtype=torch.float32,
+ generator=seed_g,
+ device=self.device)
+
+ msk = torch.ones(1, 81, lat_h, lat_w, device=self.device)
+ msk[:, 1:] = 0
+ msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]],dim=1)
+ msk = msk.view(1, msk.shape[1] // 4, 4, lat_h, lat_w)
+ msk = msk.transpose(1, 2)[0]
+
+ if n_prompt == "":
+ n_prompt = self.sample_neg_prompt
+
+ # preprocess
+ if not self.t5_cpu:
+ self.text_encoder.model.to(self.device)
+ context = self.text_encoder([input_prompt], self.device)
+ context_null = self.text_encoder([n_prompt], self.device)
+ if offload_model:
+ self.text_encoder.model.cpu()
+ else:
+ context = self.text_encoder([input_prompt], torch.device('cpu'))
+ context_null = self.text_encoder([n_prompt], torch.device('cpu'))
+ context = [t.to(self.device) for t in context]
+ context_null = [t.to(self.device) for t in context_null]
+
+ self.clip.model.to(self.device)
+ clip_context = self.clip.visual([img[:, None, :, :]])
+ if offload_model:
+ self.clip.model.cpu()
+
+ y = self.vae.encode([torch.concat([torch.nn.functional.interpolate(img[None].cpu(), size=(h, w), mode='bicubic').transpose(0, 1), torch.zeros(3, 80, h, w)],dim=1).to(self.device)])[0]
+ y = torch.concat([msk, y])
+
+ @contextmanager
+ def noop_no_sync():
+ yield
+
+ no_sync = getattr(self.model, 'no_sync', noop_no_sync)
+
+ # evaluation mode
+ with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync():
+
+ if sample_solver == 'unipc':
+ sample_scheduler = FlowUniPCMultistepScheduler(
+ num_train_timesteps=self.num_train_timesteps,
+ shift=1,
+ use_dynamic_shifting=False)
+ sample_scheduler.set_timesteps(
+ sampling_steps, device=self.device, shift=shift)
+ timesteps = sample_scheduler.timesteps
+ elif sample_solver == 'dpm++':
+ sample_scheduler = FlowDPMSolverMultistepScheduler(
+ num_train_timesteps=self.num_train_timesteps,
+ shift=1,
+ use_dynamic_shifting=False)
+ sampling_sigmas = get_sampling_sigmas(sampling_steps, shift)
+ timesteps, _ = retrieve_timesteps(
+ sample_scheduler,
+ device=self.device,
+ sigmas=sampling_sigmas)
+ else:
+ raise NotImplementedError("Unsupported solver.")
+
+ # sample videos
+ latent = noise
+
+ arg_c = {
+ 'context': [context[0]],
+ 'clip_fea': clip_context,
+ 'seq_len': max_seq_len,
+ 'y': [y],
+ }
+
+ arg_null = {
+ 'context': context_null,
+ 'clip_fea': clip_context,
+ 'seq_len': max_seq_len,
+ 'y': [y],
+ }
+
+ if offload_model:
+ torch.cuda.empty_cache()
+
+ self.model.to(self.device)
+ for _, t in enumerate(tqdm(timesteps)):
+ latent_model_input = [latent.to(self.device)]
+ timestep = [t]
+
+ timestep = torch.stack(timestep).to(self.device)
+
+ noise_pred_cond = self.model(
+ latent_model_input, t=timestep, **arg_c)[0].to(
+ torch.device('cpu') if offload_model else self.device)
+ if offload_model:
+ torch.cuda.empty_cache()
+ noise_pred_uncond = self.model(
+ latent_model_input, t=timestep, **arg_null)[0].to(
+ torch.device('cpu') if offload_model else self.device)
+ if offload_model:
+ torch.cuda.empty_cache()
+ noise_pred = noise_pred_uncond + guide_scale * (
+ noise_pred_cond - noise_pred_uncond)
+
+ latent = latent.to(
+ torch.device('cpu') if offload_model else self.device)
+
+ temp_x0 = sample_scheduler.step(
+ noise_pred.unsqueeze(0),
+ t,
+ latent.unsqueeze(0),
+ return_dict=False,
+ generator=seed_g)[0]
+ latent = temp_x0.squeeze(0)
+
+ x0 = [latent.to(self.device)]
+ del latent_model_input, timestep
+
+ if offload_model:
+ self.model.cpu()
+ torch.cuda.empty_cache()
+
+ if self.rank == 0:
+ videos = self.vae.decode(x0)
+
+ del noise, latent
+ del sample_scheduler
+ if offload_model:
+ gc.collect()
+ torch.cuda.synchronize()
+ if dist.is_initialized():
+ dist.barrier()
+
+ return videos[0] if self.rank == 0 else None
diff --git a/wan/models/__init__.py b/wan/models/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/wan/models/__pycache__/__init__.cpython-311.pyc b/wan/models/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..ced6f2b54e0b2b09f29ae4562418eb903868ffca
Binary files /dev/null and b/wan/models/__pycache__/__init__.cpython-311.pyc differ
diff --git a/wan/models/__pycache__/cache_utils.cpython-311.pyc b/wan/models/__pycache__/cache_utils.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..e4e65783434c40ace30cd82b7c73ab7efa7d2422
Binary files /dev/null and b/wan/models/__pycache__/cache_utils.cpython-311.pyc differ
diff --git a/wan/models/__pycache__/vocal_projector_fantasy.cpython-311.pyc b/wan/models/__pycache__/vocal_projector_fantasy.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..a179da7d68333658457f350db76e3bb5a16b669b
Binary files /dev/null and b/wan/models/__pycache__/vocal_projector_fantasy.cpython-311.pyc differ
diff --git a/wan/models/__pycache__/vocal_projector_fantasy_1B.cpython-311.pyc b/wan/models/__pycache__/vocal_projector_fantasy_1B.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..b8de190a0ad9cedf6852c9fae6aff948b1e7eaa6
Binary files /dev/null and b/wan/models/__pycache__/vocal_projector_fantasy_1B.cpython-311.pyc differ
diff --git a/wan/models/__pycache__/wan_fantasy_transformer3d_1B.cpython-311.pyc b/wan/models/__pycache__/wan_fantasy_transformer3d_1B.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c9824dad8e91560fff7b0bfe6452d8c543503abb
Binary files /dev/null and b/wan/models/__pycache__/wan_fantasy_transformer3d_1B.cpython-311.pyc differ
diff --git a/wan/models/__pycache__/wan_image_encoder.cpython-311.pyc b/wan/models/__pycache__/wan_image_encoder.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..822f6d280a48a6ad9aab92021836adf9f5fba612
Binary files /dev/null and b/wan/models/__pycache__/wan_image_encoder.cpython-311.pyc differ
diff --git a/wan/models/__pycache__/wan_text_encoder.cpython-311.pyc b/wan/models/__pycache__/wan_text_encoder.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..9431c093bc2b540b7ee7099a19b63c25cd1eddb2
Binary files /dev/null and b/wan/models/__pycache__/wan_text_encoder.cpython-311.pyc differ
diff --git a/wan/models/__pycache__/wan_transformer3d.cpython-311.pyc b/wan/models/__pycache__/wan_transformer3d.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..bd5aa02eeb20713059a18fcdab3cb376b69945c3
Binary files /dev/null and b/wan/models/__pycache__/wan_transformer3d.cpython-311.pyc differ
diff --git a/wan/models/__pycache__/wan_vae.cpython-311.pyc b/wan/models/__pycache__/wan_vae.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..478140eed74a8c4fe4f68cb159571a23b79d83e9
Binary files /dev/null and b/wan/models/__pycache__/wan_vae.cpython-311.pyc differ
diff --git a/wan/models/__pycache__/wan_xlm_roberta.cpython-311.pyc b/wan/models/__pycache__/wan_xlm_roberta.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..77f7e074d8d0c082eec55990c751bfd9f6e83240
Binary files /dev/null and b/wan/models/__pycache__/wan_xlm_roberta.cpython-311.pyc differ
diff --git a/wan/models/attention_processor.py b/wan/models/attention_processor.py
new file mode 100644
index 0000000000000000000000000000000000000000..23029b1e67e16871540c8ee75777d793722c965e
--- /dev/null
+++ b/wan/models/attention_processor.py
@@ -0,0 +1,6228 @@
+import inspect
+import math
+from typing import Callable, List, Optional, Tuple, Union
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from diffusers.image_processor import IPAdapterMaskProcessor
+from diffusers.utils import deprecate, is_torch_xla_available, logging
+from diffusers.utils import is_torch_npu_available, is_xformers_available
+from diffusers.utils import is_torch_version
+
+try:
+ from torch._dynamo import allow_in_graph as maybe_allow_in_graph
+except (ImportError, ModuleNotFoundError):
+ def maybe_allow_in_graph(cls):
+ return cls
+
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+if is_torch_npu_available():
+ import torch_npu
+
+if is_xformers_available():
+ import xformers
+ import xformers.ops
+else:
+ xformers = None
+
+if is_torch_xla_available():
+ # flash attention pallas kernel is introduced in the torch_xla 2.3 release.
+ if is_torch_xla_version(">", "2.2"):
+ from torch_xla.experimental.custom_kernel import flash_attention
+ from torch_xla.runtime import is_spmd
+ XLA_AVAILABLE = True
+else:
+ XLA_AVAILABLE = False
+
+
+@maybe_allow_in_graph
+class Attention(nn.Module):
+ r"""
+ A cross attention layer.
+
+ Parameters:
+ query_dim (`int`):
+ The number of channels in the query.
+ cross_attention_dim (`int`, *optional*):
+ The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`.
+ heads (`int`, *optional*, defaults to 8):
+ The number of heads to use for multi-head attention.
+ kv_heads (`int`, *optional*, defaults to `None`):
+ The number of key and value heads to use for multi-head attention. Defaults to `heads`. If
+ `kv_heads=heads`, the model will use Multi Head Attention (MHA), if `kv_heads=1` the model will use Multi
+ Query Attention (MQA) otherwise GQA is used.
+ dim_head (`int`, *optional*, defaults to 64):
+ The number of channels in each head.
+ dropout (`float`, *optional*, defaults to 0.0):
+ The dropout probability to use.
+ bias (`bool`, *optional*, defaults to False):
+ Set to `True` for the query, key, and value linear layers to contain a bias parameter.
+ upcast_attention (`bool`, *optional*, defaults to False):
+ Set to `True` to upcast the attention computation to `float32`.
+ upcast_softmax (`bool`, *optional*, defaults to False):
+ Set to `True` to upcast the softmax computation to `float32`.
+ cross_attention_norm (`str`, *optional*, defaults to `None`):
+ The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`.
+ cross_attention_norm_num_groups (`int`, *optional*, defaults to 32):
+ The number of groups to use for the group norm in the cross attention.
+ added_kv_proj_dim (`int`, *optional*, defaults to `None`):
+ The number of channels to use for the added key and value projections. If `None`, no projection is used.
+ norm_num_groups (`int`, *optional*, defaults to `None`):
+ The number of groups to use for the group norm in the attention.
+ spatial_norm_dim (`int`, *optional*, defaults to `None`):
+ The number of channels to use for the spatial normalization.
+ out_bias (`bool`, *optional*, defaults to `True`):
+ Set to `True` to use a bias in the output linear layer.
+ scale_qk (`bool`, *optional*, defaults to `True`):
+ Set to `True` to scale the query and key by `1 / sqrt(dim_head)`.
+ only_cross_attention (`bool`, *optional*, defaults to `False`):
+ Set to `True` to only use cross attention and not added_kv_proj_dim. Can only be set to `True` if
+ `added_kv_proj_dim` is not `None`.
+ eps (`float`, *optional*, defaults to 1e-5):
+ An additional value added to the denominator in group normalization that is used for numerical stability.
+ rescale_output_factor (`float`, *optional*, defaults to 1.0):
+ A factor to rescale the output by dividing it with this value.
+ residual_connection (`bool`, *optional*, defaults to `False`):
+ Set to `True` to add the residual connection to the output.
+ _from_deprecated_attn_block (`bool`, *optional*, defaults to `False`):
+ Set to `True` if the attention block is loaded from a deprecated state dict.
+ processor (`AttnProcessor`, *optional*, defaults to `None`):
+ The attention processor to use. If `None`, defaults to `AttnProcessor2_0` if `torch 2.x` is used and
+ `AttnProcessor` otherwise.
+ """
+
+ def __init__(
+ self,
+ query_dim: int,
+ cross_attention_dim: Optional[int] = None,
+ heads: int = 8,
+ kv_heads: Optional[int] = None,
+ dim_head: int = 64,
+ dropout: float = 0.0,
+ bias: bool = False,
+ upcast_attention: bool = False,
+ upcast_softmax: bool = False,
+ cross_attention_norm: Optional[str] = None,
+ cross_attention_norm_num_groups: int = 32,
+ qk_norm: Optional[str] = None,
+ added_kv_proj_dim: Optional[int] = None,
+ added_proj_bias: Optional[bool] = True,
+ norm_num_groups: Optional[int] = None,
+ spatial_norm_dim: Optional[int] = None,
+ out_bias: bool = True,
+ scale_qk: bool = True,
+ only_cross_attention: bool = False,
+ eps: float = 1e-5,
+ rescale_output_factor: float = 1.0,
+ residual_connection: bool = False,
+ _from_deprecated_attn_block: bool = False,
+ processor: Optional["AttnProcessor"] = None,
+ out_dim: int = None,
+ out_context_dim: int = None,
+ context_pre_only=None,
+ pre_only=False,
+ elementwise_affine: bool = True,
+ is_causal: bool = False,
+ ):
+ super().__init__()
+
+ # To prevent circular import.
+ from .normalization import FP32LayerNorm, LpNorm, RMSNorm
+
+ self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+ self.inner_kv_dim = self.inner_dim if kv_heads is None else dim_head * kv_heads
+ self.query_dim = query_dim
+ self.use_bias = bias
+ self.is_cross_attention = cross_attention_dim is not None
+ self.cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
+ self.upcast_attention = upcast_attention
+ self.upcast_softmax = upcast_softmax
+ self.rescale_output_factor = rescale_output_factor
+ self.residual_connection = residual_connection
+ self.dropout = dropout
+ self.fused_projections = False
+ self.out_dim = out_dim if out_dim is not None else query_dim
+ self.out_context_dim = out_context_dim if out_context_dim is not None else query_dim
+ self.context_pre_only = context_pre_only
+ self.pre_only = pre_only
+ self.is_causal = is_causal
+
+ # we make use of this private variable to know whether this class is loaded
+ # with an deprecated state dict so that we can convert it on the fly
+ self._from_deprecated_attn_block = _from_deprecated_attn_block
+
+ self.scale_qk = scale_qk
+ self.scale = dim_head**-0.5 if self.scale_qk else 1.0
+
+ self.heads = out_dim // dim_head if out_dim is not None else heads
+ # for slice_size > 0 the attention score computation
+ # is split across the batch axis to save memory
+ # You can set slice_size with `set_attention_slice`
+ self.sliceable_head_dim = heads
+
+ self.added_kv_proj_dim = added_kv_proj_dim
+ self.only_cross_attention = only_cross_attention
+
+ if self.added_kv_proj_dim is None and self.only_cross_attention:
+ raise ValueError(
+ "`only_cross_attention` can only be set to True if `added_kv_proj_dim` is not None. Make sure to set either `only_cross_attention=False` or define `added_kv_proj_dim`."
+ )
+
+ if norm_num_groups is not None:
+ self.group_norm = nn.GroupNorm(num_channels=query_dim, num_groups=norm_num_groups, eps=eps, affine=True)
+ else:
+ self.group_norm = None
+
+ if spatial_norm_dim is not None:
+ self.spatial_norm = SpatialNorm(f_channels=query_dim, zq_channels=spatial_norm_dim)
+ else:
+ self.spatial_norm = None
+
+ if qk_norm is None:
+ self.norm_q = None
+ self.norm_k = None
+ elif qk_norm == "layer_norm":
+ self.norm_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ elif qk_norm == "fp32_layer_norm":
+ self.norm_q = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)
+ self.norm_k = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)
+ elif qk_norm == "layer_norm_across_heads":
+ # Lumina applies qk norm across all heads
+ self.norm_q = nn.LayerNorm(dim_head * heads, eps=eps)
+ self.norm_k = nn.LayerNorm(dim_head * kv_heads, eps=eps)
+ elif qk_norm == "rms_norm":
+ self.norm_q = RMSNorm(dim_head, eps=eps)
+ self.norm_k = RMSNorm(dim_head, eps=eps)
+ elif qk_norm == "rms_norm_across_heads":
+ # LTX applies qk norm across all heads
+ self.norm_q = RMSNorm(dim_head * heads, eps=eps)
+ self.norm_k = RMSNorm(dim_head * kv_heads, eps=eps)
+ elif qk_norm == "l2":
+ self.norm_q = LpNorm(p=2, dim=-1, eps=eps)
+ self.norm_k = LpNorm(p=2, dim=-1, eps=eps)
+ else:
+ raise ValueError(
+ f"unknown qk_norm: {qk_norm}. Should be one of None, 'layer_norm', 'fp32_layer_norm', 'layer_norm_across_heads', 'rms_norm', 'rms_norm_across_heads', 'l2'."
+ )
+
+ if cross_attention_norm is None:
+ self.norm_cross = None
+ elif cross_attention_norm == "layer_norm":
+ self.norm_cross = nn.LayerNorm(self.cross_attention_dim)
+ elif cross_attention_norm == "group_norm":
+ if self.added_kv_proj_dim is not None:
+ # The given `encoder_hidden_states` are initially of shape
+ # (batch_size, seq_len, added_kv_proj_dim) before being projected
+ # to (batch_size, seq_len, cross_attention_dim). The norm is applied
+ # before the projection, so we need to use `added_kv_proj_dim` as
+ # the number of channels for the group norm.
+ norm_cross_num_channels = added_kv_proj_dim
+ else:
+ norm_cross_num_channels = self.cross_attention_dim
+
+ self.norm_cross = nn.GroupNorm(
+ num_channels=norm_cross_num_channels, num_groups=cross_attention_norm_num_groups, eps=1e-5, affine=True
+ )
+ else:
+ raise ValueError(
+ f"unknown cross_attention_norm: {cross_attention_norm}. Should be None, 'layer_norm' or 'group_norm'"
+ )
+
+ self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
+
+ if not self.only_cross_attention:
+ # only relevant for the `AddedKVProcessor` classes
+ self.to_k = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+ self.to_v = nn.Linear(self.cross_attention_dim, self.inner_kv_dim, bias=bias)
+ else:
+ self.to_k = None
+ self.to_v = None
+
+ self.added_proj_bias = added_proj_bias
+ if self.added_kv_proj_dim is not None:
+ self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias)
+ self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_kv_dim, bias=added_proj_bias)
+ if self.context_pre_only is not None:
+ self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ else:
+ self.add_q_proj = None
+ self.add_k_proj = None
+ self.add_v_proj = None
+
+ if not self.pre_only:
+ self.to_out = nn.ModuleList([])
+ self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+ self.to_out.append(nn.Dropout(dropout))
+ else:
+ self.to_out = None
+
+ if self.context_pre_only is not None and not self.context_pre_only:
+ self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias)
+ else:
+ self.to_add_out = None
+
+ if qk_norm is not None and added_kv_proj_dim is not None:
+ if qk_norm == "layer_norm":
+ self.norm_added_q = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ self.norm_added_k = nn.LayerNorm(dim_head, eps=eps, elementwise_affine=elementwise_affine)
+ elif qk_norm == "fp32_layer_norm":
+ self.norm_added_q = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)
+ self.norm_added_k = FP32LayerNorm(dim_head, elementwise_affine=False, bias=False, eps=eps)
+ elif qk_norm == "rms_norm":
+ self.norm_added_q = RMSNorm(dim_head, eps=eps)
+ self.norm_added_k = RMSNorm(dim_head, eps=eps)
+ elif qk_norm == "rms_norm_across_heads":
+ # Wan applies qk norm across all heads
+ # Wan also doesn't apply a q norm
+ self.norm_added_q = None
+ self.norm_added_k = RMSNorm(dim_head * kv_heads, eps=eps)
+ else:
+ raise ValueError(
+ f"unknown qk_norm: {qk_norm}. Should be one of `None,'layer_norm','fp32_layer_norm','rms_norm'`"
+ )
+ else:
+ self.norm_added_q = None
+ self.norm_added_k = None
+
+ # set attention processor
+ # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+ # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+ # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+ if processor is None:
+ processor = (
+ AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
+ )
+ self.set_processor(processor)
+
+ def set_use_xla_flash_attention(
+ self,
+ use_xla_flash_attention: bool,
+ partition_spec: Optional[Tuple[Optional[str], ...]] = None,
+ is_flux=False,
+ ) -> None:
+ r"""
+ Set whether to use xla flash attention from `torch_xla` or not.
+
+ Args:
+ use_xla_flash_attention (`bool`):
+ Whether to use pallas flash attention kernel from `torch_xla` or not.
+ partition_spec (`Tuple[]`, *optional*):
+ Specify the partition specification if using SPMD. Otherwise None.
+ """
+ if use_xla_flash_attention:
+ if not is_torch_xla_available:
+ raise "torch_xla is not available"
+ elif is_torch_xla_version("<", "2.3"):
+ raise "flash attention pallas kernel is supported from torch_xla version 2.3"
+ elif is_spmd() and is_torch_xla_version("<", "2.4"):
+ raise "flash attention pallas kernel using SPMD is supported from torch_xla version 2.4"
+ else:
+ if is_flux:
+ processor = XLAFluxFlashAttnProcessor2_0(partition_spec)
+ else:
+ processor = XLAFlashAttnProcessor2_0(partition_spec)
+ else:
+ processor = (
+ AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
+ )
+ self.set_processor(processor)
+
+ def set_use_npu_flash_attention(self, use_npu_flash_attention: bool) -> None:
+ r"""
+ Set whether to use npu flash attention from `torch_npu` or not.
+
+ """
+ if use_npu_flash_attention:
+ processor = AttnProcessorNPU()
+ else:
+ # set attention processor
+ # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+ # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+ # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+ processor = (
+ AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
+ )
+ self.set_processor(processor)
+
+ def set_use_memory_efficient_attention_xformers(
+ self, use_memory_efficient_attention_xformers: bool, attention_op: Optional[Callable] = None
+ ) -> None:
+ r"""
+ Set whether to use memory efficient attention from `xformers` or not.
+
+ Args:
+ use_memory_efficient_attention_xformers (`bool`):
+ Whether to use memory efficient attention from `xformers` or not.
+ attention_op (`Callable`, *optional*):
+ The attention operation to use. Defaults to `None` which uses the default attention operation from
+ `xformers`.
+ """
+ is_custom_diffusion = hasattr(self, "processor") and isinstance(
+ self.processor,
+ (CustomDiffusionAttnProcessor, CustomDiffusionXFormersAttnProcessor, CustomDiffusionAttnProcessor2_0),
+ )
+ is_added_kv_processor = hasattr(self, "processor") and isinstance(
+ self.processor,
+ (
+ AttnAddedKVProcessor,
+ AttnAddedKVProcessor2_0,
+ SlicedAttnAddedKVProcessor,
+ XFormersAttnAddedKVProcessor,
+ ),
+ )
+ is_ip_adapter = hasattr(self, "processor") and isinstance(
+ self.processor,
+ (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0, IPAdapterXFormersAttnProcessor),
+ )
+ is_joint_processor = hasattr(self, "processor") and isinstance(
+ self.processor,
+ (
+ JointAttnProcessor2_0,
+ XFormersJointAttnProcessor,
+ ),
+ )
+
+ if use_memory_efficient_attention_xformers:
+ if is_added_kv_processor and is_custom_diffusion:
+ raise NotImplementedError(
+ f"Memory efficient attention is currently not supported for custom diffusion for attention processor type {self.processor}"
+ )
+ if not is_xformers_available():
+ raise ModuleNotFoundError(
+ (
+ "Refer to https://github.com/facebookresearch/xformers for more information on how to install"
+ " xformers"
+ ),
+ name="xformers",
+ )
+ elif not torch.cuda.is_available():
+ raise ValueError(
+ "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is"
+ " only available for GPU "
+ )
+ else:
+ try:
+ # Make sure we can run the memory efficient attention
+ dtype = None
+ if attention_op is not None:
+ op_fw, op_bw = attention_op
+ dtype, *_ = op_fw.SUPPORTED_DTYPES
+ q = torch.randn((1, 2, 40), device="cuda", dtype=dtype)
+ _ = xformers.ops.memory_efficient_attention(q, q, q)
+ except Exception as e:
+ raise e
+
+ if is_custom_diffusion:
+ processor = CustomDiffusionXFormersAttnProcessor(
+ train_kv=self.processor.train_kv,
+ train_q_out=self.processor.train_q_out,
+ hidden_size=self.processor.hidden_size,
+ cross_attention_dim=self.processor.cross_attention_dim,
+ attention_op=attention_op,
+ )
+ processor.load_state_dict(self.processor.state_dict())
+ if hasattr(self.processor, "to_k_custom_diffusion"):
+ processor.to(self.processor.to_k_custom_diffusion.weight.device)
+ elif is_added_kv_processor:
+ # TODO(Patrick, Suraj, William) - currently xformers doesn't work for UnCLIP
+ # which uses this type of cross attention ONLY because the attention mask of format
+ # [0, ..., -10.000, ..., 0, ...,] is not supported
+ # throw warning
+ logger.info(
+ "Memory efficient attention with `xformers` might currently not work correctly if an attention mask is required for the attention operation."
+ )
+ processor = XFormersAttnAddedKVProcessor(attention_op=attention_op)
+ elif is_ip_adapter:
+ processor = IPAdapterXFormersAttnProcessor(
+ hidden_size=self.processor.hidden_size,
+ cross_attention_dim=self.processor.cross_attention_dim,
+ num_tokens=self.processor.num_tokens,
+ scale=self.processor.scale,
+ attention_op=attention_op,
+ )
+ processor.load_state_dict(self.processor.state_dict())
+ if hasattr(self.processor, "to_k_ip"):
+ processor.to(
+ device=self.processor.to_k_ip[0].weight.device, dtype=self.processor.to_k_ip[0].weight.dtype
+ )
+ elif is_joint_processor:
+ processor = XFormersJointAttnProcessor(attention_op=attention_op)
+ else:
+ processor = XFormersAttnProcessor(attention_op=attention_op)
+ else:
+ if is_custom_diffusion:
+ attn_processor_class = (
+ CustomDiffusionAttnProcessor2_0
+ if hasattr(F, "scaled_dot_product_attention")
+ else CustomDiffusionAttnProcessor
+ )
+ processor = attn_processor_class(
+ train_kv=self.processor.train_kv,
+ train_q_out=self.processor.train_q_out,
+ hidden_size=self.processor.hidden_size,
+ cross_attention_dim=self.processor.cross_attention_dim,
+ )
+ processor.load_state_dict(self.processor.state_dict())
+ if hasattr(self.processor, "to_k_custom_diffusion"):
+ processor.to(self.processor.to_k_custom_diffusion.weight.device)
+ elif is_ip_adapter:
+ processor = IPAdapterAttnProcessor2_0(
+ hidden_size=self.processor.hidden_size,
+ cross_attention_dim=self.processor.cross_attention_dim,
+ num_tokens=self.processor.num_tokens,
+ scale=self.processor.scale,
+ )
+ processor.load_state_dict(self.processor.state_dict())
+ if hasattr(self.processor, "to_k_ip"):
+ processor.to(
+ device=self.processor.to_k_ip[0].weight.device, dtype=self.processor.to_k_ip[0].weight.dtype
+ )
+ else:
+ # set attention processor
+ # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+ # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+ # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+ processor = (
+ AttnProcessor2_0()
+ if hasattr(F, "scaled_dot_product_attention") and self.scale_qk
+ else AttnProcessor()
+ )
+
+ self.set_processor(processor)
+
+ def set_attention_slice(self, slice_size: int) -> None:
+ r"""
+ Set the slice size for attention computation.
+
+ Args:
+ slice_size (`int`):
+ The slice size for attention computation.
+ """
+ if slice_size is not None and slice_size > self.sliceable_head_dim:
+ raise ValueError(f"slice_size {slice_size} has to be smaller or equal to {self.sliceable_head_dim}.")
+
+ if slice_size is not None and self.added_kv_proj_dim is not None:
+ processor = SlicedAttnAddedKVProcessor(slice_size)
+ elif slice_size is not None:
+ processor = SlicedAttnProcessor(slice_size)
+ elif self.added_kv_proj_dim is not None:
+ processor = AttnAddedKVProcessor()
+ else:
+ # set attention processor
+ # We use the AttnProcessor2_0 by default when torch 2.x is used which uses
+ # torch.nn.functional.scaled_dot_product_attention for native Flash/memory_efficient_attention
+ # but only if it has the default `scale` argument. TODO remove scale_qk check when we move to torch 2.1
+ processor = (
+ AttnProcessor2_0() if hasattr(F, "scaled_dot_product_attention") and self.scale_qk else AttnProcessor()
+ )
+
+ self.set_processor(processor)
+
+ def set_processor(self, processor: "AttnProcessor") -> None:
+ r"""
+ Set the attention processor to use.
+
+ Args:
+ processor (`AttnProcessor`):
+ The attention processor to use.
+ """
+ # if current processor is in `self._modules` and if passed `processor` is not, we need to
+ # pop `processor` from `self._modules`
+ if (
+ hasattr(self, "processor")
+ and isinstance(self.processor, torch.nn.Module)
+ and not isinstance(processor, torch.nn.Module)
+ ):
+ logger.info(f"You are removing possibly trained weights of {self.processor} with {processor}")
+ self._modules.pop("processor")
+
+ self.processor = processor
+
+ def get_processor(self, return_deprecated_lora: bool = False) -> "AttentionProcessor":
+ r"""
+ Get the attention processor in use.
+
+ Args:
+ return_deprecated_lora (`bool`, *optional*, defaults to `False`):
+ Set to `True` to return the deprecated LoRA attention processor.
+
+ Returns:
+ "AttentionProcessor": The attention processor in use.
+ """
+ if not return_deprecated_lora:
+ return self.processor
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ **cross_attention_kwargs,
+ ) -> torch.Tensor:
+ r"""
+ The forward method of the `Attention` class.
+
+ Args:
+ hidden_states (`torch.Tensor`):
+ The hidden states of the query.
+ encoder_hidden_states (`torch.Tensor`, *optional*):
+ The hidden states of the encoder.
+ attention_mask (`torch.Tensor`, *optional*):
+ The attention mask to use. If `None`, no mask is applied.
+ **cross_attention_kwargs:
+ Additional keyword arguments to pass along to the cross attention.
+
+ Returns:
+ `torch.Tensor`: The output of the attention layer.
+ """
+ # The `Attention` class can call different attention processors / attention functions
+ # here we simply pass along all tensors to the selected processor class
+ # For standard processors that are defined here, `**cross_attention_kwargs` is empty
+
+ attn_parameters = set(inspect.signature(self.processor.__call__).parameters.keys())
+ quiet_attn_parameters = {"ip_adapter_masks", "ip_hidden_states"}
+ unused_kwargs = [
+ k for k, _ in cross_attention_kwargs.items() if k not in attn_parameters and k not in quiet_attn_parameters
+ ]
+ if len(unused_kwargs) > 0:
+ logger.warning(
+ f"cross_attention_kwargs {unused_kwargs} are not expected by {self.processor.__class__.__name__} and will be ignored."
+ )
+ cross_attention_kwargs = {k: w for k, w in cross_attention_kwargs.items() if k in attn_parameters}
+
+ return self.processor(
+ self,
+ hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ attention_mask=attention_mask,
+ **cross_attention_kwargs,
+ )
+
+ def batch_to_head_dim(self, tensor: torch.Tensor) -> torch.Tensor:
+ r"""
+ Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size // heads, seq_len, dim * heads]`. `heads`
+ is the number of heads initialized while constructing the `Attention` class.
+
+ Args:
+ tensor (`torch.Tensor`): The tensor to reshape.
+
+ Returns:
+ `torch.Tensor`: The reshaped tensor.
+ """
+ head_size = self.heads
+ batch_size, seq_len, dim = tensor.shape
+ tensor = tensor.reshape(batch_size // head_size, head_size, seq_len, dim)
+ tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size, seq_len, dim * head_size)
+ return tensor
+
+ def head_to_batch_dim(self, tensor: torch.Tensor, out_dim: int = 3) -> torch.Tensor:
+ r"""
+ Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size, seq_len, heads, dim // heads]` `heads` is
+ the number of heads initialized while constructing the `Attention` class.
+
+ Args:
+ tensor (`torch.Tensor`): The tensor to reshape.
+ out_dim (`int`, *optional*, defaults to `3`): The output dimension of the tensor. If `3`, the tensor is
+ reshaped to `[batch_size * heads, seq_len, dim // heads]`.
+
+ Returns:
+ `torch.Tensor`: The reshaped tensor.
+ """
+ head_size = self.heads
+ if tensor.ndim == 3:
+ batch_size, seq_len, dim = tensor.shape
+ extra_dim = 1
+ else:
+ batch_size, extra_dim, seq_len, dim = tensor.shape
+ tensor = tensor.reshape(batch_size, seq_len * extra_dim, head_size, dim // head_size)
+ tensor = tensor.permute(0, 2, 1, 3)
+
+ if out_dim == 3:
+ tensor = tensor.reshape(batch_size * head_size, seq_len * extra_dim, dim // head_size)
+
+ return tensor
+
+ def get_attention_scores(
+ self, query: torch.Tensor, key: torch.Tensor, attention_mask: Optional[torch.Tensor] = None
+ ) -> torch.Tensor:
+ r"""
+ Compute the attention scores.
+
+ Args:
+ query (`torch.Tensor`): The query tensor.
+ key (`torch.Tensor`): The key tensor.
+ attention_mask (`torch.Tensor`, *optional*): The attention mask to use. If `None`, no mask is applied.
+
+ Returns:
+ `torch.Tensor`: The attention probabilities/scores.
+ """
+ dtype = query.dtype
+ if self.upcast_attention:
+ query = query.float()
+ key = key.float()
+
+ if attention_mask is None:
+ baddbmm_input = torch.empty(
+ query.shape[0], query.shape[1], key.shape[1], dtype=query.dtype, device=query.device
+ )
+ beta = 0
+ else:
+ baddbmm_input = attention_mask
+ beta = 1
+
+ attention_scores = torch.baddbmm(
+ baddbmm_input,
+ query,
+ key.transpose(-1, -2),
+ beta=beta,
+ alpha=self.scale,
+ )
+ del baddbmm_input
+
+ if self.upcast_softmax:
+ attention_scores = attention_scores.float()
+
+ attention_probs = attention_scores.softmax(dim=-1)
+ del attention_scores
+
+ attention_probs = attention_probs.to(dtype)
+
+ return attention_probs
+
+ def prepare_attention_mask(
+ self, attention_mask: torch.Tensor, target_length: int, batch_size: int, out_dim: int = 3
+ ) -> torch.Tensor:
+ r"""
+ Prepare the attention mask for the attention computation.
+
+ Args:
+ attention_mask (`torch.Tensor`):
+ The attention mask to prepare.
+ target_length (`int`):
+ The target length of the attention mask. This is the length of the attention mask after padding.
+ batch_size (`int`):
+ The batch size, which is used to repeat the attention mask.
+ out_dim (`int`, *optional*, defaults to `3`):
+ The output dimension of the attention mask. Can be either `3` or `4`.
+
+ Returns:
+ `torch.Tensor`: The prepared attention mask.
+ """
+ head_size = self.heads
+ if attention_mask is None:
+ return attention_mask
+
+ current_length: int = attention_mask.shape[-1]
+ if current_length != target_length:
+ if attention_mask.device.type == "mps":
+ # HACK: MPS: Does not support padding by greater than dimension of input tensor.
+ # Instead, we can manually construct the padding tensor.
+ padding_shape = (attention_mask.shape[0], attention_mask.shape[1], target_length)
+ padding = torch.zeros(padding_shape, dtype=attention_mask.dtype, device=attention_mask.device)
+ attention_mask = torch.cat([attention_mask, padding], dim=2)
+ else:
+ # TODO: for pipelines such as stable-diffusion, padding cross-attn mask:
+ # we want to instead pad by (0, remaining_length), where remaining_length is:
+ # remaining_length: int = target_length - current_length
+ # TODO: re-enable tests/models/test_models_unet_2d_condition.py#test_model_xattn_padding
+ attention_mask = F.pad(attention_mask, (0, target_length), value=0.0)
+
+ if out_dim == 3:
+ if attention_mask.shape[0] < batch_size * head_size:
+ attention_mask = attention_mask.repeat_interleave(
+ head_size, dim=0, output_size=attention_mask.shape[0] * head_size
+ )
+ elif out_dim == 4:
+ attention_mask = attention_mask.unsqueeze(1)
+ attention_mask = attention_mask.repeat_interleave(
+ head_size, dim=1, output_size=attention_mask.shape[1] * head_size
+ )
+
+ return attention_mask
+
+ def norm_encoder_hidden_states(self, encoder_hidden_states: torch.Tensor) -> torch.Tensor:
+ r"""
+ Normalize the encoder hidden states. Requires `self.norm_cross` to be specified when constructing the
+ `Attention` class.
+
+ Args:
+ encoder_hidden_states (`torch.Tensor`): Hidden states of the encoder.
+
+ Returns:
+ `torch.Tensor`: The normalized encoder hidden states.
+ """
+ assert self.norm_cross is not None, "self.norm_cross must be defined to call self.norm_encoder_hidden_states"
+
+ if isinstance(self.norm_cross, nn.LayerNorm):
+ encoder_hidden_states = self.norm_cross(encoder_hidden_states)
+ elif isinstance(self.norm_cross, nn.GroupNorm):
+ # Group norm norms along the channels dimension and expects
+ # input to be in the shape of (N, C, *). In this case, we want
+ # to norm along the hidden dimension, so we need to move
+ # (batch_size, sequence_length, hidden_size) ->
+ # (batch_size, hidden_size, sequence_length)
+ encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
+ encoder_hidden_states = self.norm_cross(encoder_hidden_states)
+ encoder_hidden_states = encoder_hidden_states.transpose(1, 2)
+ else:
+ assert False
+
+ return encoder_hidden_states
+
+ @torch.no_grad()
+ def fuse_projections(self, fuse=True):
+ device = self.to_q.weight.data.device
+ dtype = self.to_q.weight.data.dtype
+
+ if not self.is_cross_attention:
+ # fetch weight matrices.
+ concatenated_weights = torch.cat([self.to_q.weight.data, self.to_k.weight.data, self.to_v.weight.data])
+ in_features = concatenated_weights.shape[1]
+ out_features = concatenated_weights.shape[0]
+
+ # create a new single projection layer and copy over the weights.
+ self.to_qkv = nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype)
+ self.to_qkv.weight.copy_(concatenated_weights)
+ if self.use_bias:
+ concatenated_bias = torch.cat([self.to_q.bias.data, self.to_k.bias.data, self.to_v.bias.data])
+ self.to_qkv.bias.copy_(concatenated_bias)
+
+ else:
+ concatenated_weights = torch.cat([self.to_k.weight.data, self.to_v.weight.data])
+ in_features = concatenated_weights.shape[1]
+ out_features = concatenated_weights.shape[0]
+
+ self.to_kv = nn.Linear(in_features, out_features, bias=self.use_bias, device=device, dtype=dtype)
+ self.to_kv.weight.copy_(concatenated_weights)
+ if self.use_bias:
+ concatenated_bias = torch.cat([self.to_k.bias.data, self.to_v.bias.data])
+ self.to_kv.bias.copy_(concatenated_bias)
+
+ # handle added projections for SD3 and others.
+ if (
+ getattr(self, "add_q_proj", None) is not None
+ and getattr(self, "add_k_proj", None) is not None
+ and getattr(self, "add_v_proj", None) is not None
+ ):
+ concatenated_weights = torch.cat(
+ [self.add_q_proj.weight.data, self.add_k_proj.weight.data, self.add_v_proj.weight.data]
+ )
+ in_features = concatenated_weights.shape[1]
+ out_features = concatenated_weights.shape[0]
+
+ self.to_added_qkv = nn.Linear(
+ in_features, out_features, bias=self.added_proj_bias, device=device, dtype=dtype
+ )
+ self.to_added_qkv.weight.copy_(concatenated_weights)
+ if self.added_proj_bias:
+ concatenated_bias = torch.cat(
+ [self.add_q_proj.bias.data, self.add_k_proj.bias.data, self.add_v_proj.bias.data]
+ )
+ self.to_added_qkv.bias.copy_(concatenated_bias)
+
+ self.fused_projections = fuse
+
+
+class SanaMultiscaleAttentionProjection(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ num_attention_heads: int,
+ kernel_size: int,
+ ) -> None:
+ super().__init__()
+
+ channels = 3 * in_channels
+ self.proj_in = nn.Conv2d(
+ channels,
+ channels,
+ kernel_size,
+ padding=kernel_size // 2,
+ groups=channels,
+ bias=False,
+ )
+ self.proj_out = nn.Conv2d(channels, channels, 1, 1, 0, groups=3 * num_attention_heads, bias=False)
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ hidden_states = self.proj_in(hidden_states)
+ hidden_states = self.proj_out(hidden_states)
+ return hidden_states
+
+
+class SanaMultiscaleLinearAttention(nn.Module):
+ r"""Lightweight multi-scale linear attention"""
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ num_attention_heads: Optional[int] = None,
+ attention_head_dim: int = 8,
+ mult: float = 1.0,
+ norm_type: str = "batch_norm",
+ kernel_sizes: Tuple[int, ...] = (5,),
+ eps: float = 1e-15,
+ residual_connection: bool = False,
+ ):
+ super().__init__()
+
+ # To prevent circular import
+ from .normalization import get_normalization
+
+ self.eps = eps
+ self.attention_head_dim = attention_head_dim
+ self.norm_type = norm_type
+ self.residual_connection = residual_connection
+
+ num_attention_heads = (
+ int(in_channels // attention_head_dim * mult) if num_attention_heads is None else num_attention_heads
+ )
+ inner_dim = num_attention_heads * attention_head_dim
+
+ self.to_q = nn.Linear(in_channels, inner_dim, bias=False)
+ self.to_k = nn.Linear(in_channels, inner_dim, bias=False)
+ self.to_v = nn.Linear(in_channels, inner_dim, bias=False)
+
+ self.to_qkv_multiscale = nn.ModuleList()
+ for kernel_size in kernel_sizes:
+ self.to_qkv_multiscale.append(
+ SanaMultiscaleAttentionProjection(inner_dim, num_attention_heads, kernel_size)
+ )
+
+ self.nonlinearity = nn.ReLU()
+ self.to_out = nn.Linear(inner_dim * (1 + len(kernel_sizes)), out_channels, bias=False)
+ self.norm_out = get_normalization(norm_type, num_features=out_channels)
+
+ self.processor = SanaMultiscaleAttnProcessor2_0()
+
+ def apply_linear_attention(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+ value = F.pad(value, (0, 0, 0, 1), mode="constant", value=1) # Adds padding
+ scores = torch.matmul(value, key.transpose(-1, -2))
+ hidden_states = torch.matmul(scores, query)
+
+ hidden_states = hidden_states.to(dtype=torch.float32)
+ hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + self.eps)
+ return hidden_states
+
+ def apply_quadratic_attention(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
+ scores = torch.matmul(key.transpose(-1, -2), query)
+ scores = scores.to(dtype=torch.float32)
+ scores = scores / (torch.sum(scores, dim=2, keepdim=True) + self.eps)
+ hidden_states = torch.matmul(value, scores.to(value.dtype))
+ return hidden_states
+
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+ return self.processor(self, hidden_states)
+
+
+class MochiAttention(nn.Module):
+ def __init__(
+ self,
+ query_dim: int,
+ added_kv_proj_dim: int,
+ processor: "MochiAttnProcessor2_0",
+ heads: int = 8,
+ dim_head: int = 64,
+ dropout: float = 0.0,
+ bias: bool = False,
+ added_proj_bias: bool = True,
+ out_dim: Optional[int] = None,
+ out_context_dim: Optional[int] = None,
+ out_bias: bool = True,
+ context_pre_only: bool = False,
+ eps: float = 1e-5,
+ ):
+ super().__init__()
+ from .normalization import MochiRMSNorm
+
+ self.inner_dim = out_dim if out_dim is not None else dim_head * heads
+ self.out_dim = out_dim if out_dim is not None else query_dim
+ self.out_context_dim = out_context_dim if out_context_dim else query_dim
+ self.context_pre_only = context_pre_only
+
+ self.heads = out_dim // dim_head if out_dim is not None else heads
+
+ self.norm_q = MochiRMSNorm(dim_head, eps, True)
+ self.norm_k = MochiRMSNorm(dim_head, eps, True)
+ self.norm_added_q = MochiRMSNorm(dim_head, eps, True)
+ self.norm_added_k = MochiRMSNorm(dim_head, eps, True)
+
+ self.to_q = nn.Linear(query_dim, self.inner_dim, bias=bias)
+ self.to_k = nn.Linear(query_dim, self.inner_dim, bias=bias)
+ self.to_v = nn.Linear(query_dim, self.inner_dim, bias=bias)
+
+ self.add_k_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ self.add_v_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+ if self.context_pre_only is not None:
+ self.add_q_proj = nn.Linear(added_kv_proj_dim, self.inner_dim, bias=added_proj_bias)
+
+ self.to_out = nn.ModuleList([])
+ self.to_out.append(nn.Linear(self.inner_dim, self.out_dim, bias=out_bias))
+ self.to_out.append(nn.Dropout(dropout))
+
+ if not self.context_pre_only:
+ self.to_add_out = nn.Linear(self.inner_dim, self.out_context_dim, bias=out_bias)
+
+ self.processor = processor
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ **kwargs,
+ ):
+ return self.processor(
+ self,
+ hidden_states,
+ encoder_hidden_states=encoder_hidden_states,
+ attention_mask=attention_mask,
+ **kwargs,
+ )
+
+
+class MochiAttnProcessor2_0:
+ """Attention processor used in Mochi."""
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError("MochiAttnProcessor2_0 requires PyTorch 2.0. To use it, please upgrade PyTorch to 2.0.")
+
+ def __call__(
+ self,
+ attn: "MochiAttention",
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ attention_mask: torch.Tensor,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ query = query.unflatten(2, (attn.heads, -1))
+ key = key.unflatten(2, (attn.heads, -1))
+ value = value.unflatten(2, (attn.heads, -1))
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ encoder_query = attn.add_q_proj(encoder_hidden_states)
+ encoder_key = attn.add_k_proj(encoder_hidden_states)
+ encoder_value = attn.add_v_proj(encoder_hidden_states)
+
+ encoder_query = encoder_query.unflatten(2, (attn.heads, -1))
+ encoder_key = encoder_key.unflatten(2, (attn.heads, -1))
+ encoder_value = encoder_value.unflatten(2, (attn.heads, -1))
+
+ if attn.norm_added_q is not None:
+ encoder_query = attn.norm_added_q(encoder_query)
+ if attn.norm_added_k is not None:
+ encoder_key = attn.norm_added_k(encoder_key)
+
+ if image_rotary_emb is not None:
+
+ def apply_rotary_emb(x, freqs_cos, freqs_sin):
+ x_even = x[..., 0::2].float()
+ x_odd = x[..., 1::2].float()
+
+ cos = (x_even * freqs_cos - x_odd * freqs_sin).to(x.dtype)
+ sin = (x_even * freqs_sin + x_odd * freqs_cos).to(x.dtype)
+
+ return torch.stack([cos, sin], dim=-1).flatten(-2)
+
+ query = apply_rotary_emb(query, *image_rotary_emb)
+ key = apply_rotary_emb(key, *image_rotary_emb)
+
+ query, key, value = query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)
+ encoder_query, encoder_key, encoder_value = (
+ encoder_query.transpose(1, 2),
+ encoder_key.transpose(1, 2),
+ encoder_value.transpose(1, 2),
+ )
+
+ sequence_length = query.size(2)
+ encoder_sequence_length = encoder_query.size(2)
+ total_length = sequence_length + encoder_sequence_length
+
+ batch_size, heads, _, dim = query.shape
+ attn_outputs = []
+ for idx in range(batch_size):
+ mask = attention_mask[idx][None, :]
+ valid_prompt_token_indices = torch.nonzero(mask.flatten(), as_tuple=False).flatten()
+
+ valid_encoder_query = encoder_query[idx : idx + 1, :, valid_prompt_token_indices, :]
+ valid_encoder_key = encoder_key[idx : idx + 1, :, valid_prompt_token_indices, :]
+ valid_encoder_value = encoder_value[idx : idx + 1, :, valid_prompt_token_indices, :]
+
+ valid_query = torch.cat([query[idx : idx + 1], valid_encoder_query], dim=2)
+ valid_key = torch.cat([key[idx : idx + 1], valid_encoder_key], dim=2)
+ valid_value = torch.cat([value[idx : idx + 1], valid_encoder_value], dim=2)
+
+ attn_output = F.scaled_dot_product_attention(
+ valid_query, valid_key, valid_value, dropout_p=0.0, is_causal=False
+ )
+ valid_sequence_length = attn_output.size(2)
+ attn_output = F.pad(attn_output, (0, 0, 0, total_length - valid_sequence_length))
+ attn_outputs.append(attn_output)
+
+ hidden_states = torch.cat(attn_outputs, dim=0)
+ hidden_states = hidden_states.transpose(1, 2).flatten(2, 3)
+
+ hidden_states, encoder_hidden_states = hidden_states.split_with_sizes(
+ (sequence_length, encoder_sequence_length), dim=1
+ )
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if hasattr(attn, "to_add_out"):
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ return hidden_states, encoder_hidden_states
+
+
+class AttnProcessor:
+ r"""
+ Default processor for performing attention-related computations.
+ """
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.Tensor:
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+ deprecate("scale", "1.0.0", deprecation_message)
+
+ residual = hidden_states
+
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ query = attn.head_to_batch_dim(query)
+ key = attn.head_to_batch_dim(key)
+ value = attn.head_to_batch_dim(value)
+
+ attention_probs = attn.get_attention_scores(query, key, attention_mask)
+ hidden_states = torch.bmm(attention_probs, value)
+ hidden_states = attn.batch_to_head_dim(hidden_states)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class CustomDiffusionAttnProcessor(nn.Module):
+ r"""
+ Processor for implementing attention for the Custom Diffusion method.
+
+ Args:
+ train_kv (`bool`, defaults to `True`):
+ Whether to newly train the key and value matrices corresponding to the text features.
+ train_q_out (`bool`, defaults to `True`):
+ Whether to newly train query matrices corresponding to the latent image features.
+ hidden_size (`int`, *optional*, defaults to `None`):
+ The hidden size of the attention layer.
+ cross_attention_dim (`int`, *optional*, defaults to `None`):
+ The number of channels in the `encoder_hidden_states`.
+ out_bias (`bool`, defaults to `True`):
+ Whether to include the bias parameter in `train_q_out`.
+ dropout (`float`, *optional*, defaults to 0.0):
+ The dropout probability to use.
+ """
+
+ def __init__(
+ self,
+ train_kv: bool = True,
+ train_q_out: bool = True,
+ hidden_size: Optional[int] = None,
+ cross_attention_dim: Optional[int] = None,
+ out_bias: bool = True,
+ dropout: float = 0.0,
+ ):
+ super().__init__()
+ self.train_kv = train_kv
+ self.train_q_out = train_q_out
+
+ self.hidden_size = hidden_size
+ self.cross_attention_dim = cross_attention_dim
+
+ # `_custom_diffusion` id for easy serialization and loading.
+ if self.train_kv:
+ self.to_k_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
+ self.to_v_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
+ if self.train_q_out:
+ self.to_q_custom_diffusion = nn.Linear(hidden_size, hidden_size, bias=False)
+ self.to_out_custom_diffusion = nn.ModuleList([])
+ self.to_out_custom_diffusion.append(nn.Linear(hidden_size, hidden_size, bias=out_bias))
+ self.to_out_custom_diffusion.append(nn.Dropout(dropout))
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ batch_size, sequence_length, _ = hidden_states.shape
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ if self.train_q_out:
+ query = self.to_q_custom_diffusion(hidden_states).to(attn.to_q.weight.dtype)
+ else:
+ query = attn.to_q(hidden_states.to(attn.to_q.weight.dtype))
+
+ if encoder_hidden_states is None:
+ crossattn = False
+ encoder_hidden_states = hidden_states
+ else:
+ crossattn = True
+ if attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ if self.train_kv:
+ key = self.to_k_custom_diffusion(encoder_hidden_states.to(self.to_k_custom_diffusion.weight.dtype))
+ value = self.to_v_custom_diffusion(encoder_hidden_states.to(self.to_v_custom_diffusion.weight.dtype))
+ key = key.to(attn.to_q.weight.dtype)
+ value = value.to(attn.to_q.weight.dtype)
+ else:
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ if crossattn:
+ detach = torch.ones_like(key)
+ detach[:, :1, :] = detach[:, :1, :] * 0.0
+ key = detach * key + (1 - detach) * key.detach()
+ value = detach * value + (1 - detach) * value.detach()
+
+ query = attn.head_to_batch_dim(query)
+ key = attn.head_to_batch_dim(key)
+ value = attn.head_to_batch_dim(value)
+
+ attention_probs = attn.get_attention_scores(query, key, attention_mask)
+ hidden_states = torch.bmm(attention_probs, value)
+ hidden_states = attn.batch_to_head_dim(hidden_states)
+
+ if self.train_q_out:
+ # linear proj
+ hidden_states = self.to_out_custom_diffusion[0](hidden_states)
+ # dropout
+ hidden_states = self.to_out_custom_diffusion[1](hidden_states)
+ else:
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ return hidden_states
+
+
+class AttnAddedKVProcessor:
+ r"""
+ Processor for performing attention-related computations with extra learnable key and value matrices for the text
+ encoder.
+ """
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.Tensor:
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+ deprecate("scale", "1.0.0", deprecation_message)
+
+ residual = hidden_states
+
+ hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)
+ batch_size, sequence_length, _ = hidden_states.shape
+
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+ query = attn.head_to_batch_dim(query)
+
+ encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+ encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+ encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj)
+ encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj)
+
+ if not attn.only_cross_attention:
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+ key = attn.head_to_batch_dim(key)
+ value = attn.head_to_batch_dim(value)
+ key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)
+ value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)
+ else:
+ key = encoder_hidden_states_key_proj
+ value = encoder_hidden_states_value_proj
+
+ attention_probs = attn.get_attention_scores(query, key, attention_mask)
+ hidden_states = torch.bmm(attention_probs, value)
+ hidden_states = attn.batch_to_head_dim(hidden_states)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape)
+ hidden_states = hidden_states + residual
+
+ return hidden_states
+
+
+class AttnAddedKVProcessor2_0:
+ r"""
+ Processor for performing scaled dot-product attention (enabled by default if you're using PyTorch 2.0), with extra
+ learnable key and value matrices for the text encoder.
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "AttnAddedKVProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.Tensor:
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+ deprecate("scale", "1.0.0", deprecation_message)
+
+ residual = hidden_states
+
+ hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)
+ batch_size, sequence_length, _ = hidden_states.shape
+
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size, out_dim=4)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+ query = attn.head_to_batch_dim(query, out_dim=4)
+
+ encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+ encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+ encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj, out_dim=4)
+ encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj, out_dim=4)
+
+ if not attn.only_cross_attention:
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+ key = attn.head_to_batch_dim(key, out_dim=4)
+ value = attn.head_to_batch_dim(value, out_dim=4)
+ key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+ value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+ else:
+ key = encoder_hidden_states_key_proj
+ value = encoder_hidden_states_value_proj
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, residual.shape[1])
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape)
+ hidden_states = hidden_states + residual
+
+ return hidden_states
+
+
+class JointAttnProcessor2_0:
+ """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError("JointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.FloatTensor:
+ residual = hidden_states
+
+ batch_size = hidden_states.shape[0]
+
+ # `sample` projections.
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # `context` projections.
+ if encoder_hidden_states is not None:
+ encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+ encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+ encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+ encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+
+ if attn.norm_added_q is not None:
+ encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+ if attn.norm_added_k is not None:
+ encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+ query = torch.cat([query, encoder_hidden_states_query_proj], dim=2)
+ key = torch.cat([key, encoder_hidden_states_key_proj], dim=2)
+ value = torch.cat([value, encoder_hidden_states_value_proj], dim=2)
+
+ hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if encoder_hidden_states is not None:
+ # Split the attention outputs.
+ hidden_states, encoder_hidden_states = (
+ hidden_states[:, : residual.shape[1]],
+ hidden_states[:, residual.shape[1] :],
+ )
+ if not attn.context_pre_only:
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if encoder_hidden_states is not None:
+ return hidden_states, encoder_hidden_states
+ else:
+ return hidden_states
+
+
+class PAGJointAttnProcessor2_0:
+ """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "PAGJointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ ) -> torch.FloatTensor:
+ residual = hidden_states
+
+ input_ndim = hidden_states.ndim
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+ context_input_ndim = encoder_hidden_states.ndim
+ if context_input_ndim == 4:
+ batch_size, channel, height, width = encoder_hidden_states.shape
+ encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ # store the length of image patch sequences to create a mask that prevents interaction between patches
+ # similar to making the self-attention map an identity matrix
+ identity_block_size = hidden_states.shape[1]
+
+ # chunk
+ hidden_states_org, hidden_states_ptb = hidden_states.chunk(2)
+ encoder_hidden_states_org, encoder_hidden_states_ptb = encoder_hidden_states.chunk(2)
+
+ ################## original path ##################
+ batch_size = encoder_hidden_states_org.shape[0]
+
+ # `sample` projections.
+ query_org = attn.to_q(hidden_states_org)
+ key_org = attn.to_k(hidden_states_org)
+ value_org = attn.to_v(hidden_states_org)
+
+ # `context` projections.
+ encoder_hidden_states_org_query_proj = attn.add_q_proj(encoder_hidden_states_org)
+ encoder_hidden_states_org_key_proj = attn.add_k_proj(encoder_hidden_states_org)
+ encoder_hidden_states_org_value_proj = attn.add_v_proj(encoder_hidden_states_org)
+
+ # attention
+ query_org = torch.cat([query_org, encoder_hidden_states_org_query_proj], dim=1)
+ key_org = torch.cat([key_org, encoder_hidden_states_org_key_proj], dim=1)
+ value_org = torch.cat([value_org, encoder_hidden_states_org_value_proj], dim=1)
+
+ inner_dim = key_org.shape[-1]
+ head_dim = inner_dim // attn.heads
+ query_org = query_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key_org = key_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value_org = value_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ hidden_states_org = F.scaled_dot_product_attention(
+ query_org, key_org, value_org, dropout_p=0.0, is_causal=False
+ )
+ hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states_org = hidden_states_org.to(query_org.dtype)
+
+ # Split the attention outputs.
+ hidden_states_org, encoder_hidden_states_org = (
+ hidden_states_org[:, : residual.shape[1]],
+ hidden_states_org[:, residual.shape[1] :],
+ )
+
+ # linear proj
+ hidden_states_org = attn.to_out[0](hidden_states_org)
+ # dropout
+ hidden_states_org = attn.to_out[1](hidden_states_org)
+ if not attn.context_pre_only:
+ encoder_hidden_states_org = attn.to_add_out(encoder_hidden_states_org)
+
+ if input_ndim == 4:
+ hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width)
+ if context_input_ndim == 4:
+ encoder_hidden_states_org = encoder_hidden_states_org.transpose(-1, -2).reshape(
+ batch_size, channel, height, width
+ )
+
+ ################## perturbed path ##################
+
+ batch_size = encoder_hidden_states_ptb.shape[0]
+
+ # `sample` projections.
+ query_ptb = attn.to_q(hidden_states_ptb)
+ key_ptb = attn.to_k(hidden_states_ptb)
+ value_ptb = attn.to_v(hidden_states_ptb)
+
+ # `context` projections.
+ encoder_hidden_states_ptb_query_proj = attn.add_q_proj(encoder_hidden_states_ptb)
+ encoder_hidden_states_ptb_key_proj = attn.add_k_proj(encoder_hidden_states_ptb)
+ encoder_hidden_states_ptb_value_proj = attn.add_v_proj(encoder_hidden_states_ptb)
+
+ # attention
+ query_ptb = torch.cat([query_ptb, encoder_hidden_states_ptb_query_proj], dim=1)
+ key_ptb = torch.cat([key_ptb, encoder_hidden_states_ptb_key_proj], dim=1)
+ value_ptb = torch.cat([value_ptb, encoder_hidden_states_ptb_value_proj], dim=1)
+
+ inner_dim = key_ptb.shape[-1]
+ head_dim = inner_dim // attn.heads
+ query_ptb = query_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key_ptb = key_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value_ptb = value_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ # create a full mask with all entries set to 0
+ seq_len = query_ptb.size(2)
+ full_mask = torch.zeros((seq_len, seq_len), device=query_ptb.device, dtype=query_ptb.dtype)
+
+ # set the attention value between image patches to -inf
+ full_mask[:identity_block_size, :identity_block_size] = float("-inf")
+
+ # set the diagonal of the attention value between image patches to 0
+ full_mask[:identity_block_size, :identity_block_size].fill_diagonal_(0)
+
+ # expand the mask to match the attention weights shape
+ full_mask = full_mask.unsqueeze(0).unsqueeze(0) # Add batch and num_heads dimensions
+
+ hidden_states_ptb = F.scaled_dot_product_attention(
+ query_ptb, key_ptb, value_ptb, attn_mask=full_mask, dropout_p=0.0, is_causal=False
+ )
+ hidden_states_ptb = hidden_states_ptb.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states_ptb = hidden_states_ptb.to(query_ptb.dtype)
+
+ # split the attention outputs.
+ hidden_states_ptb, encoder_hidden_states_ptb = (
+ hidden_states_ptb[:, : residual.shape[1]],
+ hidden_states_ptb[:, residual.shape[1] :],
+ )
+
+ # linear proj
+ hidden_states_ptb = attn.to_out[0](hidden_states_ptb)
+ # dropout
+ hidden_states_ptb = attn.to_out[1](hidden_states_ptb)
+ if not attn.context_pre_only:
+ encoder_hidden_states_ptb = attn.to_add_out(encoder_hidden_states_ptb)
+
+ if input_ndim == 4:
+ hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width)
+ if context_input_ndim == 4:
+ encoder_hidden_states_ptb = encoder_hidden_states_ptb.transpose(-1, -2).reshape(
+ batch_size, channel, height, width
+ )
+
+ ################ concat ###############
+ hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])
+ encoder_hidden_states = torch.cat([encoder_hidden_states_org, encoder_hidden_states_ptb])
+
+ return hidden_states, encoder_hidden_states
+
+
+class PAGCFGJointAttnProcessor2_0:
+ """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "PAGCFGJointAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.FloatTensor:
+ residual = hidden_states
+
+ input_ndim = hidden_states.ndim
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+ context_input_ndim = encoder_hidden_states.ndim
+ if context_input_ndim == 4:
+ batch_size, channel, height, width = encoder_hidden_states.shape
+ encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ identity_block_size = hidden_states.shape[
+ 1
+ ] # patch embeddings width * height (correspond to self-attention map width or height)
+
+ # chunk
+ hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3)
+ hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org])
+
+ (
+ encoder_hidden_states_uncond,
+ encoder_hidden_states_org,
+ encoder_hidden_states_ptb,
+ ) = encoder_hidden_states.chunk(3)
+ encoder_hidden_states_org = torch.cat([encoder_hidden_states_uncond, encoder_hidden_states_org])
+
+ ################## original path ##################
+ batch_size = encoder_hidden_states_org.shape[0]
+
+ # `sample` projections.
+ query_org = attn.to_q(hidden_states_org)
+ key_org = attn.to_k(hidden_states_org)
+ value_org = attn.to_v(hidden_states_org)
+
+ # `context` projections.
+ encoder_hidden_states_org_query_proj = attn.add_q_proj(encoder_hidden_states_org)
+ encoder_hidden_states_org_key_proj = attn.add_k_proj(encoder_hidden_states_org)
+ encoder_hidden_states_org_value_proj = attn.add_v_proj(encoder_hidden_states_org)
+
+ # attention
+ query_org = torch.cat([query_org, encoder_hidden_states_org_query_proj], dim=1)
+ key_org = torch.cat([key_org, encoder_hidden_states_org_key_proj], dim=1)
+ value_org = torch.cat([value_org, encoder_hidden_states_org_value_proj], dim=1)
+
+ inner_dim = key_org.shape[-1]
+ head_dim = inner_dim // attn.heads
+ query_org = query_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key_org = key_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value_org = value_org.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ hidden_states_org = F.scaled_dot_product_attention(
+ query_org, key_org, value_org, dropout_p=0.0, is_causal=False
+ )
+ hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states_org = hidden_states_org.to(query_org.dtype)
+
+ # Split the attention outputs.
+ hidden_states_org, encoder_hidden_states_org = (
+ hidden_states_org[:, : residual.shape[1]],
+ hidden_states_org[:, residual.shape[1] :],
+ )
+
+ # linear proj
+ hidden_states_org = attn.to_out[0](hidden_states_org)
+ # dropout
+ hidden_states_org = attn.to_out[1](hidden_states_org)
+ if not attn.context_pre_only:
+ encoder_hidden_states_org = attn.to_add_out(encoder_hidden_states_org)
+
+ if input_ndim == 4:
+ hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width)
+ if context_input_ndim == 4:
+ encoder_hidden_states_org = encoder_hidden_states_org.transpose(-1, -2).reshape(
+ batch_size, channel, height, width
+ )
+
+ ################## perturbed path ##################
+
+ batch_size = encoder_hidden_states_ptb.shape[0]
+
+ # `sample` projections.
+ query_ptb = attn.to_q(hidden_states_ptb)
+ key_ptb = attn.to_k(hidden_states_ptb)
+ value_ptb = attn.to_v(hidden_states_ptb)
+
+ # `context` projections.
+ encoder_hidden_states_ptb_query_proj = attn.add_q_proj(encoder_hidden_states_ptb)
+ encoder_hidden_states_ptb_key_proj = attn.add_k_proj(encoder_hidden_states_ptb)
+ encoder_hidden_states_ptb_value_proj = attn.add_v_proj(encoder_hidden_states_ptb)
+
+ # attention
+ query_ptb = torch.cat([query_ptb, encoder_hidden_states_ptb_query_proj], dim=1)
+ key_ptb = torch.cat([key_ptb, encoder_hidden_states_ptb_key_proj], dim=1)
+ value_ptb = torch.cat([value_ptb, encoder_hidden_states_ptb_value_proj], dim=1)
+
+ inner_dim = key_ptb.shape[-1]
+ head_dim = inner_dim // attn.heads
+ query_ptb = query_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key_ptb = key_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value_ptb = value_ptb.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ # create a full mask with all entries set to 0
+ seq_len = query_ptb.size(2)
+ full_mask = torch.zeros((seq_len, seq_len), device=query_ptb.device, dtype=query_ptb.dtype)
+
+ # set the attention value between image patches to -inf
+ full_mask[:identity_block_size, :identity_block_size] = float("-inf")
+
+ # set the diagonal of the attention value between image patches to 0
+ full_mask[:identity_block_size, :identity_block_size].fill_diagonal_(0)
+
+ # expand the mask to match the attention weights shape
+ full_mask = full_mask.unsqueeze(0).unsqueeze(0) # Add batch and num_heads dimensions
+
+ hidden_states_ptb = F.scaled_dot_product_attention(
+ query_ptb, key_ptb, value_ptb, attn_mask=full_mask, dropout_p=0.0, is_causal=False
+ )
+ hidden_states_ptb = hidden_states_ptb.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states_ptb = hidden_states_ptb.to(query_ptb.dtype)
+
+ # split the attention outputs.
+ hidden_states_ptb, encoder_hidden_states_ptb = (
+ hidden_states_ptb[:, : residual.shape[1]],
+ hidden_states_ptb[:, residual.shape[1] :],
+ )
+
+ # linear proj
+ hidden_states_ptb = attn.to_out[0](hidden_states_ptb)
+ # dropout
+ hidden_states_ptb = attn.to_out[1](hidden_states_ptb)
+ if not attn.context_pre_only:
+ encoder_hidden_states_ptb = attn.to_add_out(encoder_hidden_states_ptb)
+
+ if input_ndim == 4:
+ hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width)
+ if context_input_ndim == 4:
+ encoder_hidden_states_ptb = encoder_hidden_states_ptb.transpose(-1, -2).reshape(
+ batch_size, channel, height, width
+ )
+
+ ################ concat ###############
+ hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])
+ encoder_hidden_states = torch.cat([encoder_hidden_states_org, encoder_hidden_states_ptb])
+
+ return hidden_states, encoder_hidden_states
+
+
+class FusedJointAttnProcessor2_0:
+ """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.FloatTensor:
+ residual = hidden_states
+
+ input_ndim = hidden_states.ndim
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+ context_input_ndim = encoder_hidden_states.ndim
+ if context_input_ndim == 4:
+ batch_size, channel, height, width = encoder_hidden_states.shape
+ encoder_hidden_states = encoder_hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size = encoder_hidden_states.shape[0]
+
+ # `sample` projections.
+ qkv = attn.to_qkv(hidden_states)
+ split_size = qkv.shape[-1] // 3
+ query, key, value = torch.split(qkv, split_size, dim=-1)
+
+ # `context` projections.
+ encoder_qkv = attn.to_added_qkv(encoder_hidden_states)
+ split_size = encoder_qkv.shape[-1] // 3
+ (
+ encoder_hidden_states_query_proj,
+ encoder_hidden_states_key_proj,
+ encoder_hidden_states_value_proj,
+ ) = torch.split(encoder_qkv, split_size, dim=-1)
+
+ # attention
+ query = torch.cat([query, encoder_hidden_states_query_proj], dim=1)
+ key = torch.cat([key, encoder_hidden_states_key_proj], dim=1)
+ value = torch.cat([value, encoder_hidden_states_value_proj], dim=1)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # Split the attention outputs.
+ hidden_states, encoder_hidden_states = (
+ hidden_states[:, : residual.shape[1]],
+ hidden_states[:, residual.shape[1] :],
+ )
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+ if not attn.context_pre_only:
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+ if context_input_ndim == 4:
+ encoder_hidden_states = encoder_hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ return hidden_states, encoder_hidden_states
+
+
+class XFormersJointAttnProcessor:
+ r"""
+ Processor for implementing memory efficient attention using xFormers.
+
+ Args:
+ attention_op (`Callable`, *optional*, defaults to `None`):
+ The base
+ [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to
+ use as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best
+ operator.
+ """
+
+ def __init__(self, attention_op: Optional[Callable] = None):
+ self.attention_op = attention_op
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.FloatTensor:
+ residual = hidden_states
+
+ # `sample` projections.
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ query = attn.head_to_batch_dim(query).contiguous()
+ key = attn.head_to_batch_dim(key).contiguous()
+ value = attn.head_to_batch_dim(value).contiguous()
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # `context` projections.
+ if encoder_hidden_states is not None:
+ encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+ encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+ encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+ encoder_hidden_states_query_proj = attn.head_to_batch_dim(encoder_hidden_states_query_proj).contiguous()
+ encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj).contiguous()
+ encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj).contiguous()
+
+ if attn.norm_added_q is not None:
+ encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+ if attn.norm_added_k is not None:
+ encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+ query = torch.cat([query, encoder_hidden_states_query_proj], dim=1)
+ key = torch.cat([key, encoder_hidden_states_key_proj], dim=1)
+ value = torch.cat([value, encoder_hidden_states_value_proj], dim=1)
+
+ hidden_states = xformers.ops.memory_efficient_attention(
+ query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale
+ )
+ hidden_states = hidden_states.to(query.dtype)
+ hidden_states = attn.batch_to_head_dim(hidden_states)
+
+ if encoder_hidden_states is not None:
+ # Split the attention outputs.
+ hidden_states, encoder_hidden_states = (
+ hidden_states[:, : residual.shape[1]],
+ hidden_states[:, residual.shape[1] :],
+ )
+ if not attn.context_pre_only:
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if encoder_hidden_states is not None:
+ return hidden_states, encoder_hidden_states
+ else:
+ return hidden_states
+
+
+class AllegroAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is
+ used in the Allegro model. It applies a normalization layer and rotary embedding on the query and key vector.
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "AllegroAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ residual = hidden_states
+
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ # Apply RoPE if needed
+ if image_rotary_emb is not None and not attn.is_cross_attention:
+ from .embeddings import apply_rotary_emb_allegro
+
+ query = apply_rotary_emb_allegro(query, image_rotary_emb[0], image_rotary_emb[1])
+ key = apply_rotary_emb_allegro(key, image_rotary_emb[0], image_rotary_emb[1])
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class AuraFlowAttnProcessor2_0:
+ """Attention processor used typically in processing Aura Flow."""
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention") and is_torch_version("<", "2.1"):
+ raise ImportError(
+ "AuraFlowAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to at least 2.1 or above as we use `scale` in `F.scaled_dot_product_attention()`. "
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ *args,
+ **kwargs,
+ ) -> torch.FloatTensor:
+ batch_size = hidden_states.shape[0]
+
+ # `sample` projections.
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ # `context` projections.
+ if encoder_hidden_states is not None:
+ encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+ encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+ encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+ # Reshape.
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+ query = query.view(batch_size, -1, attn.heads, head_dim)
+ key = key.view(batch_size, -1, attn.heads, head_dim)
+ value = value.view(batch_size, -1, attn.heads, head_dim)
+
+ # Apply QK norm.
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # Concatenate the projections.
+ if encoder_hidden_states is not None:
+ encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ )
+ encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(batch_size, -1, attn.heads, head_dim)
+ encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ )
+
+ if attn.norm_added_q is not None:
+ encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+ if attn.norm_added_k is not None:
+ encoder_hidden_states_key_proj = attn.norm_added_q(encoder_hidden_states_key_proj)
+
+ query = torch.cat([encoder_hidden_states_query_proj, query], dim=1)
+ key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)
+ value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)
+
+ query = query.transpose(1, 2)
+ key = key.transpose(1, 2)
+ value = value.transpose(1, 2)
+
+ # Attention.
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, dropout_p=0.0, scale=attn.scale, is_causal=False
+ )
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # Split the attention outputs.
+ if encoder_hidden_states is not None:
+ hidden_states, encoder_hidden_states = (
+ hidden_states[:, encoder_hidden_states.shape[1] :],
+ hidden_states[:, : encoder_hidden_states.shape[1]],
+ )
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+ if encoder_hidden_states is not None:
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ if encoder_hidden_states is not None:
+ return hidden_states, encoder_hidden_states
+ else:
+ return hidden_states
+
+
+class FusedAuraFlowAttnProcessor2_0:
+ """Attention processor used typically in processing Aura Flow with fused projections."""
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention") and is_torch_version("<", "2.1"):
+ raise ImportError(
+ "FusedAuraFlowAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to at least 2.1 or above as we use `scale` in `F.scaled_dot_product_attention()`. "
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ *args,
+ **kwargs,
+ ) -> torch.FloatTensor:
+ batch_size = hidden_states.shape[0]
+
+ # `sample` projections.
+ qkv = attn.to_qkv(hidden_states)
+ split_size = qkv.shape[-1] // 3
+ query, key, value = torch.split(qkv, split_size, dim=-1)
+
+ # `context` projections.
+ if encoder_hidden_states is not None:
+ encoder_qkv = attn.to_added_qkv(encoder_hidden_states)
+ split_size = encoder_qkv.shape[-1] // 3
+ (
+ encoder_hidden_states_query_proj,
+ encoder_hidden_states_key_proj,
+ encoder_hidden_states_value_proj,
+ ) = torch.split(encoder_qkv, split_size, dim=-1)
+
+ # Reshape.
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+ query = query.view(batch_size, -1, attn.heads, head_dim)
+ key = key.view(batch_size, -1, attn.heads, head_dim)
+ value = value.view(batch_size, -1, attn.heads, head_dim)
+
+ # Apply QK norm.
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # Concatenate the projections.
+ if encoder_hidden_states is not None:
+ encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ )
+ encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(batch_size, -1, attn.heads, head_dim)
+ encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ )
+
+ if attn.norm_added_q is not None:
+ encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+ if attn.norm_added_k is not None:
+ encoder_hidden_states_key_proj = attn.norm_added_q(encoder_hidden_states_key_proj)
+
+ query = torch.cat([encoder_hidden_states_query_proj, query], dim=1)
+ key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)
+ value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)
+
+ query = query.transpose(1, 2)
+ key = key.transpose(1, 2)
+ value = value.transpose(1, 2)
+
+ # Attention.
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, dropout_p=0.0, scale=attn.scale, is_causal=False
+ )
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # Split the attention outputs.
+ if encoder_hidden_states is not None:
+ hidden_states, encoder_hidden_states = (
+ hidden_states[:, encoder_hidden_states.shape[1] :],
+ hidden_states[:, : encoder_hidden_states.shape[1]],
+ )
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+ if encoder_hidden_states is not None:
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ if encoder_hidden_states is not None:
+ return hidden_states, encoder_hidden_states
+ else:
+ return hidden_states
+
+
+class FluxAttnProcessor2_0:
+ """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError("FluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.FloatTensor:
+ batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+ # `sample` projections.
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+ if encoder_hidden_states is not None:
+ # `context` projections.
+ encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+ encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+ encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+ encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+
+ if attn.norm_added_q is not None:
+ encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+ if attn.norm_added_k is not None:
+ encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+ # attention
+ query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+ key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+ value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+ if image_rotary_emb is not None:
+ from .embeddings import apply_rotary_emb
+
+ query = apply_rotary_emb(query, image_rotary_emb)
+ key = apply_rotary_emb(key, image_rotary_emb)
+
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if encoder_hidden_states is not None:
+ encoder_hidden_states, hidden_states = (
+ hidden_states[:, : encoder_hidden_states.shape[1]],
+ hidden_states[:, encoder_hidden_states.shape[1] :],
+ )
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ return hidden_states, encoder_hidden_states
+ else:
+ return hidden_states
+
+
+class FluxAttnProcessor2_0_NPU:
+ """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "FluxAttnProcessor2_0_NPU requires PyTorch 2.0 and torch NPU, to use it, please upgrade PyTorch to 2.0 and install torch NPU"
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.FloatTensor:
+ batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+ # `sample` projections.
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+ if encoder_hidden_states is not None:
+ # `context` projections.
+ encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+ encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+ encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+ encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+
+ if attn.norm_added_q is not None:
+ encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+ if attn.norm_added_k is not None:
+ encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+ # attention
+ query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+ key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+ value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+ if image_rotary_emb is not None:
+ from .embeddings import apply_rotary_emb
+
+ query = apply_rotary_emb(query, image_rotary_emb)
+ key = apply_rotary_emb(key, image_rotary_emb)
+
+ if query.dtype in (torch.float16, torch.bfloat16):
+ hidden_states = torch_npu.npu_fusion_attention(
+ query,
+ key,
+ value,
+ attn.heads,
+ input_layout="BNSD",
+ pse=None,
+ scale=1.0 / math.sqrt(query.shape[-1]),
+ pre_tockens=65536,
+ next_tockens=65536,
+ keep_prob=1.0,
+ sync=False,
+ inner_precise=0,
+ )[0]
+ else:
+ hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if encoder_hidden_states is not None:
+ encoder_hidden_states, hidden_states = (
+ hidden_states[:, : encoder_hidden_states.shape[1]],
+ hidden_states[:, encoder_hidden_states.shape[1] :],
+ )
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ return hidden_states, encoder_hidden_states
+ else:
+ return hidden_states
+
+
+class FusedFluxAttnProcessor2_0:
+ """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "FusedFluxAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.FloatTensor:
+ batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+ # `sample` projections.
+ qkv = attn.to_qkv(hidden_states)
+ split_size = qkv.shape[-1] // 3
+ query, key, value = torch.split(qkv, split_size, dim=-1)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+ # `context` projections.
+ if encoder_hidden_states is not None:
+ encoder_qkv = attn.to_added_qkv(encoder_hidden_states)
+ split_size = encoder_qkv.shape[-1] // 3
+ (
+ encoder_hidden_states_query_proj,
+ encoder_hidden_states_key_proj,
+ encoder_hidden_states_value_proj,
+ ) = torch.split(encoder_qkv, split_size, dim=-1)
+
+ encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+
+ if attn.norm_added_q is not None:
+ encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+ if attn.norm_added_k is not None:
+ encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+ # attention
+ query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+ key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+ value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+ if image_rotary_emb is not None:
+ from .embeddings import apply_rotary_emb
+
+ query = apply_rotary_emb(query, image_rotary_emb)
+ key = apply_rotary_emb(key, image_rotary_emb)
+
+ hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if encoder_hidden_states is not None:
+ encoder_hidden_states, hidden_states = (
+ hidden_states[:, : encoder_hidden_states.shape[1]],
+ hidden_states[:, encoder_hidden_states.shape[1] :],
+ )
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ return hidden_states, encoder_hidden_states
+ else:
+ return hidden_states
+
+
+class FusedFluxAttnProcessor2_0_NPU:
+ """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "FluxAttnProcessor2_0_NPU requires PyTorch 2.0 and torch NPU, to use it, please upgrade PyTorch to 2.0, and install torch NPU"
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.FloatTensor:
+ batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+ # `sample` projections.
+ qkv = attn.to_qkv(hidden_states)
+ split_size = qkv.shape[-1] // 3
+ query, key, value = torch.split(qkv, split_size, dim=-1)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+ # `context` projections.
+ if encoder_hidden_states is not None:
+ encoder_qkv = attn.to_added_qkv(encoder_hidden_states)
+ split_size = encoder_qkv.shape[-1] // 3
+ (
+ encoder_hidden_states_query_proj,
+ encoder_hidden_states_key_proj,
+ encoder_hidden_states_value_proj,
+ ) = torch.split(encoder_qkv, split_size, dim=-1)
+
+ encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+
+ if attn.norm_added_q is not None:
+ encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+ if attn.norm_added_k is not None:
+ encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+ # attention
+ query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+ key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+ value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+ if image_rotary_emb is not None:
+ from .embeddings import apply_rotary_emb
+
+ query = apply_rotary_emb(query, image_rotary_emb)
+ key = apply_rotary_emb(key, image_rotary_emb)
+
+ if query.dtype in (torch.float16, torch.bfloat16):
+ hidden_states = torch_npu.npu_fusion_attention(
+ query,
+ key,
+ value,
+ attn.heads,
+ input_layout="BNSD",
+ pse=None,
+ scale=1.0 / math.sqrt(query.shape[-1]),
+ pre_tockens=65536,
+ next_tockens=65536,
+ keep_prob=1.0,
+ sync=False,
+ inner_precise=0,
+ )[0]
+ else:
+ hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if encoder_hidden_states is not None:
+ encoder_hidden_states, hidden_states = (
+ hidden_states[:, : encoder_hidden_states.shape[1]],
+ hidden_states[:, encoder_hidden_states.shape[1] :],
+ )
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ return hidden_states, encoder_hidden_states
+ else:
+ return hidden_states
+
+
+class FluxIPAdapterJointAttnProcessor2_0(torch.nn.Module):
+ """Flux Attention processor for IP-Adapter."""
+
+ def __init__(
+ self, hidden_size: int, cross_attention_dim: int, num_tokens=(4,), scale=1.0, device=None, dtype=None
+ ):
+ super().__init__()
+
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ f"{self.__class__.__name__} requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+
+ self.hidden_size = hidden_size
+ self.cross_attention_dim = cross_attention_dim
+
+ if not isinstance(num_tokens, (tuple, list)):
+ num_tokens = [num_tokens]
+
+ if not isinstance(scale, list):
+ scale = [scale] * len(num_tokens)
+ if len(scale) != len(num_tokens):
+ raise ValueError("`scale` should be a list of integers with the same length as `num_tokens`.")
+ self.scale = scale
+
+ self.to_k_ip = nn.ModuleList(
+ [
+ nn.Linear(cross_attention_dim, hidden_size, bias=True, device=device, dtype=dtype)
+ for _ in range(len(num_tokens))
+ ]
+ )
+ self.to_v_ip = nn.ModuleList(
+ [
+ nn.Linear(cross_attention_dim, hidden_size, bias=True, device=device, dtype=dtype)
+ for _ in range(len(num_tokens))
+ ]
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ip_hidden_states: Optional[List[torch.Tensor]] = None,
+ ip_adapter_masks: Optional[torch.Tensor] = None,
+ ) -> torch.FloatTensor:
+ batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+ # `sample` projections.
+ hidden_states_query_proj = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ hidden_states_query_proj = hidden_states_query_proj.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ hidden_states_query_proj = attn.norm_q(hidden_states_query_proj)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+ if encoder_hidden_states is not None:
+ # `context` projections.
+ encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+ encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+ encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+ encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+
+ if attn.norm_added_q is not None:
+ encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+ if attn.norm_added_k is not None:
+ encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+ # attention
+ query = torch.cat([encoder_hidden_states_query_proj, hidden_states_query_proj], dim=2)
+ key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+ value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+ if image_rotary_emb is not None:
+ from .embeddings import apply_rotary_emb
+
+ query = apply_rotary_emb(query, image_rotary_emb)
+ key = apply_rotary_emb(key, image_rotary_emb)
+
+ hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if encoder_hidden_states is not None:
+ encoder_hidden_states, hidden_states = (
+ hidden_states[:, : encoder_hidden_states.shape[1]],
+ hidden_states[:, encoder_hidden_states.shape[1] :],
+ )
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ # IP-adapter
+ ip_query = hidden_states_query_proj
+ ip_attn_output = torch.zeros_like(hidden_states)
+
+ for current_ip_hidden_states, scale, to_k_ip, to_v_ip in zip(
+ ip_hidden_states, self.scale, self.to_k_ip, self.to_v_ip
+ ):
+ ip_key = to_k_ip(current_ip_hidden_states)
+ ip_value = to_v_ip(current_ip_hidden_states)
+
+ ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ current_ip_hidden_states = F.scaled_dot_product_attention(
+ ip_query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
+ )
+ current_ip_hidden_states = current_ip_hidden_states.transpose(1, 2).reshape(
+ batch_size, -1, attn.heads * head_dim
+ )
+ current_ip_hidden_states = current_ip_hidden_states.to(ip_query.dtype)
+ ip_attn_output += scale * current_ip_hidden_states
+
+ return hidden_states, encoder_hidden_states, ip_attn_output
+ else:
+ return hidden_states
+
+
+class CogVideoXAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on
+ query and key vectors, but does not include spatial normalization.
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ text_seq_length = encoder_hidden_states.size(1)
+
+ hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+ batch_size, sequence_length, _ = hidden_states.shape
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # Apply RoPE if needed
+ if image_rotary_emb is not None:
+ from .embeddings import apply_rotary_emb
+
+ query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
+ if not attn.is_cross_attention:
+ key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
+
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ encoder_hidden_states, hidden_states = hidden_states.split(
+ [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+ )
+ return hidden_states, encoder_hidden_states
+
+
+class FusedCogVideoXAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention for the CogVideoX model. It applies a rotary embedding on
+ query and key vectors, but does not include spatial normalization.
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError("CogVideoXAttnProcessor requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ text_seq_length = encoder_hidden_states.size(1)
+
+ hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ qkv = attn.to_qkv(hidden_states)
+ split_size = qkv.shape[-1] // 3
+ query, key, value = torch.split(qkv, split_size, dim=-1)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # Apply RoPE if needed
+ if image_rotary_emb is not None:
+ from .embeddings import apply_rotary_emb
+
+ query[:, :, text_seq_length:] = apply_rotary_emb(query[:, :, text_seq_length:], image_rotary_emb)
+ if not attn.is_cross_attention:
+ key[:, :, text_seq_length:] = apply_rotary_emb(key[:, :, text_seq_length:], image_rotary_emb)
+
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ encoder_hidden_states, hidden_states = hidden_states.split(
+ [text_seq_length, hidden_states.size(1) - text_seq_length], dim=1
+ )
+ return hidden_states, encoder_hidden_states
+
+
+class XFormersAttnAddedKVProcessor:
+ r"""
+ Processor for implementing memory efficient attention using xFormers.
+
+ Args:
+ attention_op (`Callable`, *optional*, defaults to `None`):
+ The base
+ [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to
+ use as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best
+ operator.
+ """
+
+ def __init__(self, attention_op: Optional[Callable] = None):
+ self.attention_op = attention_op
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ residual = hidden_states
+ hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)
+ batch_size, sequence_length, _ = hidden_states.shape
+
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+ query = attn.head_to_batch_dim(query)
+
+ encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+ encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+ encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj)
+ encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj)
+
+ if not attn.only_cross_attention:
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+ key = attn.head_to_batch_dim(key)
+ value = attn.head_to_batch_dim(value)
+ key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)
+ value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)
+ else:
+ key = encoder_hidden_states_key_proj
+ value = encoder_hidden_states_value_proj
+
+ hidden_states = xformers.ops.memory_efficient_attention(
+ query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale
+ )
+ hidden_states = hidden_states.to(query.dtype)
+ hidden_states = attn.batch_to_head_dim(hidden_states)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape)
+ hidden_states = hidden_states + residual
+
+ return hidden_states
+
+
+class XFormersAttnProcessor:
+ r"""
+ Processor for implementing memory efficient attention using xFormers.
+
+ Args:
+ attention_op (`Callable`, *optional*, defaults to `None`):
+ The base
+ [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to
+ use as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best
+ operator.
+ """
+
+ def __init__(self, attention_op: Optional[Callable] = None):
+ self.attention_op = attention_op
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.Tensor:
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+ deprecate("scale", "1.0.0", deprecation_message)
+
+ residual = hidden_states
+
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, key_tokens, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ attention_mask = attn.prepare_attention_mask(attention_mask, key_tokens, batch_size)
+ if attention_mask is not None:
+ # expand our mask's singleton query_tokens dimension:
+ # [batch*heads, 1, key_tokens] ->
+ # [batch*heads, query_tokens, key_tokens]
+ # so that it can be added as a bias onto the attention scores that xformers computes:
+ # [batch*heads, query_tokens, key_tokens]
+ # we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
+ _, query_tokens, _ = hidden_states.shape
+ attention_mask = attention_mask.expand(-1, query_tokens, -1)
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ query = attn.head_to_batch_dim(query).contiguous()
+ key = attn.head_to_batch_dim(key).contiguous()
+ value = attn.head_to_batch_dim(value).contiguous()
+
+ hidden_states = xformers.ops.memory_efficient_attention(
+ query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale
+ )
+ hidden_states = hidden_states.to(query.dtype)
+ hidden_states = attn.batch_to_head_dim(hidden_states)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class AttnProcessorNPU:
+ r"""
+ Processor for implementing flash attention using torch_npu. Torch_npu supports only fp16 and bf16 data types. If
+ fp32 is used, F.scaled_dot_product_attention will be used for computation, but the acceleration effect on NPU is
+ not significant.
+
+ """
+
+ def __init__(self):
+ if not is_torch_npu_available():
+ raise ImportError("AttnProcessorNPU requires torch_npu extensions and is supported only on npu devices.")
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.Tensor:
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+ deprecate("scale", "1.0.0", deprecation_message)
+
+ residual = hidden_states
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+ attention_mask = attention_mask.repeat(1, 1, hidden_states.shape[1], 1)
+ if attention_mask.dtype == torch.bool:
+ attention_mask = torch.logical_not(attention_mask.bool())
+ else:
+ attention_mask = attention_mask.bool()
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ if query.dtype in (torch.float16, torch.bfloat16):
+ hidden_states = torch_npu.npu_fusion_attention(
+ query,
+ key,
+ value,
+ attn.heads,
+ input_layout="BNSD",
+ pse=None,
+ atten_mask=attention_mask,
+ scale=1.0 / math.sqrt(query.shape[-1]),
+ pre_tockens=65536,
+ next_tockens=65536,
+ keep_prob=1.0,
+ sync=False,
+ inner_precise=0,
+ )[0]
+ else:
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class AttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.Tensor:
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+ deprecate("scale", "1.0.0", deprecation_message)
+
+ residual = hidden_states
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class XLAFlashAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention with pallas flash attention kernel if using `torch_xla`.
+ """
+
+ def __init__(self, partition_spec: Optional[Tuple[Optional[str], ...]] = None):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "XLAFlashAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+ if is_torch_xla_version("<", "2.3"):
+ raise ImportError("XLA flash attention requires torch_xla version >= 2.3.")
+ if is_spmd() and is_torch_xla_version("<", "2.4"):
+ raise ImportError("SPMD support for XLA flash attention needs torch_xla version >= 2.4.")
+ self.partition_spec = partition_spec
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.Tensor:
+ residual = hidden_states
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ if all(tensor.shape[2] >= 4096 for tensor in [query, key, value]):
+ if attention_mask is not None:
+ attention_mask = attention_mask.view(batch_size, 1, 1, attention_mask.shape[-1])
+ # Convert mask to float and replace 0s with -inf and 1s with 0
+ attention_mask = (
+ attention_mask.float()
+ .masked_fill(attention_mask == 0, float("-inf"))
+ .masked_fill(attention_mask == 1, float(0.0))
+ )
+
+ # Apply attention mask to key
+ key = key + attention_mask
+ query /= math.sqrt(query.shape[3])
+ partition_spec = self.partition_spec if is_spmd() else None
+ hidden_states = flash_attention(query, key, value, causal=False, partition_spec=partition_spec)
+ else:
+ logger.warning(
+ "Unable to use the flash attention pallas kernel API call due to QKV sequence length < 4096."
+ )
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class XLAFluxFlashAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention with pallas flash attention kernel if using `torch_xla`.
+ """
+
+ def __init__(self, partition_spec: Optional[Tuple[Optional[str], ...]] = None):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "XLAFlashAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+ if is_torch_xla_version("<", "2.3"):
+ raise ImportError("XLA flash attention requires torch_xla version >= 2.3.")
+ if is_spmd() and is_torch_xla_version("<", "2.4"):
+ raise ImportError("SPMD support for XLA flash attention needs torch_xla version >= 2.4.")
+ self.partition_spec = partition_spec
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.FloatTensor:
+ batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+ # `sample` projections.
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+ if encoder_hidden_states is not None:
+ # `context` projections.
+ encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+ encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+ encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+ encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+
+ if attn.norm_added_q is not None:
+ encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+ if attn.norm_added_k is not None:
+ encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+ # attention
+ query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+ key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+ value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+ if image_rotary_emb is not None:
+ from .embeddings import apply_rotary_emb
+
+ query = apply_rotary_emb(query, image_rotary_emb)
+ key = apply_rotary_emb(key, image_rotary_emb)
+
+ query /= math.sqrt(head_dim)
+ hidden_states = flash_attention(query, key, value, causal=False)
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if encoder_hidden_states is not None:
+ encoder_hidden_states, hidden_states = (
+ hidden_states[:, : encoder_hidden_states.shape[1]],
+ hidden_states[:, encoder_hidden_states.shape[1] :],
+ )
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ return hidden_states, encoder_hidden_states
+ else:
+ return hidden_states
+
+
+class MochiVaeAttnProcessor2_0:
+ r"""
+ Attention processor used in Mochi VAE.
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ residual = hidden_states
+ is_single_frame = hidden_states.shape[1] == 1
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if is_single_frame:
+ hidden_states = attn.to_v(hidden_states)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+ return hidden_states
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=attn.is_causal
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class StableAudioAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is
+ used in the Stable Audio model. It applies rotary embedding on query and key vector, and allows MHA, GQA or MQA.
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "StableAudioAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+
+ def apply_partial_rotary_emb(
+ self,
+ x: torch.Tensor,
+ freqs_cis: Tuple[torch.Tensor],
+ ) -> torch.Tensor:
+ from .embeddings import apply_rotary_emb
+
+ rot_dim = freqs_cis[0].shape[-1]
+ x_to_rotate, x_unrotated = x[..., :rot_dim], x[..., rot_dim:]
+
+ x_rotated = apply_rotary_emb(x_to_rotate, freqs_cis, use_real=True, use_real_unbind_dim=-2)
+
+ out = torch.cat((x_rotated, x_unrotated), dim=-1)
+ return out
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ from .embeddings import apply_rotary_emb
+
+ residual = hidden_states
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ head_dim = query.shape[-1] // attn.heads
+ kv_heads = key.shape[-1] // head_dim
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ key = key.view(batch_size, -1, kv_heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, kv_heads, head_dim).transpose(1, 2)
+
+ if kv_heads != attn.heads:
+ # if GQA or MQA, repeat the key/value heads to reach the number of query heads.
+ heads_per_kv_head = attn.heads // kv_heads
+ key = torch.repeat_interleave(key, heads_per_kv_head, dim=1, output_size=key.shape[1] * heads_per_kv_head)
+ value = torch.repeat_interleave(
+ value, heads_per_kv_head, dim=1, output_size=value.shape[1] * heads_per_kv_head
+ )
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # Apply RoPE if needed
+ if rotary_emb is not None:
+ query_dtype = query.dtype
+ key_dtype = key.dtype
+ query = query.to(torch.float32)
+ key = key.to(torch.float32)
+
+ rot_dim = rotary_emb[0].shape[-1]
+ query_to_rotate, query_unrotated = query[..., :rot_dim], query[..., rot_dim:]
+ query_rotated = apply_rotary_emb(query_to_rotate, rotary_emb, use_real=True, use_real_unbind_dim=-2)
+
+ query = torch.cat((query_rotated, query_unrotated), dim=-1)
+
+ if not attn.is_cross_attention:
+ key_to_rotate, key_unrotated = key[..., :rot_dim], key[..., rot_dim:]
+ key_rotated = apply_rotary_emb(key_to_rotate, rotary_emb, use_real=True, use_real_unbind_dim=-2)
+
+ key = torch.cat((key_rotated, key_unrotated), dim=-1)
+
+ query = query.to(query_dtype)
+ key = key.to(key_dtype)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class HunyuanAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is
+ used in the HunyuanDiT model. It applies a s normalization layer and rotary embedding on query and key vector.
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ from .embeddings import apply_rotary_emb
+
+ residual = hidden_states
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # Apply RoPE if needed
+ if image_rotary_emb is not None:
+ query = apply_rotary_emb(query, image_rotary_emb)
+ if not attn.is_cross_attention:
+ key = apply_rotary_emb(key, image_rotary_emb)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class FusedHunyuanAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0) with fused
+ projection layers. This is used in the HunyuanDiT model. It applies a s normalization layer and rotary embedding on
+ query and key vector.
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "FusedHunyuanAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ from .embeddings import apply_rotary_emb
+
+ residual = hidden_states
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ if encoder_hidden_states is None:
+ qkv = attn.to_qkv(hidden_states)
+ split_size = qkv.shape[-1] // 3
+ query, key, value = torch.split(qkv, split_size, dim=-1)
+ else:
+ if attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+ query = attn.to_q(hidden_states)
+
+ kv = attn.to_kv(encoder_hidden_states)
+ split_size = kv.shape[-1] // 2
+ key, value = torch.split(kv, split_size, dim=-1)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # Apply RoPE if needed
+ if image_rotary_emb is not None:
+ query = apply_rotary_emb(query, image_rotary_emb)
+ if not attn.is_cross_attention:
+ key = apply_rotary_emb(key, image_rotary_emb)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class PAGHunyuanAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is
+ used in the HunyuanDiT model. It applies a normalization layer and rotary embedding on query and key vector. This
+ variant of the processor employs [Pertubed Attention Guidance](https://arxiv.org/abs/2403.17377).
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "PAGHunyuanAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ from .embeddings import apply_rotary_emb
+
+ residual = hidden_states
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ # chunk
+ hidden_states_org, hidden_states_ptb = hidden_states.chunk(2)
+
+ # 1. Original Path
+ batch_size, sequence_length, _ = (
+ hidden_states_org.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if attn.group_norm is not None:
+ hidden_states_org = attn.group_norm(hidden_states_org.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states_org)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states_org
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # Apply RoPE if needed
+ if image_rotary_emb is not None:
+ query = apply_rotary_emb(query, image_rotary_emb)
+ if not attn.is_cross_attention:
+ key = apply_rotary_emb(key, image_rotary_emb)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states_org = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states_org = hidden_states_org.to(query.dtype)
+
+ # linear proj
+ hidden_states_org = attn.to_out[0](hidden_states_org)
+ # dropout
+ hidden_states_org = attn.to_out[1](hidden_states_org)
+
+ if input_ndim == 4:
+ hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ # 2. Perturbed Path
+ if attn.group_norm is not None:
+ hidden_states_ptb = attn.group_norm(hidden_states_ptb.transpose(1, 2)).transpose(1, 2)
+
+ hidden_states_ptb = attn.to_v(hidden_states_ptb)
+ hidden_states_ptb = hidden_states_ptb.to(query.dtype)
+
+ # linear proj
+ hidden_states_ptb = attn.to_out[0](hidden_states_ptb)
+ # dropout
+ hidden_states_ptb = attn.to_out[1](hidden_states_ptb)
+
+ if input_ndim == 4:
+ hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ # cat
+ hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class PAGCFGHunyuanAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is
+ used in the HunyuanDiT model. It applies a normalization layer and rotary embedding on query and key vector. This
+ variant of the processor employs [Pertubed Attention Guidance](https://arxiv.org/abs/2403.17377).
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "PAGCFGHunyuanAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ image_rotary_emb: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ from .embeddings import apply_rotary_emb
+
+ residual = hidden_states
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ # chunk
+ hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3)
+ hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org])
+
+ # 1. Original Path
+ batch_size, sequence_length, _ = (
+ hidden_states_org.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if attn.group_norm is not None:
+ hidden_states_org = attn.group_norm(hidden_states_org.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states_org)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states_org
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # Apply RoPE if needed
+ if image_rotary_emb is not None:
+ query = apply_rotary_emb(query, image_rotary_emb)
+ if not attn.is_cross_attention:
+ key = apply_rotary_emb(key, image_rotary_emb)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states_org = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states_org = hidden_states_org.to(query.dtype)
+
+ # linear proj
+ hidden_states_org = attn.to_out[0](hidden_states_org)
+ # dropout
+ hidden_states_org = attn.to_out[1](hidden_states_org)
+
+ if input_ndim == 4:
+ hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ # 2. Perturbed Path
+ if attn.group_norm is not None:
+ hidden_states_ptb = attn.group_norm(hidden_states_ptb.transpose(1, 2)).transpose(1, 2)
+
+ hidden_states_ptb = attn.to_v(hidden_states_ptb)
+ hidden_states_ptb = hidden_states_ptb.to(query.dtype)
+
+ # linear proj
+ hidden_states_ptb = attn.to_out[0](hidden_states_ptb)
+ # dropout
+ hidden_states_ptb = attn.to_out[1](hidden_states_ptb)
+
+ if input_ndim == 4:
+ hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ # cat
+ hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class LuminaAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). This is
+ used in the LuminaNextDiT model. It applies a s normalization layer and rotary embedding on query and key vector.
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: torch.Tensor,
+ attention_mask: Optional[torch.Tensor] = None,
+ query_rotary_emb: Optional[torch.Tensor] = None,
+ key_rotary_emb: Optional[torch.Tensor] = None,
+ base_sequence_length: Optional[int] = None,
+ ) -> torch.Tensor:
+ from .embeddings import apply_rotary_emb
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = hidden_states.shape
+
+ # Get Query-Key-Value Pair
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ query_dim = query.shape[-1]
+ inner_dim = key.shape[-1]
+ head_dim = query_dim // attn.heads
+ dtype = query.dtype
+
+ # Get key-value heads
+ kv_heads = inner_dim // head_dim
+
+ # Apply Query-Key Norm if needed
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ query = query.view(batch_size, -1, attn.heads, head_dim)
+
+ key = key.view(batch_size, -1, kv_heads, head_dim)
+ value = value.view(batch_size, -1, kv_heads, head_dim)
+
+ # Apply RoPE if needed
+ if query_rotary_emb is not None:
+ query = apply_rotary_emb(query, query_rotary_emb, use_real=False)
+ if key_rotary_emb is not None:
+ key = apply_rotary_emb(key, key_rotary_emb, use_real=False)
+
+ query, key = query.to(dtype), key.to(dtype)
+
+ # Apply proportional attention if true
+ if key_rotary_emb is None:
+ softmax_scale = None
+ else:
+ if base_sequence_length is not None:
+ softmax_scale = math.sqrt(math.log(sequence_length, base_sequence_length)) * attn.scale
+ else:
+ softmax_scale = attn.scale
+
+ # perform Grouped-qurey Attention (GQA)
+ n_rep = attn.heads // kv_heads
+ if n_rep >= 1:
+ key = key.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3)
+ value = value.unsqueeze(3).repeat(1, 1, 1, n_rep, 1).flatten(2, 3)
+
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.bool().view(batch_size, 1, 1, -1)
+ attention_mask = attention_mask.expand(-1, attn.heads, sequence_length, -1)
+
+ query = query.transpose(1, 2)
+ key = key.transpose(1, 2)
+ value = value.transpose(1, 2)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, scale=softmax_scale
+ )
+ hidden_states = hidden_states.transpose(1, 2).to(dtype)
+
+ return hidden_states
+
+
+class FusedAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0). It uses
+ fused projection layers. For self-attention modules, all projection matrices (i.e., query, key, value) are fused.
+ For cross-attention modules, key and value projection matrices are fused.
+
+
+
+ This API is currently 🧪 experimental in nature and can change in future.
+
+
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "FusedAttnProcessor2_0 requires at least PyTorch 2.0, to use it. Please upgrade PyTorch to > 2.0."
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ *args,
+ **kwargs,
+ ) -> torch.Tensor:
+ if len(args) > 0 or kwargs.get("scale", None) is not None:
+ deprecation_message = "The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`."
+ deprecate("scale", "1.0.0", deprecation_message)
+
+ residual = hidden_states
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ if encoder_hidden_states is None:
+ qkv = attn.to_qkv(hidden_states)
+ split_size = qkv.shape[-1] // 3
+ query, key, value = torch.split(qkv, split_size, dim=-1)
+ else:
+ if attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+ query = attn.to_q(hidden_states)
+
+ kv = attn.to_kv(encoder_hidden_states)
+ split_size = kv.shape[-1] // 2
+ key, value = torch.split(kv, split_size, dim=-1)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class CustomDiffusionXFormersAttnProcessor(nn.Module):
+ r"""
+ Processor for implementing memory efficient attention using xFormers for the Custom Diffusion method.
+
+ Args:
+ train_kv (`bool`, defaults to `True`):
+ Whether to newly train the key and value matrices corresponding to the text features.
+ train_q_out (`bool`, defaults to `True`):
+ Whether to newly train query matrices corresponding to the latent image features.
+ hidden_size (`int`, *optional*, defaults to `None`):
+ The hidden size of the attention layer.
+ cross_attention_dim (`int`, *optional*, defaults to `None`):
+ The number of channels in the `encoder_hidden_states`.
+ out_bias (`bool`, defaults to `True`):
+ Whether to include the bias parameter in `train_q_out`.
+ dropout (`float`, *optional*, defaults to 0.0):
+ The dropout probability to use.
+ attention_op (`Callable`, *optional*, defaults to `None`):
+ The base
+ [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to use
+ as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best operator.
+ """
+
+ def __init__(
+ self,
+ train_kv: bool = True,
+ train_q_out: bool = False,
+ hidden_size: Optional[int] = None,
+ cross_attention_dim: Optional[int] = None,
+ out_bias: bool = True,
+ dropout: float = 0.0,
+ attention_op: Optional[Callable] = None,
+ ):
+ super().__init__()
+ self.train_kv = train_kv
+ self.train_q_out = train_q_out
+
+ self.hidden_size = hidden_size
+ self.cross_attention_dim = cross_attention_dim
+ self.attention_op = attention_op
+
+ # `_custom_diffusion` id for easy serialization and loading.
+ if self.train_kv:
+ self.to_k_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
+ self.to_v_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
+ if self.train_q_out:
+ self.to_q_custom_diffusion = nn.Linear(hidden_size, hidden_size, bias=False)
+ self.to_out_custom_diffusion = nn.ModuleList([])
+ self.to_out_custom_diffusion.append(nn.Linear(hidden_size, hidden_size, bias=out_bias))
+ self.to_out_custom_diffusion.append(nn.Dropout(dropout))
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+ if self.train_q_out:
+ query = self.to_q_custom_diffusion(hidden_states).to(attn.to_q.weight.dtype)
+ else:
+ query = attn.to_q(hidden_states.to(attn.to_q.weight.dtype))
+
+ if encoder_hidden_states is None:
+ crossattn = False
+ encoder_hidden_states = hidden_states
+ else:
+ crossattn = True
+ if attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ if self.train_kv:
+ key = self.to_k_custom_diffusion(encoder_hidden_states.to(self.to_k_custom_diffusion.weight.dtype))
+ value = self.to_v_custom_diffusion(encoder_hidden_states.to(self.to_v_custom_diffusion.weight.dtype))
+ key = key.to(attn.to_q.weight.dtype)
+ value = value.to(attn.to_q.weight.dtype)
+ else:
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ if crossattn:
+ detach = torch.ones_like(key)
+ detach[:, :1, :] = detach[:, :1, :] * 0.0
+ key = detach * key + (1 - detach) * key.detach()
+ value = detach * value + (1 - detach) * value.detach()
+
+ query = attn.head_to_batch_dim(query).contiguous()
+ key = attn.head_to_batch_dim(key).contiguous()
+ value = attn.head_to_batch_dim(value).contiguous()
+
+ hidden_states = xformers.ops.memory_efficient_attention(
+ query, key, value, attn_bias=attention_mask, op=self.attention_op, scale=attn.scale
+ )
+ hidden_states = hidden_states.to(query.dtype)
+ hidden_states = attn.batch_to_head_dim(hidden_states)
+
+ if self.train_q_out:
+ # linear proj
+ hidden_states = self.to_out_custom_diffusion[0](hidden_states)
+ # dropout
+ hidden_states = self.to_out_custom_diffusion[1](hidden_states)
+ else:
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ return hidden_states
+
+
+class CustomDiffusionAttnProcessor2_0(nn.Module):
+ r"""
+ Processor for implementing attention for the Custom Diffusion method using PyTorch 2.0’s memory-efficient scaled
+ dot-product attention.
+
+ Args:
+ train_kv (`bool`, defaults to `True`):
+ Whether to newly train the key and value matrices corresponding to the text features.
+ train_q_out (`bool`, defaults to `True`):
+ Whether to newly train query matrices corresponding to the latent image features.
+ hidden_size (`int`, *optional*, defaults to `None`):
+ The hidden size of the attention layer.
+ cross_attention_dim (`int`, *optional*, defaults to `None`):
+ The number of channels in the `encoder_hidden_states`.
+ out_bias (`bool`, defaults to `True`):
+ Whether to include the bias parameter in `train_q_out`.
+ dropout (`float`, *optional*, defaults to 0.0):
+ The dropout probability to use.
+ """
+
+ def __init__(
+ self,
+ train_kv: bool = True,
+ train_q_out: bool = True,
+ hidden_size: Optional[int] = None,
+ cross_attention_dim: Optional[int] = None,
+ out_bias: bool = True,
+ dropout: float = 0.0,
+ ):
+ super().__init__()
+ self.train_kv = train_kv
+ self.train_q_out = train_q_out
+
+ self.hidden_size = hidden_size
+ self.cross_attention_dim = cross_attention_dim
+
+ # `_custom_diffusion` id for easy serialization and loading.
+ if self.train_kv:
+ self.to_k_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
+ self.to_v_custom_diffusion = nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False)
+ if self.train_q_out:
+ self.to_q_custom_diffusion = nn.Linear(hidden_size, hidden_size, bias=False)
+ self.to_out_custom_diffusion = nn.ModuleList([])
+ self.to_out_custom_diffusion.append(nn.Linear(hidden_size, hidden_size, bias=out_bias))
+ self.to_out_custom_diffusion.append(nn.Dropout(dropout))
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ batch_size, sequence_length, _ = hidden_states.shape
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ if self.train_q_out:
+ query = self.to_q_custom_diffusion(hidden_states)
+ else:
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ crossattn = False
+ encoder_hidden_states = hidden_states
+ else:
+ crossattn = True
+ if attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ if self.train_kv:
+ key = self.to_k_custom_diffusion(encoder_hidden_states.to(self.to_k_custom_diffusion.weight.dtype))
+ value = self.to_v_custom_diffusion(encoder_hidden_states.to(self.to_v_custom_diffusion.weight.dtype))
+ key = key.to(attn.to_q.weight.dtype)
+ value = value.to(attn.to_q.weight.dtype)
+
+ else:
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ if crossattn:
+ detach = torch.ones_like(key)
+ detach[:, :1, :] = detach[:, :1, :] * 0.0
+ key = detach * key + (1 - detach) * key.detach()
+ value = detach * value + (1 - detach) * value.detach()
+
+ inner_dim = hidden_states.shape[-1]
+
+ head_dim = inner_dim // attn.heads
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if self.train_q_out:
+ # linear proj
+ hidden_states = self.to_out_custom_diffusion[0](hidden_states)
+ # dropout
+ hidden_states = self.to_out_custom_diffusion[1](hidden_states)
+ else:
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ return hidden_states
+
+
+class SlicedAttnProcessor:
+ r"""
+ Processor for implementing sliced attention.
+
+ Args:
+ slice_size (`int`, *optional*):
+ The number of steps to compute attention. Uses as many slices as `attention_head_dim // slice_size`, and
+ `attention_head_dim` must be a multiple of the `slice_size`.
+ """
+
+ def __init__(self, slice_size: int):
+ self.slice_size = slice_size
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ residual = hidden_states
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+ dim = query.shape[-1]
+ query = attn.head_to_batch_dim(query)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+ key = attn.head_to_batch_dim(key)
+ value = attn.head_to_batch_dim(value)
+
+ batch_size_attention, query_tokens, _ = query.shape
+ hidden_states = torch.zeros(
+ (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype
+ )
+
+ for i in range((batch_size_attention - 1) // self.slice_size + 1):
+ start_idx = i * self.slice_size
+ end_idx = (i + 1) * self.slice_size
+
+ query_slice = query[start_idx:end_idx]
+ key_slice = key[start_idx:end_idx]
+ attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None
+
+ attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)
+
+ attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx])
+
+ hidden_states[start_idx:end_idx] = attn_slice
+
+ hidden_states = attn.batch_to_head_dim(hidden_states)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class SlicedAttnAddedKVProcessor:
+ r"""
+ Processor for implementing sliced attention with extra learnable key and value matrices for the text encoder.
+
+ Args:
+ slice_size (`int`, *optional*):
+ The number of steps to compute attention. Uses as many slices as `attention_head_dim // slice_size`, and
+ `attention_head_dim` must be a multiple of the `slice_size`.
+ """
+
+ def __init__(self, slice_size):
+ self.slice_size = slice_size
+
+ def __call__(
+ self,
+ attn: "Attention",
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ residual = hidden_states
+
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)
+
+ batch_size, sequence_length, _ = hidden_states.shape
+
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+ dim = query.shape[-1]
+ query = attn.head_to_batch_dim(query)
+
+ encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+ encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+ encoder_hidden_states_key_proj = attn.head_to_batch_dim(encoder_hidden_states_key_proj)
+ encoder_hidden_states_value_proj = attn.head_to_batch_dim(encoder_hidden_states_value_proj)
+
+ if not attn.only_cross_attention:
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+ key = attn.head_to_batch_dim(key)
+ value = attn.head_to_batch_dim(value)
+ key = torch.cat([encoder_hidden_states_key_proj, key], dim=1)
+ value = torch.cat([encoder_hidden_states_value_proj, value], dim=1)
+ else:
+ key = encoder_hidden_states_key_proj
+ value = encoder_hidden_states_value_proj
+
+ batch_size_attention, query_tokens, _ = query.shape
+ hidden_states = torch.zeros(
+ (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype
+ )
+
+ for i in range((batch_size_attention - 1) // self.slice_size + 1):
+ start_idx = i * self.slice_size
+ end_idx = (i + 1) * self.slice_size
+
+ query_slice = query[start_idx:end_idx]
+ key_slice = key[start_idx:end_idx]
+ attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None
+
+ attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice)
+
+ attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx])
+
+ hidden_states[start_idx:end_idx] = attn_slice
+
+ hidden_states = attn.batch_to_head_dim(hidden_states)
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape)
+ hidden_states = hidden_states + residual
+
+ return hidden_states
+
+
+class SpatialNorm(nn.Module):
+ """
+ Spatially conditioned normalization as defined in https://arxiv.org/abs/2209.09002.
+
+ Args:
+ f_channels (`int`):
+ The number of channels for input to group normalization layer, and output of the spatial norm layer.
+ zq_channels (`int`):
+ The number of channels for the quantized vector as described in the paper.
+ """
+
+ def __init__(
+ self,
+ f_channels: int,
+ zq_channels: int,
+ ):
+ super().__init__()
+ self.norm_layer = nn.GroupNorm(num_channels=f_channels, num_groups=32, eps=1e-6, affine=True)
+ self.conv_y = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0)
+ self.conv_b = nn.Conv2d(zq_channels, f_channels, kernel_size=1, stride=1, padding=0)
+
+ def forward(self, f: torch.Tensor, zq: torch.Tensor) -> torch.Tensor:
+ f_size = f.shape[-2:]
+ zq = F.interpolate(zq, size=f_size, mode="nearest")
+ norm_f = self.norm_layer(f)
+ new_f = norm_f * self.conv_y(zq) + self.conv_b(zq)
+ return new_f
+
+
+class IPAdapterAttnProcessor(nn.Module):
+ r"""
+ Attention processor for Multiple IP-Adapters.
+
+ Args:
+ hidden_size (`int`):
+ The hidden size of the attention layer.
+ cross_attention_dim (`int`):
+ The number of channels in the `encoder_hidden_states`.
+ num_tokens (`int`, `Tuple[int]` or `List[int]`, defaults to `(4,)`):
+ The context length of the image features.
+ scale (`float` or List[`float`], defaults to 1.0):
+ the weight scale of image prompt.
+ """
+
+ def __init__(self, hidden_size, cross_attention_dim=None, num_tokens=(4,), scale=1.0):
+ super().__init__()
+
+ self.hidden_size = hidden_size
+ self.cross_attention_dim = cross_attention_dim
+
+ if not isinstance(num_tokens, (tuple, list)):
+ num_tokens = [num_tokens]
+ self.num_tokens = num_tokens
+
+ if not isinstance(scale, list):
+ scale = [scale] * len(num_tokens)
+ if len(scale) != len(num_tokens):
+ raise ValueError("`scale` should be a list of integers with the same length as `num_tokens`.")
+ self.scale = scale
+
+ self.to_k_ip = nn.ModuleList(
+ [nn.Linear(cross_attention_dim, hidden_size, bias=False) for _ in range(len(num_tokens))]
+ )
+ self.to_v_ip = nn.ModuleList(
+ [nn.Linear(cross_attention_dim, hidden_size, bias=False) for _ in range(len(num_tokens))]
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ scale: float = 1.0,
+ ip_adapter_masks: Optional[torch.Tensor] = None,
+ ):
+ residual = hidden_states
+
+ # separate ip_hidden_states from encoder_hidden_states
+ if encoder_hidden_states is not None:
+ if isinstance(encoder_hidden_states, tuple):
+ encoder_hidden_states, ip_hidden_states = encoder_hidden_states
+ else:
+ deprecation_message = (
+ "You have passed a tensor as `encoder_hidden_states`. This is deprecated and will be removed in a future release."
+ " Please make sure to update your script to pass `encoder_hidden_states` as a tuple to suppress this warning."
+ )
+ deprecate("encoder_hidden_states not a tuple", "1.0.0", deprecation_message, standard_warn=False)
+ end_pos = encoder_hidden_states.shape[1] - self.num_tokens[0]
+ encoder_hidden_states, ip_hidden_states = (
+ encoder_hidden_states[:, :end_pos, :],
+ [encoder_hidden_states[:, end_pos:, :]],
+ )
+
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ query = attn.head_to_batch_dim(query)
+ key = attn.head_to_batch_dim(key)
+ value = attn.head_to_batch_dim(value)
+
+ attention_probs = attn.get_attention_scores(query, key, attention_mask)
+ hidden_states = torch.bmm(attention_probs, value)
+ hidden_states = attn.batch_to_head_dim(hidden_states)
+
+ if ip_adapter_masks is not None:
+ if not isinstance(ip_adapter_masks, List):
+ # for backward compatibility, we accept `ip_adapter_mask` as a tensor of shape [num_ip_adapter, 1, height, width]
+ ip_adapter_masks = list(ip_adapter_masks.unsqueeze(1))
+ if not (len(ip_adapter_masks) == len(self.scale) == len(ip_hidden_states)):
+ raise ValueError(
+ f"Length of ip_adapter_masks array ({len(ip_adapter_masks)}) must match "
+ f"length of self.scale array ({len(self.scale)}) and number of ip_hidden_states "
+ f"({len(ip_hidden_states)})"
+ )
+ else:
+ for index, (mask, scale, ip_state) in enumerate(zip(ip_adapter_masks, self.scale, ip_hidden_states)):
+ if mask is None:
+ continue
+ if not isinstance(mask, torch.Tensor) or mask.ndim != 4:
+ raise ValueError(
+ "Each element of the ip_adapter_masks array should be a tensor with shape "
+ "[1, num_images_for_ip_adapter, height, width]."
+ " Please use `IPAdapterMaskProcessor` to preprocess your mask"
+ )
+ if mask.shape[1] != ip_state.shape[1]:
+ raise ValueError(
+ f"Number of masks ({mask.shape[1]}) does not match "
+ f"number of ip images ({ip_state.shape[1]}) at index {index}"
+ )
+ if isinstance(scale, list) and not len(scale) == mask.shape[1]:
+ raise ValueError(
+ f"Number of masks ({mask.shape[1]}) does not match "
+ f"number of scales ({len(scale)}) at index {index}"
+ )
+ else:
+ ip_adapter_masks = [None] * len(self.scale)
+
+ # for ip-adapter
+ for current_ip_hidden_states, scale, to_k_ip, to_v_ip, mask in zip(
+ ip_hidden_states, self.scale, self.to_k_ip, self.to_v_ip, ip_adapter_masks
+ ):
+ skip = False
+ if isinstance(scale, list):
+ if all(s == 0 for s in scale):
+ skip = True
+ elif scale == 0:
+ skip = True
+ if not skip:
+ if mask is not None:
+ if not isinstance(scale, list):
+ scale = [scale] * mask.shape[1]
+
+ current_num_images = mask.shape[1]
+ for i in range(current_num_images):
+ ip_key = to_k_ip(current_ip_hidden_states[:, i, :, :])
+ ip_value = to_v_ip(current_ip_hidden_states[:, i, :, :])
+
+ ip_key = attn.head_to_batch_dim(ip_key)
+ ip_value = attn.head_to_batch_dim(ip_value)
+
+ ip_attention_probs = attn.get_attention_scores(query, ip_key, None)
+ _current_ip_hidden_states = torch.bmm(ip_attention_probs, ip_value)
+ _current_ip_hidden_states = attn.batch_to_head_dim(_current_ip_hidden_states)
+
+ mask_downsample = IPAdapterMaskProcessor.downsample(
+ mask[:, i, :, :],
+ batch_size,
+ _current_ip_hidden_states.shape[1],
+ _current_ip_hidden_states.shape[2],
+ )
+
+ mask_downsample = mask_downsample.to(dtype=query.dtype, device=query.device)
+
+ hidden_states = hidden_states + scale[i] * (_current_ip_hidden_states * mask_downsample)
+ else:
+ ip_key = to_k_ip(current_ip_hidden_states)
+ ip_value = to_v_ip(current_ip_hidden_states)
+
+ ip_key = attn.head_to_batch_dim(ip_key)
+ ip_value = attn.head_to_batch_dim(ip_value)
+
+ ip_attention_probs = attn.get_attention_scores(query, ip_key, None)
+ current_ip_hidden_states = torch.bmm(ip_attention_probs, ip_value)
+ current_ip_hidden_states = attn.batch_to_head_dim(current_ip_hidden_states)
+
+ hidden_states = hidden_states + scale * current_ip_hidden_states
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class IPAdapterAttnProcessor2_0(torch.nn.Module):
+ r"""
+ Attention processor for IP-Adapter for PyTorch 2.0.
+
+ Args:
+ hidden_size (`int`):
+ The hidden size of the attention layer.
+ cross_attention_dim (`int`):
+ The number of channels in the `encoder_hidden_states`.
+ num_tokens (`int`, `Tuple[int]` or `List[int]`, defaults to `(4,)`):
+ The context length of the image features.
+ scale (`float` or `List[float]`, defaults to 1.0):
+ the weight scale of image prompt.
+ """
+
+ def __init__(self, hidden_size, cross_attention_dim=None, num_tokens=(4,), scale=1.0):
+ super().__init__()
+
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ f"{self.__class__.__name__} requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+
+ self.hidden_size = hidden_size
+ self.cross_attention_dim = cross_attention_dim
+
+ if not isinstance(num_tokens, (tuple, list)):
+ num_tokens = [num_tokens]
+ self.num_tokens = num_tokens
+
+ if not isinstance(scale, list):
+ scale = [scale] * len(num_tokens)
+ if len(scale) != len(num_tokens):
+ raise ValueError("`scale` should be a list of integers with the same length as `num_tokens`.")
+ self.scale = scale
+
+ self.to_k_ip = nn.ModuleList(
+ [nn.Linear(cross_attention_dim, hidden_size, bias=False) for _ in range(len(num_tokens))]
+ )
+ self.to_v_ip = nn.ModuleList(
+ [nn.Linear(cross_attention_dim, hidden_size, bias=False) for _ in range(len(num_tokens))]
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ temb: Optional[torch.Tensor] = None,
+ scale: float = 1.0,
+ ip_adapter_masks: Optional[torch.Tensor] = None,
+ ):
+ residual = hidden_states
+
+ # separate ip_hidden_states from encoder_hidden_states
+ if encoder_hidden_states is not None:
+ if isinstance(encoder_hidden_states, tuple):
+ encoder_hidden_states, ip_hidden_states = encoder_hidden_states
+ else:
+ deprecation_message = (
+ "You have passed a tensor as `encoder_hidden_states`. This is deprecated and will be removed in a future release."
+ " Please make sure to update your script to pass `encoder_hidden_states` as a tuple to suppress this warning."
+ )
+ deprecate("encoder_hidden_states not a tuple", "1.0.0", deprecation_message, standard_warn=False)
+ end_pos = encoder_hidden_states.shape[1] - self.num_tokens[0]
+ encoder_hidden_states, ip_hidden_states = (
+ encoder_hidden_states[:, :end_pos, :],
+ [encoder_hidden_states[:, end_pos:, :]],
+ )
+
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if ip_adapter_masks is not None:
+ if not isinstance(ip_adapter_masks, List):
+ # for backward compatibility, we accept `ip_adapter_mask` as a tensor of shape [num_ip_adapter, 1, height, width]
+ ip_adapter_masks = list(ip_adapter_masks.unsqueeze(1))
+ if not (len(ip_adapter_masks) == len(self.scale) == len(ip_hidden_states)):
+ raise ValueError(
+ f"Length of ip_adapter_masks array ({len(ip_adapter_masks)}) must match "
+ f"length of self.scale array ({len(self.scale)}) and number of ip_hidden_states "
+ f"({len(ip_hidden_states)})"
+ )
+ else:
+ for index, (mask, scale, ip_state) in enumerate(zip(ip_adapter_masks, self.scale, ip_hidden_states)):
+ if mask is None:
+ continue
+ if not isinstance(mask, torch.Tensor) or mask.ndim != 4:
+ raise ValueError(
+ "Each element of the ip_adapter_masks array should be a tensor with shape "
+ "[1, num_images_for_ip_adapter, height, width]."
+ " Please use `IPAdapterMaskProcessor` to preprocess your mask"
+ )
+ if mask.shape[1] != ip_state.shape[1]:
+ raise ValueError(
+ f"Number of masks ({mask.shape[1]}) does not match "
+ f"number of ip images ({ip_state.shape[1]}) at index {index}"
+ )
+ if isinstance(scale, list) and not len(scale) == mask.shape[1]:
+ raise ValueError(
+ f"Number of masks ({mask.shape[1]}) does not match "
+ f"number of scales ({len(scale)}) at index {index}"
+ )
+ else:
+ ip_adapter_masks = [None] * len(self.scale)
+
+ # for ip-adapter
+ for current_ip_hidden_states, scale, to_k_ip, to_v_ip, mask in zip(
+ ip_hidden_states, self.scale, self.to_k_ip, self.to_v_ip, ip_adapter_masks
+ ):
+ skip = False
+ if isinstance(scale, list):
+ if all(s == 0 for s in scale):
+ skip = True
+ elif scale == 0:
+ skip = True
+ if not skip:
+ if mask is not None:
+ if not isinstance(scale, list):
+ scale = [scale] * mask.shape[1]
+
+ current_num_images = mask.shape[1]
+ for i in range(current_num_images):
+ ip_key = to_k_ip(current_ip_hidden_states[:, i, :, :])
+ ip_value = to_v_ip(current_ip_hidden_states[:, i, :, :])
+
+ ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ _current_ip_hidden_states = F.scaled_dot_product_attention(
+ query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
+ )
+
+ _current_ip_hidden_states = _current_ip_hidden_states.transpose(1, 2).reshape(
+ batch_size, -1, attn.heads * head_dim
+ )
+ _current_ip_hidden_states = _current_ip_hidden_states.to(query.dtype)
+
+ mask_downsample = IPAdapterMaskProcessor.downsample(
+ mask[:, i, :, :],
+ batch_size,
+ _current_ip_hidden_states.shape[1],
+ _current_ip_hidden_states.shape[2],
+ )
+
+ mask_downsample = mask_downsample.to(dtype=query.dtype, device=query.device)
+ hidden_states = hidden_states + scale[i] * (_current_ip_hidden_states * mask_downsample)
+ else:
+ ip_key = to_k_ip(current_ip_hidden_states)
+ ip_value = to_v_ip(current_ip_hidden_states)
+
+ ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ current_ip_hidden_states = F.scaled_dot_product_attention(
+ query, ip_key, ip_value, attn_mask=None, dropout_p=0.0, is_causal=False
+ )
+
+ current_ip_hidden_states = current_ip_hidden_states.transpose(1, 2).reshape(
+ batch_size, -1, attn.heads * head_dim
+ )
+ current_ip_hidden_states = current_ip_hidden_states.to(query.dtype)
+
+ hidden_states = hidden_states + scale * current_ip_hidden_states
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class IPAdapterXFormersAttnProcessor(torch.nn.Module):
+ r"""
+ Attention processor for IP-Adapter using xFormers.
+
+ Args:
+ hidden_size (`int`):
+ The hidden size of the attention layer.
+ cross_attention_dim (`int`):
+ The number of channels in the `encoder_hidden_states`.
+ num_tokens (`int`, `Tuple[int]` or `List[int]`, defaults to `(4,)`):
+ The context length of the image features.
+ scale (`float` or `List[float]`, defaults to 1.0):
+ the weight scale of image prompt.
+ attention_op (`Callable`, *optional*, defaults to `None`):
+ The base
+ [operator](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.AttentionOpBase) to
+ use as the attention operator. It is recommended to set to `None`, and allow xFormers to choose the best
+ operator.
+ """
+
+ def __init__(
+ self,
+ hidden_size,
+ cross_attention_dim=None,
+ num_tokens=(4,),
+ scale=1.0,
+ attention_op: Optional[Callable] = None,
+ ):
+ super().__init__()
+
+ self.hidden_size = hidden_size
+ self.cross_attention_dim = cross_attention_dim
+ self.attention_op = attention_op
+
+ if not isinstance(num_tokens, (tuple, list)):
+ num_tokens = [num_tokens]
+ self.num_tokens = num_tokens
+
+ if not isinstance(scale, list):
+ scale = [scale] * len(num_tokens)
+ if len(scale) != len(num_tokens):
+ raise ValueError("`scale` should be a list of integers with the same length as `num_tokens`.")
+ self.scale = scale
+
+ self.to_k_ip = nn.ModuleList(
+ [nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) for _ in range(len(num_tokens))]
+ )
+ self.to_v_ip = nn.ModuleList(
+ [nn.Linear(cross_attention_dim or hidden_size, hidden_size, bias=False) for _ in range(len(num_tokens))]
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ temb: Optional[torch.FloatTensor] = None,
+ scale: float = 1.0,
+ ip_adapter_masks: Optional[torch.FloatTensor] = None,
+ ):
+ residual = hidden_states
+
+ # separate ip_hidden_states from encoder_hidden_states
+ if encoder_hidden_states is not None:
+ if isinstance(encoder_hidden_states, tuple):
+ encoder_hidden_states, ip_hidden_states = encoder_hidden_states
+ else:
+ deprecation_message = (
+ "You have passed a tensor as `encoder_hidden_states`. This is deprecated and will be removed in a future release."
+ " Please make sure to update your script to pass `encoder_hidden_states` as a tuple to suppress this warning."
+ )
+ deprecate("encoder_hidden_states not a tuple", "1.0.0", deprecation_message, standard_warn=False)
+ end_pos = encoder_hidden_states.shape[1] - self.num_tokens[0]
+ encoder_hidden_states, ip_hidden_states = (
+ encoder_hidden_states[:, :end_pos, :],
+ [encoder_hidden_states[:, end_pos:, :]],
+ )
+
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ batch_size, sequence_length, _ = (
+ hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+ )
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # expand our mask's singleton query_tokens dimension:
+ # [batch*heads, 1, key_tokens] ->
+ # [batch*heads, query_tokens, key_tokens]
+ # so that it can be added as a bias onto the attention scores that xformers computes:
+ # [batch*heads, query_tokens, key_tokens]
+ # we do this explicitly because xformers doesn't broadcast the singleton dimension for us.
+ _, query_tokens, _ = hidden_states.shape
+ attention_mask = attention_mask.expand(-1, query_tokens, -1)
+
+ if attn.group_norm is not None:
+ hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states)
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+ elif attn.norm_cross:
+ encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)
+
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ query = attn.head_to_batch_dim(query).contiguous()
+ key = attn.head_to_batch_dim(key).contiguous()
+ value = attn.head_to_batch_dim(value).contiguous()
+
+ hidden_states = xformers.ops.memory_efficient_attention(
+ query, key, value, attn_bias=attention_mask, op=self.attention_op
+ )
+ hidden_states = hidden_states.to(query.dtype)
+ hidden_states = attn.batch_to_head_dim(hidden_states)
+
+ if ip_hidden_states:
+ if ip_adapter_masks is not None:
+ if not isinstance(ip_adapter_masks, List):
+ # for backward compatibility, we accept `ip_adapter_mask` as a tensor of shape [num_ip_adapter, 1, height, width]
+ ip_adapter_masks = list(ip_adapter_masks.unsqueeze(1))
+ if not (len(ip_adapter_masks) == len(self.scale) == len(ip_hidden_states)):
+ raise ValueError(
+ f"Length of ip_adapter_masks array ({len(ip_adapter_masks)}) must match "
+ f"length of self.scale array ({len(self.scale)}) and number of ip_hidden_states "
+ f"({len(ip_hidden_states)})"
+ )
+ else:
+ for index, (mask, scale, ip_state) in enumerate(
+ zip(ip_adapter_masks, self.scale, ip_hidden_states)
+ ):
+ if mask is None:
+ continue
+ if not isinstance(mask, torch.Tensor) or mask.ndim != 4:
+ raise ValueError(
+ "Each element of the ip_adapter_masks array should be a tensor with shape "
+ "[1, num_images_for_ip_adapter, height, width]."
+ " Please use `IPAdapterMaskProcessor` to preprocess your mask"
+ )
+ if mask.shape[1] != ip_state.shape[1]:
+ raise ValueError(
+ f"Number of masks ({mask.shape[1]}) does not match "
+ f"number of ip images ({ip_state.shape[1]}) at index {index}"
+ )
+ if isinstance(scale, list) and not len(scale) == mask.shape[1]:
+ raise ValueError(
+ f"Number of masks ({mask.shape[1]}) does not match "
+ f"number of scales ({len(scale)}) at index {index}"
+ )
+ else:
+ ip_adapter_masks = [None] * len(self.scale)
+
+ # for ip-adapter
+ for current_ip_hidden_states, scale, to_k_ip, to_v_ip, mask in zip(
+ ip_hidden_states, self.scale, self.to_k_ip, self.to_v_ip, ip_adapter_masks
+ ):
+ skip = False
+ if isinstance(scale, list):
+ if all(s == 0 for s in scale):
+ skip = True
+ elif scale == 0:
+ skip = True
+ if not skip:
+ if mask is not None:
+ mask = mask.to(torch.float16)
+ if not isinstance(scale, list):
+ scale = [scale] * mask.shape[1]
+
+ current_num_images = mask.shape[1]
+ for i in range(current_num_images):
+ ip_key = to_k_ip(current_ip_hidden_states[:, i, :, :])
+ ip_value = to_v_ip(current_ip_hidden_states[:, i, :, :])
+
+ ip_key = attn.head_to_batch_dim(ip_key).contiguous()
+ ip_value = attn.head_to_batch_dim(ip_value).contiguous()
+
+ _current_ip_hidden_states = xformers.ops.memory_efficient_attention(
+ query, ip_key, ip_value, op=self.attention_op
+ )
+ _current_ip_hidden_states = _current_ip_hidden_states.to(query.dtype)
+ _current_ip_hidden_states = attn.batch_to_head_dim(_current_ip_hidden_states)
+
+ mask_downsample = IPAdapterMaskProcessor.downsample(
+ mask[:, i, :, :],
+ batch_size,
+ _current_ip_hidden_states.shape[1],
+ _current_ip_hidden_states.shape[2],
+ )
+
+ mask_downsample = mask_downsample.to(dtype=query.dtype, device=query.device)
+ hidden_states = hidden_states + scale[i] * (_current_ip_hidden_states * mask_downsample)
+ else:
+ ip_key = to_k_ip(current_ip_hidden_states)
+ ip_value = to_v_ip(current_ip_hidden_states)
+
+ ip_key = attn.head_to_batch_dim(ip_key).contiguous()
+ ip_value = attn.head_to_batch_dim(ip_value).contiguous()
+
+ current_ip_hidden_states = xformers.ops.memory_efficient_attention(
+ query, ip_key, ip_value, op=self.attention_op
+ )
+ current_ip_hidden_states = current_ip_hidden_states.to(query.dtype)
+ current_ip_hidden_states = attn.batch_to_head_dim(current_ip_hidden_states)
+
+ hidden_states = hidden_states + scale * current_ip_hidden_states
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if input_ndim == 4:
+ hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class SD3IPAdapterJointAttnProcessor2_0(torch.nn.Module):
+ """
+ Attention processor for IP-Adapter used typically in processing the SD3-like self-attention projections, with
+ additional image-based information and timestep embeddings.
+
+ Args:
+ hidden_size (`int`):
+ The number of hidden channels.
+ ip_hidden_states_dim (`int`):
+ The image feature dimension.
+ head_dim (`int`):
+ The number of head channels.
+ timesteps_emb_dim (`int`, defaults to 1280):
+ The number of input channels for timestep embedding.
+ scale (`float`, defaults to 0.5):
+ IP-Adapter scale.
+ """
+
+ def __init__(
+ self,
+ hidden_size: int,
+ ip_hidden_states_dim: int,
+ head_dim: int,
+ timesteps_emb_dim: int = 1280,
+ scale: float = 0.5,
+ ):
+ super().__init__()
+
+ # To prevent circular import
+ from .normalization import AdaLayerNorm, RMSNorm
+
+ self.norm_ip = AdaLayerNorm(timesteps_emb_dim, output_dim=ip_hidden_states_dim * 2, norm_eps=1e-6, chunk_dim=1)
+ self.to_k_ip = nn.Linear(ip_hidden_states_dim, hidden_size, bias=False)
+ self.to_v_ip = nn.Linear(ip_hidden_states_dim, hidden_size, bias=False)
+ self.norm_q = RMSNorm(head_dim, 1e-6)
+ self.norm_k = RMSNorm(head_dim, 1e-6)
+ self.norm_ip_k = RMSNorm(head_dim, 1e-6)
+ self.scale = scale
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: torch.FloatTensor = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ ip_hidden_states: torch.FloatTensor = None,
+ temb: torch.FloatTensor = None,
+ ) -> torch.FloatTensor:
+ """
+ Perform the attention computation, integrating image features (if provided) and timestep embeddings.
+
+ If `ip_hidden_states` is `None`, this is equivalent to using JointAttnProcessor2_0.
+
+ Args:
+ attn (`Attention`):
+ Attention instance.
+ hidden_states (`torch.FloatTensor`):
+ Input `hidden_states`.
+ encoder_hidden_states (`torch.FloatTensor`, *optional*):
+ The encoder hidden states.
+ attention_mask (`torch.FloatTensor`, *optional*):
+ Attention mask.
+ ip_hidden_states (`torch.FloatTensor`, *optional*):
+ Image embeddings.
+ temb (`torch.FloatTensor`, *optional*):
+ Timestep embeddings.
+
+ Returns:
+ `torch.FloatTensor`: Output hidden states.
+ """
+ residual = hidden_states
+
+ batch_size = hidden_states.shape[0]
+
+ # `sample` projections.
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ img_query = query
+ img_key = key
+ img_value = value
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ # `context` projections.
+ if encoder_hidden_states is not None:
+ encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+ encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+ encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+ encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+ encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+ batch_size, -1, attn.heads, head_dim
+ ).transpose(1, 2)
+
+ if attn.norm_added_q is not None:
+ encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+ if attn.norm_added_k is not None:
+ encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+ query = torch.cat([query, encoder_hidden_states_query_proj], dim=2)
+ key = torch.cat([key, encoder_hidden_states_key_proj], dim=2)
+ value = torch.cat([value, encoder_hidden_states_value_proj], dim=2)
+
+ hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+ hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states = hidden_states.to(query.dtype)
+
+ if encoder_hidden_states is not None:
+ # Split the attention outputs.
+ hidden_states, encoder_hidden_states = (
+ hidden_states[:, : residual.shape[1]],
+ hidden_states[:, residual.shape[1] :],
+ )
+ if not attn.context_pre_only:
+ encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+ # IP Adapter
+ if self.scale != 0 and ip_hidden_states is not None:
+ # Norm image features
+ norm_ip_hidden_states = self.norm_ip(ip_hidden_states, temb=temb)
+
+ # To k and v
+ ip_key = self.to_k_ip(norm_ip_hidden_states)
+ ip_value = self.to_v_ip(norm_ip_hidden_states)
+
+ # Reshape
+ ip_key = ip_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ ip_value = ip_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ # Norm
+ query = self.norm_q(img_query)
+ img_key = self.norm_k(img_key)
+ ip_key = self.norm_ip_k(ip_key)
+
+ # cat img
+ key = torch.cat([img_key, ip_key], dim=2)
+ value = torch.cat([img_value, ip_value], dim=2)
+
+ ip_hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
+ ip_hidden_states = ip_hidden_states.transpose(1, 2).view(batch_size, -1, attn.heads * head_dim)
+ ip_hidden_states = ip_hidden_states.to(query.dtype)
+
+ hidden_states = hidden_states + ip_hidden_states * self.scale
+
+ # linear proj
+ hidden_states = attn.to_out[0](hidden_states)
+ # dropout
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if encoder_hidden_states is not None:
+ return hidden_states, encoder_hidden_states
+ else:
+ return hidden_states
+
+
+class PAGIdentitySelfAttnProcessor2_0:
+ r"""
+ Processor for implementing PAG using scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+ PAG reference: https://arxiv.org/abs/2403.17377
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "PAGIdentitySelfAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ temb: Optional[torch.FloatTensor] = None,
+ ) -> torch.Tensor:
+ residual = hidden_states
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ # chunk
+ hidden_states_org, hidden_states_ptb = hidden_states.chunk(2)
+
+ # original path
+ batch_size, sequence_length, _ = hidden_states_org.shape
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if attn.group_norm is not None:
+ hidden_states_org = attn.group_norm(hidden_states_org.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states_org)
+ key = attn.to_k(hidden_states_org)
+ value = attn.to_v(hidden_states_org)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states_org = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+ hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states_org = hidden_states_org.to(query.dtype)
+
+ # linear proj
+ hidden_states_org = attn.to_out[0](hidden_states_org)
+ # dropout
+ hidden_states_org = attn.to_out[1](hidden_states_org)
+
+ if input_ndim == 4:
+ hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ # perturbed path (identity attention)
+ batch_size, sequence_length, _ = hidden_states_ptb.shape
+
+ if attn.group_norm is not None:
+ hidden_states_ptb = attn.group_norm(hidden_states_ptb.transpose(1, 2)).transpose(1, 2)
+
+ hidden_states_ptb = attn.to_v(hidden_states_ptb)
+ hidden_states_ptb = hidden_states_ptb.to(query.dtype)
+
+ # linear proj
+ hidden_states_ptb = attn.to_out[0](hidden_states_ptb)
+ # dropout
+ hidden_states_ptb = attn.to_out[1](hidden_states_ptb)
+
+ if input_ndim == 4:
+ hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ # cat
+ hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class PAGCFGIdentitySelfAttnProcessor2_0:
+ r"""
+ Processor for implementing PAG using scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+ PAG reference: https://arxiv.org/abs/2403.17377
+ """
+
+ def __init__(self):
+ if not hasattr(F, "scaled_dot_product_attention"):
+ raise ImportError(
+ "PAGCFGIdentitySelfAttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0."
+ )
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.FloatTensor,
+ encoder_hidden_states: Optional[torch.FloatTensor] = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ temb: Optional[torch.FloatTensor] = None,
+ ) -> torch.Tensor:
+ residual = hidden_states
+ if attn.spatial_norm is not None:
+ hidden_states = attn.spatial_norm(hidden_states, temb)
+
+ input_ndim = hidden_states.ndim
+ if input_ndim == 4:
+ batch_size, channel, height, width = hidden_states.shape
+ hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)
+
+ # chunk
+ hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3)
+ hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org])
+
+ # original path
+ batch_size, sequence_length, _ = hidden_states_org.shape
+
+ if attention_mask is not None:
+ attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
+ # scaled_dot_product_attention expects attention_mask shape to be
+ # (batch, heads, source_length, target_length)
+ attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
+
+ if attn.group_norm is not None:
+ hidden_states_org = attn.group_norm(hidden_states_org.transpose(1, 2)).transpose(1, 2)
+
+ query = attn.to_q(hidden_states_org)
+ key = attn.to_k(hidden_states_org)
+ value = attn.to_v(hidden_states_org)
+
+ inner_dim = key.shape[-1]
+ head_dim = inner_dim // attn.heads
+
+ query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+ value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+ # the output of sdp = (batch, num_heads, seq_len, head_dim)
+ # TODO: add support for attn.scale when we move to Torch 2.1
+ hidden_states_org = F.scaled_dot_product_attention(
+ query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+ )
+
+ hidden_states_org = hidden_states_org.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+ hidden_states_org = hidden_states_org.to(query.dtype)
+
+ # linear proj
+ hidden_states_org = attn.to_out[0](hidden_states_org)
+ # dropout
+ hidden_states_org = attn.to_out[1](hidden_states_org)
+
+ if input_ndim == 4:
+ hidden_states_org = hidden_states_org.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ # perturbed path (identity attention)
+ batch_size, sequence_length, _ = hidden_states_ptb.shape
+
+ if attn.group_norm is not None:
+ hidden_states_ptb = attn.group_norm(hidden_states_ptb.transpose(1, 2)).transpose(1, 2)
+
+ value = attn.to_v(hidden_states_ptb)
+ hidden_states_ptb = value
+ hidden_states_ptb = hidden_states_ptb.to(query.dtype)
+
+ # linear proj
+ hidden_states_ptb = attn.to_out[0](hidden_states_ptb)
+ # dropout
+ hidden_states_ptb = attn.to_out[1](hidden_states_ptb)
+
+ if input_ndim == 4:
+ hidden_states_ptb = hidden_states_ptb.transpose(-1, -2).reshape(batch_size, channel, height, width)
+
+ # cat
+ hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ hidden_states = hidden_states / attn.rescale_output_factor
+
+ return hidden_states
+
+
+class SanaMultiscaleAttnProcessor2_0:
+ r"""
+ Processor for implementing multiscale quadratic attention.
+ """
+
+ def __call__(self, attn: SanaMultiscaleLinearAttention, hidden_states: torch.Tensor) -> torch.Tensor:
+ height, width = hidden_states.shape[-2:]
+ if height * width > attn.attention_head_dim:
+ use_linear_attention = True
+ else:
+ use_linear_attention = False
+
+ residual = hidden_states
+
+ batch_size, _, height, width = list(hidden_states.size())
+ original_dtype = hidden_states.dtype
+
+ hidden_states = hidden_states.movedim(1, -1)
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(hidden_states)
+ value = attn.to_v(hidden_states)
+ hidden_states = torch.cat([query, key, value], dim=3)
+ hidden_states = hidden_states.movedim(-1, 1)
+
+ multi_scale_qkv = [hidden_states]
+ for block in attn.to_qkv_multiscale:
+ multi_scale_qkv.append(block(hidden_states))
+
+ hidden_states = torch.cat(multi_scale_qkv, dim=1)
+
+ if use_linear_attention:
+ # for linear attention upcast hidden_states to float32
+ hidden_states = hidden_states.to(dtype=torch.float32)
+
+ hidden_states = hidden_states.reshape(batch_size, -1, 3 * attn.attention_head_dim, height * width)
+
+ query, key, value = hidden_states.chunk(3, dim=2)
+ query = attn.nonlinearity(query)
+ key = attn.nonlinearity(key)
+
+ if use_linear_attention:
+ hidden_states = attn.apply_linear_attention(query, key, value)
+ hidden_states = hidden_states.to(dtype=original_dtype)
+ else:
+ hidden_states = attn.apply_quadratic_attention(query, key, value)
+
+ hidden_states = torch.reshape(hidden_states, (batch_size, -1, height, width))
+ hidden_states = attn.to_out(hidden_states.movedim(1, -1)).movedim(-1, 1)
+
+ if attn.norm_type == "rms_norm":
+ hidden_states = attn.norm_out(hidden_states.movedim(1, -1)).movedim(-1, 1)
+ else:
+ hidden_states = attn.norm_out(hidden_states)
+
+ if attn.residual_connection:
+ hidden_states = hidden_states + residual
+
+ return hidden_states
+
+
+class LoRAAttnProcessor:
+ r"""
+ Processor for implementing attention with LoRA.
+ """
+
+ def __init__(self):
+ pass
+
+
+class LoRAAttnProcessor2_0:
+ r"""
+ Processor for implementing attention with LoRA (enabled by default if you're using PyTorch 2.0).
+ """
+
+ def __init__(self):
+ pass
+
+
+class LoRAXFormersAttnProcessor:
+ r"""
+ Processor for implementing attention with LoRA using xFormers.
+ """
+
+ def __init__(self):
+ pass
+
+
+class LoRAAttnAddedKVProcessor:
+ r"""
+ Processor for implementing attention with LoRA with extra learnable key and value matrices for the text encoder.
+ """
+
+ def __init__(self):
+ pass
+
+
+class FluxSingleAttnProcessor2_0(FluxAttnProcessor2_0):
+ r"""
+ Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
+ """
+
+ def __init__(self):
+ deprecation_message = "`FluxSingleAttnProcessor2_0` is deprecated and will be removed in a future version. Please use `FluxAttnProcessor2_0` instead."
+ deprecate("FluxSingleAttnProcessor2_0", "0.32.0", deprecation_message)
+ super().__init__()
+
+
+class SanaLinearAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product linear attention.
+ """
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ original_dtype = hidden_states.dtype
+
+ if encoder_hidden_states is None:
+ encoder_hidden_states = hidden_states
+
+ query = attn.to_q(hidden_states)
+ key = attn.to_k(encoder_hidden_states)
+ value = attn.to_v(encoder_hidden_states)
+
+ if attn.norm_q is not None:
+ query = attn.norm_q(query)
+ if attn.norm_k is not None:
+ key = attn.norm_k(key)
+
+ query = query.transpose(1, 2).unflatten(1, (attn.heads, -1))
+ key = key.transpose(1, 2).unflatten(1, (attn.heads, -1)).transpose(2, 3)
+ value = value.transpose(1, 2).unflatten(1, (attn.heads, -1))
+
+ query = F.relu(query)
+ key = F.relu(key)
+
+ query, key, value = query.float(), key.float(), value.float()
+
+ value = F.pad(value, (0, 0, 0, 1), mode="constant", value=1.0)
+ scores = torch.matmul(value, key)
+ hidden_states = torch.matmul(scores, query)
+
+ hidden_states = hidden_states[:, :, :-1] / (hidden_states[:, :, -1:] + 1e-15)
+ hidden_states = hidden_states.flatten(1, 2).transpose(1, 2)
+ hidden_states = hidden_states.to(original_dtype)
+
+ hidden_states = attn.to_out[0](hidden_states)
+ hidden_states = attn.to_out[1](hidden_states)
+
+ if original_dtype == torch.float16:
+ hidden_states = hidden_states.clip(-65504, 65504)
+
+ return hidden_states
+
+
+class PAGCFGSanaLinearAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product linear attention.
+ """
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ original_dtype = hidden_states.dtype
+
+ hidden_states_uncond, hidden_states_org, hidden_states_ptb = hidden_states.chunk(3)
+ hidden_states_org = torch.cat([hidden_states_uncond, hidden_states_org])
+
+ query = attn.to_q(hidden_states_org)
+ key = attn.to_k(hidden_states_org)
+ value = attn.to_v(hidden_states_org)
+
+ query = query.transpose(1, 2).unflatten(1, (attn.heads, -1))
+ key = key.transpose(1, 2).unflatten(1, (attn.heads, -1)).transpose(2, 3)
+ value = value.transpose(1, 2).unflatten(1, (attn.heads, -1))
+
+ query = F.relu(query)
+ key = F.relu(key)
+
+ query, key, value = query.float(), key.float(), value.float()
+
+ value = F.pad(value, (0, 0, 0, 1), mode="constant", value=1.0)
+ scores = torch.matmul(value, key)
+ hidden_states_org = torch.matmul(scores, query)
+
+ hidden_states_org = hidden_states_org[:, :, :-1] / (hidden_states_org[:, :, -1:] + 1e-15)
+ hidden_states_org = hidden_states_org.flatten(1, 2).transpose(1, 2)
+ hidden_states_org = hidden_states_org.to(original_dtype)
+
+ hidden_states_org = attn.to_out[0](hidden_states_org)
+ hidden_states_org = attn.to_out[1](hidden_states_org)
+
+ # perturbed path (identity attention)
+ hidden_states_ptb = attn.to_v(hidden_states_ptb).to(original_dtype)
+
+ hidden_states_ptb = attn.to_out[0](hidden_states_ptb)
+ hidden_states_ptb = attn.to_out[1](hidden_states_ptb)
+
+ hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])
+
+ if original_dtype == torch.float16:
+ hidden_states = hidden_states.clip(-65504, 65504)
+
+ return hidden_states
+
+
+class PAGIdentitySanaLinearAttnProcessor2_0:
+ r"""
+ Processor for implementing scaled dot-product linear attention.
+ """
+
+ def __call__(
+ self,
+ attn: Attention,
+ hidden_states: torch.Tensor,
+ encoder_hidden_states: Optional[torch.Tensor] = None,
+ attention_mask: Optional[torch.Tensor] = None,
+ ) -> torch.Tensor:
+ original_dtype = hidden_states.dtype
+
+ hidden_states_org, hidden_states_ptb = hidden_states.chunk(2)
+
+ query = attn.to_q(hidden_states_org)
+ key = attn.to_k(hidden_states_org)
+ value = attn.to_v(hidden_states_org)
+
+ query = query.transpose(1, 2).unflatten(1, (attn.heads, -1))
+ key = key.transpose(1, 2).unflatten(1, (attn.heads, -1)).transpose(2, 3)
+ value = value.transpose(1, 2).unflatten(1, (attn.heads, -1))
+
+ query = F.relu(query)
+ key = F.relu(key)
+
+ query, key, value = query.float(), key.float(), value.float()
+
+ value = F.pad(value, (0, 0, 0, 1), mode="constant", value=1.0)
+ scores = torch.matmul(value, key)
+ hidden_states_org = torch.matmul(scores, query)
+
+ if hidden_states_org.dtype in [torch.float16, torch.bfloat16]:
+ hidden_states_org = hidden_states_org.float()
+
+ hidden_states_org = hidden_states_org[:, :, :-1] / (hidden_states_org[:, :, -1:] + 1e-15)
+ hidden_states_org = hidden_states_org.flatten(1, 2).transpose(1, 2)
+ hidden_states_org = hidden_states_org.to(original_dtype)
+
+ hidden_states_org = attn.to_out[0](hidden_states_org)
+ hidden_states_org = attn.to_out[1](hidden_states_org)
+
+ # perturbed path (identity attention)
+ hidden_states_ptb = attn.to_v(hidden_states_ptb).to(original_dtype)
+
+ hidden_states_ptb = attn.to_out[0](hidden_states_ptb)
+ hidden_states_ptb = attn.to_out[1](hidden_states_ptb)
+
+ hidden_states = torch.cat([hidden_states_org, hidden_states_ptb])
+
+ if original_dtype == torch.float16:
+ hidden_states = hidden_states.clip(-65504, 65504)
+
+ return hidden_states
+
+
+ADDED_KV_ATTENTION_PROCESSORS = (
+ AttnAddedKVProcessor,
+ SlicedAttnAddedKVProcessor,
+ AttnAddedKVProcessor2_0,
+ XFormersAttnAddedKVProcessor,
+)
+
+CROSS_ATTENTION_PROCESSORS = (
+ AttnProcessor,
+ AttnProcessor2_0,
+ XFormersAttnProcessor,
+ SlicedAttnProcessor,
+ IPAdapterAttnProcessor,
+ IPAdapterAttnProcessor2_0,
+ FluxIPAdapterJointAttnProcessor2_0,
+)
+
+AttentionProcessor = Union[
+ AttnProcessor,
+ CustomDiffusionAttnProcessor,
+ AttnAddedKVProcessor,
+ AttnAddedKVProcessor2_0,
+ JointAttnProcessor2_0,
+ PAGJointAttnProcessor2_0,
+ PAGCFGJointAttnProcessor2_0,
+ FusedJointAttnProcessor2_0,
+ AllegroAttnProcessor2_0,
+ AuraFlowAttnProcessor2_0,
+ FusedAuraFlowAttnProcessor2_0,
+ FluxAttnProcessor2_0,
+ FluxAttnProcessor2_0_NPU,
+ FusedFluxAttnProcessor2_0,
+ FusedFluxAttnProcessor2_0_NPU,
+ CogVideoXAttnProcessor2_0,
+ FusedCogVideoXAttnProcessor2_0,
+ XFormersAttnAddedKVProcessor,
+ XFormersAttnProcessor,
+ XLAFlashAttnProcessor2_0,
+ AttnProcessorNPU,
+ AttnProcessor2_0,
+ MochiVaeAttnProcessor2_0,
+ MochiAttnProcessor2_0,
+ StableAudioAttnProcessor2_0,
+ HunyuanAttnProcessor2_0,
+ FusedHunyuanAttnProcessor2_0,
+ PAGHunyuanAttnProcessor2_0,
+ PAGCFGHunyuanAttnProcessor2_0,
+ LuminaAttnProcessor2_0,
+ FusedAttnProcessor2_0,
+ CustomDiffusionXFormersAttnProcessor,
+ CustomDiffusionAttnProcessor2_0,
+ SlicedAttnProcessor,
+ SlicedAttnAddedKVProcessor,
+ SanaLinearAttnProcessor2_0,
+ PAGCFGSanaLinearAttnProcessor2_0,
+ PAGIdentitySanaLinearAttnProcessor2_0,
+ SanaMultiscaleLinearAttention,
+ SanaMultiscaleAttnProcessor2_0,
+ SanaMultiscaleAttentionProjection,
+ IPAdapterAttnProcessor,
+ IPAdapterAttnProcessor2_0,
+ IPAdapterXFormersAttnProcessor,
+ SD3IPAdapterJointAttnProcessor2_0,
+ PAGIdentitySelfAttnProcessor2_0,
+ PAGCFGIdentitySelfAttnProcessor2_0,
+ LoRAAttnProcessor,
+ LoRAAttnProcessor2_0,
+ LoRAXFormersAttnProcessor,
+ LoRAAttnAddedKVProcessor,
+]
\ No newline at end of file
diff --git a/wan/models/cache_utils.py b/wan/models/cache_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..d55d87f61fd8521fbafcd79b2474ddd49ee7c85d
--- /dev/null
+++ b/wan/models/cache_utils.py
@@ -0,0 +1,74 @@
+import numpy as np
+import torch
+
+
+def get_teacache_coefficients(model_name):
+ if "wan2.1-t2v-1.3b" or "wan2.1-fun-1.3b" or "Wan2.1-Fun-V1.1-1.3B" in model_name.lower():
+ return [-5.21862437e+04, 9.23041404e+03, -5.28275948e+02, 1.36987616e+01, -4.99875664e-02]
+ elif "wan2.1-t2v-14b" in model_name.lower():
+ return [-3.03318725e+05, 4.90537029e+04, -2.65530556e+03, 5.87365115e+01, -3.15583525e-01]
+ elif "wan2.1-i2v-14b-480p" in model_name.lower():
+ return [2.57151496e+05, -3.54229917e+04, 1.40286849e+03, -1.35890334e+01, 1.32517977e-01]
+ elif "wan2.1-i2v-14b-720p" or "wan2.1-fun-14b" in model_name.lower():
+ return [8.10705460e+03, 2.13393892e+03, -3.72934672e+02, 1.66203073e+01, -4.17769401e-02]
+ else:
+ print(f"The model {model_name} is not supported by TeaCache.")
+ return None
+
+
+class TeaCache():
+ """
+ Timestep Embedding Aware Cache, a training-free caching approach that estimates and leverages
+ the fluctuating differences among model outputs across timesteps, thereby accelerating the inference.
+ Please refer to:
+ 1. https://github.com/ali-vilab/TeaCache.
+ 2. Liu, Feng, et al. "Timestep Embedding Tells: It's Time to Cache for Video Diffusion Model." arXiv preprint arXiv:2411.19108 (2024).
+ """
+ def __init__(
+ self,
+ coefficients: list[float],
+ num_steps: int,
+ rel_l1_thresh: float = 0.0,
+ num_skip_start_steps: int = 0,
+ offload: bool = True,
+ ):
+ if num_steps < 1:
+ raise ValueError(f"`num_steps` must be greater than 0 but is {num_steps}.")
+ if rel_l1_thresh < 0:
+ raise ValueError(f"`rel_l1_thresh` must be greater than or equal to 0 but is {rel_l1_thresh}.")
+ if num_skip_start_steps < 0 or num_skip_start_steps > num_steps:
+ raise ValueError(
+ "`num_skip_start_steps` must be great than or equal to 0 and "
+ f"less than or equal to `num_steps={num_steps}` but is {num_skip_start_steps}."
+ )
+ self.coefficients = coefficients
+ self.num_steps = num_steps
+ self.rel_l1_thresh = rel_l1_thresh
+ self.num_skip_start_steps = num_skip_start_steps
+ self.offload = offload
+ self.rescale_func = np.poly1d(self.coefficients)
+
+ self.cnt = 0
+ self.should_calc = True
+ self.accumulated_rel_l1_distance = 0
+ self.previous_modulated_input = None
+ # Some pipelines concatenate the unconditional and text guide in forward.
+ self.previous_residual = None
+ # Some pipelines perform forward propagation separately on the unconditional and text guide.
+ self.previous_residual_cond = None
+ self.previous_residual_uncond = None
+
+ @staticmethod
+ def compute_rel_l1_distance(prev: torch.Tensor, cur: torch.Tensor) -> torch.Tensor:
+ rel_l1_distance = (torch.abs(cur - prev).mean()) / torch.abs(prev).mean()
+
+ return rel_l1_distance.cpu().item()
+
+ def reset(self):
+ self.cnt = 0
+ self.should_calc = True
+ self.accumulated_rel_l1_distance = 0
+ self.previous_modulated_input = None
+ self.previous_residual = None
+ self.previous_residual_cond = None
+ self.previous_residual_uncond = None
\ No newline at end of file
diff --git a/wan/models/motion_controller.py b/wan/models/motion_controller.py
new file mode 100644
index 0000000000000000000000000000000000000000..9529bef13e71f5e4cf8cf62bd0e60356a75e7c6d
--- /dev/null
+++ b/wan/models/motion_controller.py
@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+
+
+def sinusoidal_embedding_1d(dim, position):
+ sinusoid = torch.outer(position.type(torch.float64), torch.pow(10000, -torch.arange(dim//2, dtype=torch.float64, device=position.device).div(dim//2)))
+ x = torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1)
+ return x.to(position.dtype)
+
+class WanMotionControllerModel(torch.nn.Module):
+ def __init__(self, freq_dim=256, dim=1536):
+ super().__init__()
+ self.freq_dim = freq_dim
+ self.linear = nn.Sequential(
+ nn.Linear(freq_dim, dim),
+ nn.SiLU(),
+ nn.Linear(dim, dim),
+ nn.SiLU(),
+ nn.Linear(dim, dim * 6),
+ )
+ self.init_weight()
+
+ def forward(self, motion_bucket_id):
+ emb = sinusoidal_embedding_1d(self.freq_dim, motion_bucket_id * 10)
+ emb = self.linear(emb)
+ return emb
+
+ def init_weight(self):
+ state_dict = self.linear[-1].state_dict()
+ state_dict = {i: state_dict[i] * 0 for i in state_dict}
+ self.linear[-1].load_state_dict(state_dict)
+
+
+if __name__ == "__main__":
+ dim = 1536
+ motion_controller = WanMotionControllerModel()
+ motion_bucket_id = 100.0
+ motion_bucket_id = torch.Tensor((motion_bucket_id,)).to(dtype=torch.float32, device='cpu')
+ out = motion_controller(motion_bucket_id).unflatten(1, (6, dim))
+ print(out.size())
\ No newline at end of file
diff --git a/wan/models/motion_to_bucket.py b/wan/models/motion_to_bucket.py
new file mode 100644
index 0000000000000000000000000000000000000000..8425b45710c5a7c6cd965d8c17059e8f98113c06
--- /dev/null
+++ b/wan/models/motion_to_bucket.py
@@ -0,0 +1,81 @@
+import torch
+from diffusers import ModelMixin
+from einops import rearrange
+from torch import nn
+
+
+class Motion2bucketModel(ModelMixin):
+ def __init__(self, window_size=5, blocks=12, channels=1024, clip_channels=1280, intermediate_dim=512, output_dim=768, context_tokens=32, clip_token_num=1, final_output_dim=5120):
+ super().__init__()
+ self.window_size = window_size
+ self.clip_token_num = clip_token_num
+ self.blocks = blocks
+ self.channels = channels
+ # self.input_dim = (window_size * blocks * channels + clip_channels*clip_token_num)
+ self.input_dim = (window_size * channels + clip_channels * clip_token_num)
+ self.intermediate_dim = intermediate_dim
+ self.context_tokens = context_tokens
+ self.output_dim = output_dim
+
+ # define multiple linear layers
+ self.proj1 = nn.Linear(self.input_dim, intermediate_dim)
+ self.proj2 = nn.Linear(intermediate_dim, intermediate_dim)
+ self.proj3 = nn.Linear(intermediate_dim, context_tokens * output_dim)
+ self.act = nn.SiLU()
+
+
+ self.final_proj = torch.nn.Linear(output_dim, final_output_dim)
+ self.final_norm = torch.nn.LayerNorm(final_output_dim)
+
+ nn.init.constant_(self.final_proj.weight, 0)
+ if self.final_proj.bias is not None:
+ nn.init.constant_(self.final_proj.bias, 0)
+
+ def forward(self, audio_embeds, clip_embeds):
+ """
+ Defines the forward pass for the AudioProjModel.
+
+ Parameters:
+ audio_embeds (torch.Tensor): The input audio embeddings with shape (batch_size, video_length, blocks, channels).
+
+ Returns:
+ context_tokens (torch.Tensor): The output context tokens with shape (batch_size, video_length, context_tokens, output_dim).
+ """
+ # merge
+ video_length = audio_embeds.shape[1]
+ # audio_embeds = rearrange(audio_embeds, "bz f w b c -> (bz f) w b c")
+ audio_embeds = rearrange(audio_embeds, "bz f w c -> (bz f) w c")
+ clip_embeds = clip_embeds.repeat(audio_embeds.size()[0]//clip_embeds.size()[0], 1, 1)
+ clip_embeds = rearrange(clip_embeds, "b n d -> b (n d)")
+ # batch_size, window_size, blocks, channels = audio_embeds.shape
+ # audio_embeds = audio_embeds.view(batch_size, window_size * blocks * channels)
+ batch_size, window_size, channels = audio_embeds.shape
+ audio_embeds = audio_embeds.view(batch_size, window_size * channels)
+ audio_embeds = torch.cat([audio_embeds, clip_embeds], dim=-1)
+
+ audio_embeds = self.act(self.proj1(audio_embeds))
+ audio_embeds = self.act(self.proj2(audio_embeds))
+
+ context_tokens = self.proj3(audio_embeds).reshape(
+ batch_size, self.context_tokens, self.output_dim
+ )
+
+ # context_tokens = self.norm(context_tokens)
+ context_tokens = rearrange(
+ context_tokens, "(bz f) m c -> bz f m c", f=video_length
+ )
+
+ context_tokens = self.act(context_tokens)
+ context_tokens = self.final_norm(self.final_proj(context_tokens))
+
+ return context_tokens
+
+
+if __name__ == '__main__':
+ model = Motion2bucketModel(window_size=5)
+ # audio_features = torch.randn(1, 81, 5, 12, 768)
+ audio_features = torch.randn(1, 81, 5, 1024)
+ clip_image_features = torch.randn(1, 1, 1280)
+
+ out = model(audio_features, clip_image_features).mean(dim=2).mean(dim=1)
+ print(out.size())
diff --git a/wan/models/vocal_projector_fantasy.py b/wan/models/vocal_projector_fantasy.py
new file mode 100644
index 0000000000000000000000000000000000000000..921a0f78eb111c1556c42a44d5004276eb87a0f0
--- /dev/null
+++ b/wan/models/vocal_projector_fantasy.py
@@ -0,0 +1,131 @@
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class VocalProjModel(nn.Module):
+ def __init__(self, audio_in_dim=1024, cross_attention_dim=1024):
+ super().__init__()
+ self.cross_attention_dim = cross_attention_dim
+ self.proj = torch.nn.Linear(audio_in_dim, cross_attention_dim, bias=False)
+ self.norm = torch.nn.LayerNorm(cross_attention_dim)
+
+ def forward(self, audio_embeds):
+ context_tokens = self.proj(audio_embeds)
+ context_tokens = self.norm(context_tokens)
+ return context_tokens # [B,L,C]
+
+
+class FantasyTalkingVocalConditionModel(nn.Module):
+ def __init__(self, audio_in_dim: int, audio_proj_dim: int):
+ super().__init__()
+
+ self.audio_in_dim = audio_in_dim
+ self.audio_proj_dim = audio_proj_dim
+ # audio proj model
+ self.proj_model = self.init_proj(self.audio_proj_dim)
+
+ def init_proj(self, cross_attention_dim=5120):
+ proj_model = VocalProjModel(
+ audio_in_dim=self.audio_in_dim, cross_attention_dim=cross_attention_dim
+ )
+ return proj_model
+
+ def forward(self, audio_fea=None):
+ return self.proj_model(audio_fea) if audio_fea is not None else None
+
+
+def split_audio_sequence(audio_proj_length, num_frames=81):
+ """
+ Map the audio feature sequence to corresponding latent frame slices.
+
+ Args:
+ audio_proj_length (int): The total length of the audio feature sequence
+ (e.g., 173 in audio_proj[1, 173, 768]).
+ num_frames (int): The number of video frames in the training data (default: 81).
+
+ Returns:
+ list: A list of [start_idx, end_idx] pairs. Each pair represents the index range
+ (within the audio feature sequence) corresponding to a latent frame.
+ """
+ # Average number of tokens per original video frame
+ tokens_per_frame = audio_proj_length / num_frames
+
+ # Each latent frame covers 4 video frames, and we want the center
+ tokens_per_latent_frame = tokens_per_frame * 4
+ half_tokens = int(tokens_per_latent_frame / 2)
+
+ pos_indices = []
+ for i in range(int((num_frames - 1) / 4) + 1):
+ if i == 0:
+ pos_indices.append(0)
+ else:
+ start_token = tokens_per_frame * ((i - 1) * 4 + 1)
+ end_token = tokens_per_frame * (i * 4 + 1)
+ center_token = int((start_token + end_token) / 2) - 1
+ pos_indices.append(center_token)
+
+ # Build index ranges centered around each position
+ pos_idx_ranges = [[idx - half_tokens, idx + half_tokens] for idx in pos_indices]
+
+ # Adjust the first range to avoid negative start index
+ pos_idx_ranges[0] = [
+ -(half_tokens * 2 - pos_idx_ranges[1][0]),
+ pos_idx_ranges[1][0],
+ ]
+
+ return pos_idx_ranges
+
+
+def split_tensor_with_padding(input_tensor, pos_idx_ranges, expand_length=0):
+ """
+ Split the input tensor into subsequences based on index ranges, and apply right-side zero-padding
+ if the range exceeds the input boundaries.
+
+ Args:
+ input_tensor (Tensor): Input audio tensor of shape [1, L, 768].
+ pos_idx_ranges (list): A list of index ranges, e.g. [[-7, 1], [1, 9], ..., [165, 173]].
+ expand_length (int): Number of tokens to expand on both sides of each subsequence.
+
+ Returns:
+ sub_sequences (Tensor): A tensor of shape [1, F, L, 768], where L is the length after padding.
+ Each element is a padded subsequence.
+ k_lens (Tensor): A tensor of shape [F], representing the actual (unpadded) length of each subsequence.
+ Useful for ignoring padding tokens in attention masks.
+ """
+ pos_idx_ranges = [
+ [idx[0] - expand_length, idx[1] + expand_length] for idx in pos_idx_ranges
+ ]
+ sub_sequences = []
+ seq_len = input_tensor.size(1) # 173
+ max_valid_idx = seq_len - 1 # 172
+ k_lens_list = []
+ for start, end in pos_idx_ranges:
+ # Calculate the fill amount
+ pad_front = max(-start, 0)
+ pad_back = max(end - max_valid_idx, 0)
+
+ # Calculate the start and end indices of the valid part
+ valid_start = max(start, 0)
+ valid_end = min(end, max_valid_idx)
+
+ # Extract the valid part
+ if valid_start <= valid_end:
+ valid_part = input_tensor[:, valid_start: valid_end + 1, :]
+ else:
+ valid_part = input_tensor.new_zeros((1, 0, input_tensor.size(2)))
+
+ # In the sequence dimension (the 1st dimension) perform padding
+ padded_subseq = F.pad(
+ valid_part,
+ (0, 0, 0, pad_back + pad_front, 0, 0),
+ mode="constant",
+ value=0,
+ )
+ k_lens_list.append(padded_subseq.size(-2) - pad_back - pad_front)
+
+ sub_sequences.append(padded_subseq)
+ return torch.stack(sub_sequences, dim=1), torch.tensor(
+ k_lens_list, dtype=torch.long
+ )
diff --git a/wan/models/vocal_projector_fantasy_14B.py b/wan/models/vocal_projector_fantasy_14B.py
new file mode 100644
index 0000000000000000000000000000000000000000..52eebf4738f17bbf30d932e5a7c26e663446b8b3
--- /dev/null
+++ b/wan/models/vocal_projector_fantasy_14B.py
@@ -0,0 +1,458 @@
+import os
+import warnings
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from wan.models.vocal_projector_fantasy import split_audio_sequence, split_tensor_with_padding
+
+try:
+ import flash_attn_interface
+ FLASH_ATTN_3_AVAILABLE = True
+except ModuleNotFoundError:
+ FLASH_ATTN_3_AVAILABLE = False
+
+try:
+ import flash_attn
+ FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+ FLASH_ATTN_2_AVAILABLE = False
+
+
+def flash_attention(
+ q,
+ k,
+ v,
+ q_lens=None,
+ k_lens=None,
+ dropout_p=0.,
+ softmax_scale=None,
+ q_scale=None,
+ causal=False,
+ window_size=(-1, -1),
+ deterministic=False,
+ dtype=torch.bfloat16,
+ version=None,
+):
+ """
+ q: [B, Lq, Nq, C1].
+ k: [B, Lk, Nk, C1].
+ v: [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+ q_lens: [B].
+ k_lens: [B].
+ dropout_p: float. Dropout probability.
+ softmax_scale: float. The scaling of QK^T before applying softmax.
+ causal: bool. Whether to apply causal attention mask.
+ window_size: (left right). If not (-1, -1), apply sliding window local attention.
+ deterministic: bool. If True, slightly slower and uses more memory.
+ dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+ """
+ half_dtypes = (torch.float16, torch.bfloat16)
+ assert dtype in half_dtypes
+ assert q.device.type == 'cuda' and q.size(-1) <= 256
+
+ # params
+ b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+ def half(x):
+ return x if x.dtype in half_dtypes else x.to(dtype)
+
+ # preprocess query
+ if q_lens is None:
+ q = half(q.flatten(0, 1))
+ q_lens = torch.tensor(
+ [lq] * b, dtype=torch.int32).to(
+ device=q.device, non_blocking=True)
+ else:
+ q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
+
+ # preprocess key, value
+ if k_lens is None:
+ k = half(k.flatten(0, 1))
+ v = half(v.flatten(0, 1))
+ k_lens = torch.tensor(
+ [lk] * b, dtype=torch.int32).to(
+ device=k.device, non_blocking=True)
+ else:
+ k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
+ v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
+
+ q = q.to(v.dtype)
+ k = k.to(v.dtype)
+
+ if q_scale is not None:
+ q = q * q_scale
+
+ if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+ warnings.warn(
+ 'Flash attention 3 is not available, use flash attention 2 instead.'
+ )
+
+ # apply attention
+ if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+ # Note: dropout_p, window_size are not supported in FA3 now.
+ x = flash_attn_interface.flash_attn_varlen_func(
+ q=q,
+ k=k,
+ v=v,
+ cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ seqused_q=None,
+ seqused_k=None,
+ max_seqlen_q=lq,
+ max_seqlen_k=lk,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ deterministic=deterministic)[0].unflatten(0, (b, lq))
+ else:
+ assert FLASH_ATTN_2_AVAILABLE
+ x = flash_attn.flash_attn_varlen_func(
+ q=q,
+ k=k,
+ v=v,
+ cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ max_seqlen_q=lq,
+ max_seqlen_k=lk,
+ dropout_p=dropout_p,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ window_size=window_size,
+ deterministic=deterministic).unflatten(0, (b, lq))
+
+ # output
+ return x.type(out_dtype)
+
+
+def attention(
+ q,
+ k,
+ v,
+ q_lens=None,
+ k_lens=None,
+ dropout_p=0.,
+ softmax_scale=None,
+ q_scale=None,
+ causal=False,
+ window_size=(-1, -1),
+ deterministic=False,
+ dtype=torch.bfloat16,
+ fa_version=None,
+):
+ if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+ return flash_attention(
+ q=q,
+ k=k,
+ v=v,
+ q_lens=q_lens,
+ k_lens=k_lens,
+ dropout_p=dropout_p,
+ softmax_scale=softmax_scale,
+ q_scale=q_scale,
+ causal=causal,
+ window_size=window_size,
+ deterministic=deterministic,
+ dtype=dtype,
+ version=fa_version,
+ )
+ else:
+ if q_lens is not None or k_lens is not None:
+ warnings.warn(
+ 'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.'
+ )
+ attn_mask = None
+
+ q = q.transpose(1, 2)
+ k = k.transpose(1, 2)
+ v = v.transpose(1, 2)
+
+ if torch.backends.cuda.flash_sdp_enabled() is False or torch.backends.cuda.enable_flash_sdp is False:
+ print(1/0)
+
+ out = torch.nn.functional.scaled_dot_product_attention(
+ q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p)
+
+ out = out.transpose(1, 2).contiguous()
+ return out
+
+
+class WanRMSNorm(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):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ """
+ 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 WanLayerNorm(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):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ """
+ return super().forward(x.float()).type_as(x)
+
+
+class VocalCrossAttention(nn.Module):
+ def __init__(self,
+ vocal_dim,
+ dit_dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ eps=1e-6,
+
+ ):
+ assert vocal_dim % num_heads == 0
+ super().__init__()
+ self.vocal_dim = vocal_dim
+ self.dit_dim = dit_dim
+ self.num_heads = num_heads
+ self.head_dim = vocal_dim // num_heads
+ self.window_size = window_size
+ self.qk_norm = qk_norm
+ self.eps = eps
+
+ # layers
+ self.q = nn.Linear(vocal_dim, vocal_dim)
+ self.k = nn.Linear(dit_dim, vocal_dim)
+ self.v = nn.Linear(dit_dim, vocal_dim)
+ self.o = nn.Linear(vocal_dim, vocal_dim)
+ self.norm_q = WanRMSNorm(vocal_dim, eps=eps) if qk_norm else nn.Identity()
+ self.norm_k = WanRMSNorm(vocal_dim, eps=eps) if qk_norm else nn.Identity()
+
+ def forward(self, x, context, q_lens, dtype):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ context(Tensor): Shape [B, L2, C]
+ context_lens(Tensor): Shape [B]
+ """
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ latents_num_frames = 21
+ q = self.norm_q(self.q(x.to(dtype))).view(b * latents_num_frames, -1, n, d)
+ k = self.norm_k(self.k(context.to(dtype))).view(b * latents_num_frames, -1, n, d)
+ v = self.v(context.to(dtype)).view(b * latents_num_frames, -1, n, d)
+ # compute attention
+
+ x = attention(
+ q.to(dtype),
+ k.to(dtype),
+ v.to(dtype),
+ q_lens=None,
+ k_lens=None,
+ )
+ x = x.to(dtype)
+ x = x.view(b, -1, n, d)
+
+ # output
+ x = x.flatten(2)
+ x = self.o(x)
+ return x
+
+
+class VocalAttentionBlock(nn.Module):
+
+ def __init__(self,
+ vocal_dim,
+ dit_dim,
+ ffn_dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ cross_attn_norm=True,
+ eps=1e-6):
+ super().__init__()
+ self.vocal_dim = vocal_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 = WanLayerNorm(vocal_dim, eps)
+ self.norm3 = WanLayerNorm(
+ vocal_dim, eps,
+ elementwise_affine=True) if cross_attn_norm else nn.Identity()
+
+ self.cross_attn = VocalCrossAttention(vocal_dim=vocal_dim,
+ dit_dim=dit_dim,
+ num_heads=num_heads,
+ window_size=(-1, -1),
+ qk_norm=qk_norm,
+ eps=eps)
+ self.norm2 = WanLayerNorm(vocal_dim, eps)
+ self.ffn = nn.Sequential(
+ nn.Linear(vocal_dim, ffn_dim),
+ nn.GELU(approximate='tanh'),
+ nn.Linear(ffn_dim, vocal_dim))
+
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 6, vocal_dim) / vocal_dim ** 0.5)
+
+ def forward(
+ self,
+ x,
+ e,
+ context,
+ q_lens,
+ dtype=torch.float32,
+ ):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ e(Tensor): Shape [B, 6, C]
+ seq_lens(Tensor): Shape [B], length of each sequence in batch
+ grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+ freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+ """
+ e = (self.modulation + e).chunk(6, dim=1)
+
+ # self-attention
+ if len(x.shape) == 4:
+ b, t, n, d = x.size()
+ x = rearrange(x, "b t n d -> b (t n) d", t=t)
+
+ temp_x = self.norm1(x) * (1 + e[1]) + e[0]
+ temp_x = temp_x.to(dtype)
+ x = x + temp_x * e[2]
+
+ # cross-attention & ffn function
+ def cross_attn_ffn(x, context, q_lens, e):
+ # cross-attention
+ x = x + self.cross_attn(self.norm3(x), context, q_lens, dtype)
+ # ffn function
+ temp_x = self.norm2(x) * (1 + e[4]) + e[3]
+ temp_x = temp_x.to(dtype)
+
+ y = self.ffn(temp_x)
+ x = x + y * e[5]
+ return x
+
+ x = cross_attn_ffn(x, context, q_lens, e)
+ return x
+
+
+class Final_Head(nn.Module):
+
+ def __init__(self, dim, out_dim, eps=1e-6):
+ super().__init__()
+ self.dim = dim
+ self.eps = eps
+
+ # layers
+ self.norm = WanLayerNorm(dim, eps)
+ self.final_proj = nn.Linear(dim, out_dim)
+
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+ def forward(self, x, e):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ e(Tensor): Shape [B, C]
+ """
+ e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
+ x = (self.final_proj(self.norm(x) * (1 + e[1]) + e[0]))
+ return x
+
+class VocalProjModel(nn.Module):
+ def __init__(self, audio_in_dim=1024, cross_attention_dim=1024):
+ super().__init__()
+ self.cross_attention_dim = cross_attention_dim
+ self.proj_1 = torch.nn.Linear(audio_in_dim, 2048, bias=False)
+ self.norm_1 = torch.nn.LayerNorm(2048)
+ self.proj_2 = torch.nn.Linear(2048, cross_attention_dim, bias=False)
+ self.norm_2 = torch.nn.LayerNorm(cross_attention_dim)
+
+ def forward(self, audio_embeds):
+ context_tokens = self.proj_1(audio_embeds)
+ context_tokens = self.norm_1(context_tokens)
+ context_tokens = self.proj_2(context_tokens)
+ context_tokens = self.norm_2(context_tokens)
+ return context_tokens # [B,L,C]
+
+
+class FantasyTalkingVocalCondition14BModel(nn.Module):
+ def __init__(self, audio_in_dim: int, audio_proj_dim: int, dit_dim: int):
+ super().__init__()
+
+ self.audio_in_dim = audio_in_dim
+ self.audio_proj_dim = audio_proj_dim
+ # audio proj model
+ self.proj_model = self.init_proj(self.audio_proj_dim)
+
+ num_layers = 2
+ self.blocks = nn.ModuleList([
+ VocalAttentionBlock(
+ vocal_dim=audio_proj_dim,
+ dit_dim=dit_dim,
+ ffn_dim=audio_proj_dim * 2,
+ num_heads=8,
+ window_size=(-1, -1),
+ qk_norm=True,
+ cross_attn_norm=True,
+ )
+ for _ in range(num_layers)
+ ])
+
+ self.final_head = Final_Head(dim=audio_proj_dim, out_dim=audio_proj_dim)
+
+ def init_proj(self, cross_attention_dim=5120):
+ proj_model = VocalProjModel(
+ audio_in_dim=self.audio_in_dim, cross_attention_dim=cross_attention_dim
+ )
+ return proj_model
+
+ def forward(self, vocal_embeddings=None, video_sample_n_frames=81, latents=None, e0=None, e=None):
+ vocal_proj_feature = self.proj_model(vocal_embeddings)
+ pos_idx_ranges = split_audio_sequence(vocal_proj_feature.size(1), num_frames=video_sample_n_frames)
+ vocal_proj_split, vocal_context_lens = split_tensor_with_padding(vocal_proj_feature, pos_idx_ranges, expand_length=4)
+ latents_num_frames = vocal_proj_split.size()[1]
+ for block in self.blocks:
+ vocal_proj_split = block(
+ x=vocal_proj_split,
+ e=e0,
+ context=latents,
+ q_lens=vocal_context_lens,
+ )
+ context_tokens = self.final_head(vocal_proj_split, e)
+ context_tokens = rearrange(context_tokens, "b (f n) c -> b f n c", f=latents_num_frames)
+ if vocal_embeddings.size()[0] > 1:
+ vocal_context_lens = torch.cat([vocal_context_lens] * 3)
+ return context_tokens, vocal_context_lens
+
+
+if __name__ == "__main__":
+ model = FantasyTalkingVocalCondition14BModel(audio_in_dim=768, audio_proj_dim=5120, dit_dim=5120)
+ vocal_embeddings = torch.randn(1, 134, 768)
+ latents = torch.randn(1, 21504, 5120)
+ e0 = torch.randn(1, 6, 5120)
+ e = torch.randn(1, 5120)
+ out, _ = model(vocal_embeddings=vocal_embeddings, video_sample_n_frames=81, latents=latents, e0=e0, e=e)
+ print(out.size())
\ No newline at end of file
diff --git a/wan/models/vocal_projector_fantasy_1B.py b/wan/models/vocal_projector_fantasy_1B.py
new file mode 100644
index 0000000000000000000000000000000000000000..5690841eed27246a2bc04ae061a0caf235c6495e
--- /dev/null
+++ b/wan/models/vocal_projector_fantasy_1B.py
@@ -0,0 +1,455 @@
+import os
+import warnings
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+
+from wan.models.vocal_projector_fantasy import split_audio_sequence, split_tensor_with_padding
+
+try:
+ import flash_attn_interface
+ FLASH_ATTN_3_AVAILABLE = True
+except ModuleNotFoundError:
+ FLASH_ATTN_3_AVAILABLE = False
+
+try:
+ import flash_attn
+ FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+ FLASH_ATTN_2_AVAILABLE = False
+
+
+def flash_attention(
+ q,
+ k,
+ v,
+ q_lens=None,
+ k_lens=None,
+ dropout_p=0.,
+ softmax_scale=None,
+ q_scale=None,
+ causal=False,
+ window_size=(-1, -1),
+ deterministic=False,
+ dtype=torch.bfloat16,
+ version=None,
+):
+ """
+ q: [B, Lq, Nq, C1].
+ k: [B, Lk, Nk, C1].
+ v: [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+ q_lens: [B].
+ k_lens: [B].
+ dropout_p: float. Dropout probability.
+ softmax_scale: float. The scaling of QK^T before applying softmax.
+ causal: bool. Whether to apply causal attention mask.
+ window_size: (left right). If not (-1, -1), apply sliding window local attention.
+ deterministic: bool. If True, slightly slower and uses more memory.
+ dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+ """
+ half_dtypes = (torch.float16, torch.bfloat16)
+ assert dtype in half_dtypes
+ assert q.device.type == 'cuda' and q.size(-1) <= 256
+
+ # params
+ b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+ def half(x):
+ return x if x.dtype in half_dtypes else x.to(dtype)
+
+ # preprocess query
+ if q_lens is None:
+ q = half(q.flatten(0, 1))
+ q_lens = torch.tensor(
+ [lq] * b, dtype=torch.int32).to(
+ device=q.device, non_blocking=True)
+ else:
+ q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
+
+ # preprocess key, value
+ if k_lens is None:
+ k = half(k.flatten(0, 1))
+ v = half(v.flatten(0, 1))
+ k_lens = torch.tensor(
+ [lk] * b, dtype=torch.int32).to(
+ device=k.device, non_blocking=True)
+ else:
+ k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
+ v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
+
+ q = q.to(v.dtype)
+ k = k.to(v.dtype)
+
+ if q_scale is not None:
+ q = q * q_scale
+
+ if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+ warnings.warn(
+ 'Flash attention 3 is not available, use flash attention 2 instead.'
+ )
+
+ # apply attention
+ if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+ # Note: dropout_p, window_size are not supported in FA3 now.
+ x = flash_attn_interface.flash_attn_varlen_func(
+ q=q,
+ k=k,
+ v=v,
+ cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ seqused_q=None,
+ seqused_k=None,
+ max_seqlen_q=lq,
+ max_seqlen_k=lk,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ deterministic=deterministic)[0].unflatten(0, (b, lq))
+ else:
+ assert FLASH_ATTN_2_AVAILABLE
+ x = flash_attn.flash_attn_varlen_func(
+ q=q,
+ k=k,
+ v=v,
+ cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ max_seqlen_q=lq,
+ max_seqlen_k=lk,
+ dropout_p=dropout_p,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ window_size=window_size,
+ deterministic=deterministic).unflatten(0, (b, lq))
+
+ # output
+ return x.type(out_dtype)
+
+
+def attention(
+ q,
+ k,
+ v,
+ q_lens=None,
+ k_lens=None,
+ dropout_p=0.,
+ softmax_scale=None,
+ q_scale=None,
+ causal=False,
+ window_size=(-1, -1),
+ deterministic=False,
+ dtype=torch.bfloat16,
+ fa_version=None,
+):
+ if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+ return flash_attention(
+ q=q,
+ k=k,
+ v=v,
+ q_lens=q_lens,
+ k_lens=k_lens,
+ dropout_p=dropout_p,
+ softmax_scale=softmax_scale,
+ q_scale=q_scale,
+ causal=causal,
+ window_size=window_size,
+ deterministic=deterministic,
+ dtype=dtype,
+ version=fa_version,
+ )
+ else:
+ if q_lens is not None or k_lens is not None:
+ warnings.warn(
+ 'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.'
+ )
+ attn_mask = None
+
+ q = q.transpose(1, 2)
+ k = k.transpose(1, 2)
+ v = v.transpose(1, 2)
+
+ if torch.backends.cuda.flash_sdp_enabled() is False or torch.backends.cuda.enable_flash_sdp is False:
+ print(1/0)
+
+ out = torch.nn.functional.scaled_dot_product_attention(
+ q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p)
+
+ out = out.transpose(1, 2).contiguous()
+ return out
+
+
+class WanRMSNorm(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):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ """
+ 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 WanLayerNorm(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):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ """
+ return super().forward(x.float()).type_as(x)
+
+
+class VocalCrossAttention(nn.Module):
+ def __init__(self,
+ vocal_dim,
+ dit_dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ eps=1e-6,
+
+ ):
+ assert vocal_dim % num_heads == 0
+ super().__init__()
+ self.vocal_dim = vocal_dim
+ self.dit_dim = dit_dim
+ self.num_heads = num_heads
+ self.head_dim = vocal_dim // num_heads
+ self.window_size = window_size
+ self.qk_norm = qk_norm
+ self.eps = eps
+
+ # layers
+ self.q = nn.Linear(vocal_dim, vocal_dim)
+ self.k = nn.Linear(dit_dim, vocal_dim)
+ self.v = nn.Linear(dit_dim, vocal_dim)
+ self.o = nn.Linear(vocal_dim, vocal_dim)
+ self.norm_q = WanRMSNorm(vocal_dim, eps=eps) if qk_norm else nn.Identity()
+ self.norm_k = WanRMSNorm(vocal_dim, eps=eps) if qk_norm else nn.Identity()
+
+ def forward(self, x, context, q_lens, dtype):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ context(Tensor): Shape [B, L2, C]
+ context_lens(Tensor): Shape [B]
+ """
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ latents_num_frames = 21
+ q = self.norm_q(self.q(x.to(dtype))).view(b * latents_num_frames, -1, n, d)
+ k = self.norm_k(self.k(context.to(dtype))).view(b * latents_num_frames, -1, n, d)
+ v = self.v(context.to(dtype)).view(b * latents_num_frames, -1, n, d)
+ # compute attention
+
+ x = attention(
+ q.to(dtype),
+ k.to(dtype),
+ v.to(dtype),
+ q_lens=None,
+ k_lens=None,
+ )
+ x = x.to(dtype)
+ x = x.view(b, -1, n, d)
+
+ # output
+ x = x.flatten(2)
+ x = self.o(x)
+ return x
+
+
+class VocalAttentionBlock(nn.Module):
+
+ def __init__(self,
+ vocal_dim,
+ dit_dim,
+ ffn_dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ cross_attn_norm=True,
+ eps=1e-6):
+ super().__init__()
+ self.vocal_dim = vocal_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 = WanLayerNorm(vocal_dim, eps)
+ self.norm3 = WanLayerNorm(
+ vocal_dim, eps,
+ elementwise_affine=True) if cross_attn_norm else nn.Identity()
+
+ self.cross_attn = VocalCrossAttention(vocal_dim=vocal_dim,
+ dit_dim=dit_dim,
+ num_heads=num_heads,
+ window_size=(-1, -1),
+ qk_norm=qk_norm,
+ eps=eps)
+ self.norm2 = WanLayerNorm(vocal_dim, eps)
+ self.ffn = nn.Sequential(
+ nn.Linear(vocal_dim, ffn_dim),
+ nn.GELU(approximate='tanh'),
+ nn.Linear(ffn_dim, vocal_dim))
+
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 6, vocal_dim) / vocal_dim ** 0.5)
+
+ def forward(
+ self,
+ x,
+ e,
+ context,
+ q_lens,
+ dtype=torch.float32,
+ ):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ e(Tensor): Shape [B, 6, C]
+ seq_lens(Tensor): Shape [B], length of each sequence in batch
+ grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+ freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+ """
+ e = (self.modulation + e).chunk(6, dim=1)
+
+ # self-attention
+ if len(x.shape) == 4:
+ b, t, n, d = x.size()
+ x = rearrange(x, "b t n d -> b (t n) d", t=t)
+
+ temp_x = self.norm1(x) * (1 + e[1]) + e[0]
+ temp_x = temp_x.to(dtype)
+ x = x + temp_x * e[2]
+
+ # cross-attention & ffn function
+ def cross_attn_ffn(x, context, q_lens, e):
+ # cross-attention
+ x = x + self.cross_attn(self.norm3(x), context, q_lens, dtype)
+ # ffn function
+ temp_x = self.norm2(x) * (1 + e[4]) + e[3]
+ temp_x = temp_x.to(dtype)
+
+ y = self.ffn(temp_x)
+ x = x + y * e[5]
+ return x
+
+ x = cross_attn_ffn(x, context, q_lens, e)
+ return x
+
+
+class Final_Head(nn.Module):
+
+ def __init__(self, dim, out_dim, eps=1e-6):
+ super().__init__()
+ self.dim = dim
+ self.eps = eps
+
+ # layers
+ self.norm = WanLayerNorm(dim, eps)
+ self.final_proj = nn.Linear(dim, out_dim)
+
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+ def forward(self, x, e):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ e(Tensor): Shape [B, C]
+ """
+ e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
+ x = (self.final_proj(self.norm(x) * (1 + e[1]) + e[0]))
+ return x
+
+class VocalProjModel(nn.Module):
+ def __init__(self, audio_in_dim=1024, cross_attention_dim=1024):
+ super().__init__()
+ self.cross_attention_dim = cross_attention_dim
+ self.proj = torch.nn.Linear(audio_in_dim, cross_attention_dim, bias=False)
+ self.norm = torch.nn.LayerNorm(cross_attention_dim)
+
+ def forward(self, audio_embeds):
+ context_tokens = self.proj(audio_embeds)
+ context_tokens = self.norm(context_tokens)
+ return context_tokens # [B,L,C]
+
+
+class FantasyTalkingVocalCondition1BModel(nn.Module):
+ def __init__(self, audio_in_dim: int, audio_proj_dim: int, dit_dim: int):
+ super().__init__()
+
+ self.audio_in_dim = audio_in_dim
+ self.audio_proj_dim = audio_proj_dim
+ # audio proj model
+ self.proj_model = self.init_proj(self.audio_proj_dim)
+
+ num_layers = 2
+ self.blocks = nn.ModuleList([
+ VocalAttentionBlock(
+ vocal_dim=audio_proj_dim,
+ dit_dim=dit_dim,
+ ffn_dim=audio_proj_dim * 2,
+ num_heads=8,
+ window_size=(-1, -1),
+ qk_norm=True,
+ cross_attn_norm=True,
+ )
+ for _ in range(num_layers)
+ ])
+
+ self.final_head = Final_Head(dim=audio_proj_dim, out_dim=audio_proj_dim)
+
+ def init_proj(self, cross_attention_dim=5120):
+ proj_model = VocalProjModel(
+ audio_in_dim=self.audio_in_dim, cross_attention_dim=cross_attention_dim
+ )
+ return proj_model
+
+ def forward(self, vocal_embeddings=None, video_sample_n_frames=81, latents=None, e0=None, e=None):
+ vocal_proj_feature = self.proj_model(vocal_embeddings)
+ pos_idx_ranges = split_audio_sequence(vocal_proj_feature.size(1), num_frames=video_sample_n_frames)
+ vocal_proj_split, vocal_context_lens = split_tensor_with_padding(vocal_proj_feature, pos_idx_ranges, expand_length=4)
+ latents_num_frames = vocal_proj_split.size()[1]
+ for block in self.blocks:
+ vocal_proj_split = block(
+ x=vocal_proj_split,
+ e=e0,
+ context=latents,
+ q_lens=vocal_context_lens,
+ )
+ context_tokens = self.final_head(vocal_proj_split, e)
+ context_tokens = rearrange(context_tokens, "b (f n) c -> b f n c", f=latents_num_frames)
+ if vocal_embeddings.size()[0] > 1:
+ vocal_context_lens = torch.cat([vocal_context_lens] * 3)
+ return context_tokens, vocal_context_lens
+
+
+if __name__ == "__main__":
+ model = FantasyTalkingVocalCondition1BModel(audio_in_dim=768, audio_proj_dim=1536, dit_dim=1536)
+ vocal_embeddings = torch.randn(3, 134, 768)
+ latents = torch.randn(3, 21504, 1536)
+ e0 = torch.randn(3, 6, 1536)
+ e = torch.randn(3, 1536)
+ out, seq_len = model(vocal_embeddings=vocal_embeddings, video_sample_n_frames=81, latents=latents, e0=e0, e=e)
+ print(out.size())
+ print(seq_len.size())
\ No newline at end of file
diff --git a/wan/models/wan_fantasy_transformer3d_14B.py b/wan/models/wan_fantasy_transformer3d_14B.py
new file mode 100644
index 0000000000000000000000000000000000000000..0eb651196bf4da64eec7128e1fbb14667bf7d88e
--- /dev/null
+++ b/wan/models/wan_fantasy_transformer3d_14B.py
@@ -0,0 +1,1324 @@
+# Modified from https://github.com/Wan-Video/Wan2.1/blob/main/wan/modules/model.py
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+
+import glob
+import json
+import math
+import os
+import types
+import warnings
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+import torch
+import torch.cuda.amp as amp
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import is_torch_version, logging
+from torch import nn
+from einops import rearrange
+
+from .vocal_projector_fantasy_14B import FantasyTalkingVocalCondition14BModel
+from ..dist import (get_sequence_parallel_rank,
+ get_sequence_parallel_world_size, get_sp_group,
+ xFuserLongContextAttention)
+
+from .cache_utils import TeaCache
+from ..dist.wan_xfuser import usp_attn_forward
+
+try:
+ import flash_attn_interface
+
+ FLASH_ATTN_3_AVAILABLE = True
+except ModuleNotFoundError:
+ FLASH_ATTN_3_AVAILABLE = False
+
+try:
+ import flash_attn
+
+ FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+ FLASH_ATTN_2_AVAILABLE = False
+
+
+def flash_attention(
+ q,
+ k,
+ v,
+ q_lens=None,
+ k_lens=None,
+ dropout_p=0.,
+ softmax_scale=None,
+ q_scale=None,
+ causal=False,
+ window_size=(-1, -1),
+ deterministic=False,
+ dtype=torch.bfloat16,
+ version=None,
+):
+ """
+ q: [B, Lq, Nq, C1].
+ k: [B, Lk, Nk, C1].
+ v: [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+ q_lens: [B].
+ k_lens: [B].
+ dropout_p: float. Dropout probability.
+ softmax_scale: float. The scaling of QK^T before applying softmax.
+ causal: bool. Whether to apply causal attention mask.
+ window_size: (left right). If not (-1, -1), apply sliding window local attention.
+ deterministic: bool. If True, slightly slower and uses more memory.
+ dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+ """
+ half_dtypes = (torch.float16, torch.bfloat16)
+ assert dtype in half_dtypes
+ assert q.device.type == 'cuda' and q.size(-1) <= 256
+
+ # params
+ b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+ def half(x):
+ return x if x.dtype in half_dtypes else x.to(dtype)
+
+ # preprocess query
+ if q_lens is None:
+ q = half(q.flatten(0, 1))
+ q_lens = torch.tensor(
+ [lq] * b, dtype=torch.int32).to(
+ device=q.device, non_blocking=True)
+ else:
+ q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
+
+ # preprocess key, value
+ if k_lens is None:
+ k = half(k.flatten(0, 1))
+ v = half(v.flatten(0, 1))
+ k_lens = torch.tensor(
+ [lk] * b, dtype=torch.int32).to(
+ device=k.device, non_blocking=True)
+ else:
+ k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
+ v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
+
+ q = q.to(v.dtype)
+ k = k.to(v.dtype)
+
+ if q_scale is not None:
+ q = q * q_scale
+
+ if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+ warnings.warn(
+ 'Flash attention 3 is not available, use flash attention 2 instead.'
+ )
+
+ # apply attention
+ if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+ # Note: dropout_p, window_size are not supported in FA3 now.
+ x = flash_attn_interface.flash_attn_varlen_func(
+ q=q,
+ k=k,
+ v=v,
+ cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ seqused_q=None,
+ seqused_k=None,
+ max_seqlen_q=lq,
+ max_seqlen_k=lk,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ deterministic=deterministic)[0].unflatten(0, (b, lq))
+ else:
+ assert FLASH_ATTN_2_AVAILABLE
+ x = flash_attn.flash_attn_varlen_func(
+ q=q,
+ k=k,
+ v=v,
+ cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ max_seqlen_q=lq,
+ max_seqlen_k=lk,
+ dropout_p=dropout_p,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ window_size=window_size,
+ deterministic=deterministic).unflatten(0, (b, lq))
+
+ # output
+ return x.type(out_dtype)
+
+
+def attention(
+ q,
+ k,
+ v,
+ q_lens=None,
+ k_lens=None,
+ dropout_p=0.,
+ softmax_scale=None,
+ q_scale=None,
+ causal=False,
+ window_size=(-1, -1),
+ deterministic=False,
+ dtype=torch.bfloat16,
+ fa_version=None,
+):
+ if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+ return flash_attention(
+ q=q,
+ k=k,
+ v=v,
+ q_lens=q_lens,
+ k_lens=k_lens,
+ dropout_p=dropout_p,
+ softmax_scale=softmax_scale,
+ q_scale=q_scale,
+ causal=causal,
+ window_size=window_size,
+ deterministic=deterministic,
+ dtype=dtype,
+ version=fa_version,
+ )
+ else:
+ if q_lens is not None or k_lens is not None:
+ warnings.warn(
+ 'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.'
+ )
+ attn_mask = None
+
+ q = q.transpose(1, 2)
+ k = k.transpose(1, 2)
+ v = v.transpose(1, 2)
+
+ if torch.backends.cuda.flash_sdp_enabled() is False or torch.backends.cuda.enable_flash_sdp is False:
+ print(1/0)
+
+ out = torch.nn.functional.scaled_dot_product_attention(
+ q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p)
+
+ out = out.transpose(1, 2).contiguous()
+ return out
+
+
+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
+
+
+# modified from https://github.com/thu-ml/RIFLEx/blob/main/riflex_utils.py
+@amp.autocast(enabled=False)
+def get_1d_rotary_pos_embed_riflex(
+ pos: Union[np.ndarray, int],
+ dim: int,
+ theta: float = 10000.0,
+ use_real=False,
+ k: Optional[int] = None,
+ L_test: Optional[int] = None,
+ L_test_scale: Optional[int] = None,
+):
+ """
+ RIFLEx: Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
+
+ This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' and the end
+ index 'end'. The 'theta' parameter scales the frequencies. The returned tensor contains complex values in complex64
+ data type.
+
+ Args:
+ dim (`int`): Dimension of the frequency tensor.
+ pos (`np.ndarray` or `int`): Position indices for the frequency tensor. [S] or scalar
+ theta (`float`, *optional*, defaults to 10000.0):
+ Scaling factor for frequency computation. Defaults to 10000.0.
+ use_real (`bool`, *optional*):
+ If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+ k (`int`, *optional*, defaults to None): the index for the intrinsic frequency in RoPE
+ L_test (`int`, *optional*, defaults to None): the number of frames for inference
+ Returns:
+ `torch.Tensor`: Precomputed frequency tensor with complex exponentials. [S, D/2]
+ """
+ assert dim % 2 == 0
+
+ if isinstance(pos, int):
+ pos = torch.arange(pos)
+ if isinstance(pos, np.ndarray):
+ pos = torch.from_numpy(pos) # type: ignore # [S]
+
+ freqs = 1.0 / torch.pow(theta,
+ torch.arange(0, dim, 2).to(torch.float64).div(dim))
+
+ # === Riflex modification start ===
+ # Reduce the intrinsic frequency to stay within a single period after extrapolation (see Eq. (8)).
+ # Empirical observations show that a few videos may exhibit repetition in the tail frames.
+ # To be conservative, we multiply by 0.9 to keep the extrapolated length below 90% of a single period.
+ if k is not None:
+ freqs[k - 1] = 0.9 * 2 * torch.pi / L_test
+ # === Riflex modification end ===
+ if L_test_scale is not None:
+ freqs[k - 1] = freqs[k - 1] / L_test_scale
+
+ freqs = torch.outer(pos, freqs) # type: ignore # [S, D/2]
+ if use_real:
+ freqs_cos = freqs.cos().repeat_interleave(2, dim=1).float() # [S, D]
+ freqs_sin = freqs.sin().repeat_interleave(2, dim=1).float() # [S, D]
+ return freqs_cos, freqs_sin
+ else:
+ # lumina
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 # [S, D/2]
+ return freqs_cis
+
+
+@amp.autocast(enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+ n, c = x.size(2), x.size(3) // 2
+
+ # split freqs
+ freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+ # loop over samples
+ output = []
+ for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+ seq_len = f * h * w
+
+ # precompute multipliers
+ x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float32).reshape(
+ seq_len, n, -1, 2))
+ freqs_i = torch.cat([
+ freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+ freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+ freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+ ],
+ dim=-1).reshape(seq_len, 1, -1)
+
+ # apply rotary embedding
+ x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+ x_i = torch.cat([x_i, x[i, seq_len:]])
+
+ # append to collection
+ output.append(x_i)
+ return torch.stack(output).float()
+
+
+class WanRMSNorm(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):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ """
+ 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 WanLayerNorm(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):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ """
+ return super().forward(x.float()).type_as(x)
+
+
+class WanSelfAttention(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 = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+ self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+ def forward(self, x, seq_lens, grid_sizes, freqs, dtype):
+ r"""
+ 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.to(dtype))).view(b, s, n, d)
+ k = self.norm_k(self.k(x.to(dtype))).view(b, s, n, d)
+ v = self.v(x.to(dtype)).view(b, s, n, d)
+ return q, k, v
+
+ q, k, v = qkv_fn(x)
+
+ x = attention(
+ q=rope_apply(q, grid_sizes, freqs).to(dtype),
+ k=rope_apply(k, grid_sizes, freqs).to(dtype),
+ v=v.to(dtype),
+ k_lens=seq_lens,
+ window_size=self.window_size)
+ x = x.to(dtype)
+
+ # output
+ x = x.flatten(2)
+ x = self.o(x)
+ return x
+
+
+class WanT2VCrossAttention(WanSelfAttention):
+
+ def forward(self, x, context, context_lens, dtype):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ context(Tensor): Shape [B, L2, C]
+ context_lens(Tensor): Shape [B]
+ """
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.norm_q(self.q(x.to(dtype))).view(b, -1, n, d)
+ k = self.norm_k(self.k(context.to(dtype))).view(b, -1, n, d)
+ v = self.v(context.to(dtype)).view(b, -1, n, d)
+
+ # compute attention
+ x = attention(
+ q.to(dtype),
+ k.to(dtype),
+ v.to(dtype),
+ k_lens=context_lens
+ )
+ x = x.to(dtype)
+
+ # output
+ x = x.flatten(2)
+ x = self.o(x)
+ return x
+
+
+class WanI2VCrossAttention(WanSelfAttention):
+
+ def __init__(self,
+ dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ eps=1e-6):
+ super().__init__(dim, num_heads, window_size, qk_norm, eps)
+
+ # self.alpha = nn.Parameter(torch.zeros((1, )))
+ self.k_img = nn.Linear(dim, dim)
+ self.v_img = nn.Linear(dim, dim)
+ self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+ def forward(self, x, context, context_lens, dtype):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ context(Tensor): Shape [B, L2, C]
+ context_lens(Tensor): Shape [B]
+ """
+ context_img = context[:, :257]
+ context = context[:, 257:]
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.norm_q(self.q(x.to(dtype))).view(b, -1, n, d)
+ k = self.norm_k(self.k(context.to(dtype))).view(b, -1, n, d)
+ v = self.v(context.to(dtype)).view(b, -1, n, d)
+ k_img = self.norm_k_img(self.k_img(context_img.to(dtype))).view(b, -1, n, d)
+ v_img = self.v_img(context_img.to(dtype)).view(b, -1, n, d)
+
+ img_x = attention(
+ q.to(dtype),
+ k_img.to(dtype),
+ v_img.to(dtype),
+ k_lens=None
+ )
+ img_x = img_x.to(dtype)
+ # compute attention
+ x = attention(
+ q.to(dtype),
+ k.to(dtype),
+ v.to(dtype),
+ k_lens=context_lens
+ )
+ x = x.to(dtype)
+
+ # output
+ x = x.flatten(2)
+ img_x = img_x.flatten(2)
+ x = x + img_x
+ x = self.o(x)
+ return x
+
+
+class WanI2VTalkingCrossAttention(WanSelfAttention):
+ def __init__(self,
+ dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ eps=1e-6,
+ audio_context_dim=1024
+ ):
+ super().__init__(dim, num_heads, window_size, qk_norm, eps)
+
+ # 14B should be audio_context_dim=2048
+ # 1.3B should be audio_context_dim=1024
+
+ # self.alpha = nn.Parameter(torch.zeros((1, )))
+ self.k_img = nn.Linear(dim, dim)
+ self.v_img = nn.Linear(dim, dim)
+ self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+ self.k_vocal = nn.Linear(dim, dim)
+ self.v_vocal = nn.Linear(dim, dim)
+
+ nn.init.zeros_(self.k_vocal.weight)
+ nn.init.zeros_(self.v_vocal.weight)
+ if self.k_vocal.bias is not None:
+ nn.init.zeros_(self.k_vocal.bias)
+ if self.v_vocal.bias is not None:
+ nn.init.zeros_(self.v_vocal.bias)
+
+
+ def forward(self, x, context, context_lens, dtype, vocal_context=None, vocal_context_lens=None):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ context(Tensor): Shape [B, L2, C]
+ context_lens(Tensor): Shape [B]
+ """
+
+ # print(f"The size of x in cross attention: {x.size()}") # [1, 21504, 5120]
+ # print(f"The size of context_clip in cross attention: {context_img.size()}") # [1, 257, 5120]
+ context_img = context[:, :257]
+ context = context[:, 257:]
+
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.norm_q(self.q(x.to(dtype))).view(b, -1, n, d)
+ k = self.norm_k(self.k(context.to(dtype))).view(b, -1, n, d)
+ v = self.v(context.to(dtype)).view(b, -1, n, d)
+ k_img = self.norm_k_img(self.k_img(context_img.to(dtype))).view(b, -1, n, d)
+ v_img = self.v_img(context_img.to(dtype)).view(b, -1, n, d)
+
+ img_x = attention(
+ q.to(dtype),
+ k_img.to(dtype),
+ v_img.to(dtype),
+ k_lens=None
+ )
+ img_x = img_x.to(dtype)
+ # compute attention
+ x = attention(
+ q.to(dtype),
+ k.to(dtype),
+ v.to(dtype),
+ k_lens=context_lens
+ )
+ x = x.to(dtype)
+
+ latents_num_frames = 21
+ if len(vocal_context.shape) == 4:
+ vocal_q = q.view(b * latents_num_frames, -1, n, d)
+ vocal_ip_key = self.k_vocal(vocal_context).view(b * latents_num_frames, -1, n, d)
+ vocal_ip_value = self.v_vocal(vocal_context).view(b * latents_num_frames, -1, n, d)
+ vocal_x = attention(
+ vocal_q.to(dtype),
+ vocal_ip_key.to(dtype),
+ vocal_ip_value.to(dtype),
+ k_lens=vocal_context_lens
+ )
+ vocal_x = vocal_x.view(b, q.size(1), n, d)
+ vocal_x = vocal_x.flatten(2)
+ else:
+ vocal_ip_key = self.k_vocal(vocal_context).view(b, -1, n, d)
+ vocal_ip_value = self.v_vocal(vocal_context).view(b, -1, n, d)
+ vocal_x = attention(
+ q.to(dtype),
+ vocal_ip_key.to(dtype),
+ vocal_ip_value.to(dtype),
+ k_lens=None,
+ )
+ vocal_x = vocal_x.flatten(2)
+
+
+ # output
+ x = x.flatten(2)
+ img_x = img_x.flatten(2)
+
+ x = x + img_x + vocal_x
+ x = self.o(x)
+ return x
+
+
+WAN_CROSSATTENTION_CLASSES = {
+ 't2v_cross_attn': WanT2VCrossAttention,
+ 'i2v_cross_attn': WanI2VCrossAttention,
+}
+
+
+class WanAttentionBlock(nn.Module):
+
+ def __init__(self,
+ cross_attn_type,
+ dim,
+ ffn_dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ cross_attn_norm=False,
+ eps=1e-6):
+ 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 = WanLayerNorm(dim, eps)
+ self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm,
+ eps)
+ self.norm3 = WanLayerNorm(
+ dim, eps,
+ elementwise_affine=True) if cross_attn_norm else nn.Identity()
+ # self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim, num_heads, (-1, -1), qk_norm, eps)
+ self.cross_attn = WanI2VTalkingCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps)
+ self.norm2 = WanLayerNorm(dim, eps)
+ self.ffn = nn.Sequential(
+ nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
+ nn.Linear(ffn_dim, dim))
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim ** 0.5)
+
+ def forward(
+ self,
+ x,
+ e,
+ seq_lens,
+ grid_sizes,
+ freqs,
+ context,
+ context_lens,
+ dtype=torch.float32,
+ vocal_context=None,
+ vocal_context_lens=None,
+ ):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ e(Tensor): Shape [B, 6, C]
+ seq_lens(Tensor): Shape [B], length of each sequence in batch
+ grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+ freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+ """
+ e = (self.modulation + e).chunk(6, dim=1)
+
+ # self-attention
+ temp_x = self.norm1(x) * (1 + e[1]) + e[0]
+ temp_x = temp_x.to(dtype)
+
+ y = self.self_attn(temp_x, seq_lens, grid_sizes, freqs, dtype)
+ x = x + y * e[2]
+
+ # cross-attention & ffn function
+ def cross_attn_ffn(x, context, context_lens, e, vocal_context, vocal_context_lens):
+ # cross-attention
+ x = x + self.cross_attn(self.norm3(x), context, context_lens, dtype, vocal_context, vocal_context_lens)
+
+ # ffn function
+ temp_x = self.norm2(x) * (1 + e[4]) + e[3]
+ temp_x = temp_x.to(dtype)
+
+ y = self.ffn(temp_x)
+ x = x + y * e[5]
+ return x
+
+ x = cross_attn_ffn(x, context, context_lens, e, vocal_context, vocal_context_lens)
+ 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 = WanLayerNorm(dim, eps)
+ self.head = nn.Linear(dim, out_dim)
+
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim ** 0.5)
+
+ def forward(self, x, e):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ e(Tensor): Shape [B, C]
+ """
+ e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
+ x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))
+ return x
+
+
+class MLPProj(torch.nn.Module):
+
+ def __init__(self, in_dim, out_dim):
+ super().__init__()
+
+ self.proj = torch.nn.Sequential(
+ torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim),
+ torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim),
+ torch.nn.LayerNorm(out_dim))
+
+ def forward(self, image_embeds):
+ clip_extra_context_tokens = self.proj(image_embeds)
+ return clip_extra_context_tokens
+
+
+class WanTransformer3DFantasy14BModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+ r"""
+ Wan diffusion backbone supporting both text-to-video and image-to-video.
+ """
+
+ # ignore_for_config = [
+ # 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size'
+ # ]
+ # _no_split_modules = ['WanAttentionBlock']
+ _supports_gradient_checkpointing = True
+
+ @register_to_config
+ def __init__(
+ self,
+ model_type='t2v',
+ 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,
+ in_channels=16,
+ hidden_size=2048,
+ ):
+ r"""
+ Initialize the diffusion model backbone.
+
+ Args:
+ model_type (`str`, *optional*, defaults to 't2v'):
+ Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+ patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+ 3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+ text_len (`int`, *optional*, defaults to 512):
+ Fixed length for text embeddings
+ in_dim (`int`, *optional*, defaults to 16):
+ Input video channels (C_in)
+ dim (`int`, *optional*, defaults to 2048):
+ Hidden dimension of the transformer
+ ffn_dim (`int`, *optional*, defaults to 8192):
+ Intermediate dimension in feed-forward network
+ freq_dim (`int`, *optional*, defaults to 256):
+ Dimension for sinusoidal time embeddings
+ text_dim (`int`, *optional*, defaults to 4096):
+ Input dimension for text embeddings
+ out_dim (`int`, *optional*, defaults to 16):
+ Output video channels (C_out)
+ num_heads (`int`, *optional*, defaults to 16):
+ Number of attention heads
+ num_layers (`int`, *optional*, defaults to 32):
+ Number of transformer blocks
+ window_size (`tuple`, *optional*, defaults to (-1, -1)):
+ Window size for local attention (-1 indicates global attention)
+ qk_norm (`bool`, *optional*, defaults to True):
+ Enable query/key normalization
+ cross_attn_norm (`bool`, *optional*, defaults to False):
+ Enable cross-attention normalization
+ eps (`float`, *optional*, defaults to 1e-6):
+ Epsilon value for normalization layers
+ """
+
+ super().__init__()
+
+ assert model_type in ['t2v', 'i2v']
+ 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
+ cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'
+ self.blocks = nn.ModuleList([
+ WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,
+ window_size, qk_norm, cross_attn_norm, eps)
+ for _ in range(num_layers)
+ ])
+
+ # head
+ self.head = Head(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.d = d
+ 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
+ )
+
+ if model_type == 'i2v':
+ self.img_emb = MLPProj(1280, dim)
+
+ self.teacache = None
+ self.gradient_checkpointing = False
+ self.sp_world_size = 1
+ self.sp_world_rank = 0
+
+ self.vocal_projector = FantasyTalkingVocalCondition14BModel(audio_in_dim=768, audio_proj_dim=dim, dit_dim=dim)
+
+ def enable_teacache(
+ self,
+ coefficients,
+ num_steps: int,
+ rel_l1_thresh: float,
+ num_skip_start_steps: int = 0,
+ offload: bool = True
+ ):
+ self.teacache = TeaCache(
+ coefficients, num_steps, rel_l1_thresh=rel_l1_thresh, num_skip_start_steps=num_skip_start_steps,
+ offload=offload
+ )
+
+ def disable_teacache(self):
+ self.teacache = None
+
+ def enable_riflex(
+ self,
+ k=6,
+ L_test=66,
+ L_test_scale=4.886,
+ ):
+ device = self.freqs.device
+ self.freqs = torch.cat(
+ [
+ get_1d_rotary_pos_embed_riflex(1024, self.d - 4 * (self.d // 6), use_real=False, k=k, L_test=L_test,
+ L_test_scale=L_test_scale),
+ rope_params(1024, 2 * (self.d // 6)),
+ rope_params(1024, 2 * (self.d // 6))
+ ],
+ dim=1
+ ).to(device)
+
+ def disable_riflex(self):
+ device = self.freqs.device
+ self.freqs = torch.cat(
+ [
+ rope_params(1024, self.d - 4 * (self.d // 6)),
+ rope_params(1024, 2 * (self.d // 6)),
+ rope_params(1024, 2 * (self.d // 6))
+ ],
+ dim=1
+ ).to(device)
+
+ def enable_multi_gpus_inference(self, ):
+ self.sp_world_size = get_sequence_parallel_world_size()
+ self.sp_world_rank = get_sequence_parallel_rank()
+ for block in self.blocks:
+ block.self_attn.forward = types.MethodType(
+ usp_attn_forward, block.self_attn)
+
+ def _set_gradient_checkpointing(self, module, value=False):
+ self.gradient_checkpointing = value
+
+ def forward(
+ self,
+ x,
+ t,
+ context,
+ seq_len,
+ clip_fea=None,
+ y=None,
+ cond_flag=True,
+ vocal_embeddings=None,
+ is_clip_level_modeling=False,
+ ):
+ r"""
+ Forward pass through the diffusion model
+
+ Args:
+ x (List[Tensor]):
+ List of input video tensors, each with shape [C_in, F, H, W]
+ t (Tensor):
+ Diffusion timesteps tensor of shape [B]
+ context (List[Tensor]):
+ List of text embeddings each with shape [L, C]
+ seq_len (`int`):
+ Maximum sequence length for positional encoding
+ clip_fea (Tensor, *optional*):
+ CLIP image features for image-to-video mode
+ y (List[Tensor], *optional*):
+ Conditional video inputs for image-to-video mode, same shape as x
+ cond_flag (`bool`, *optional*, defaults to True):
+ Flag to indicate whether to forward the condition input
+
+ Returns:
+ List[Tensor]:
+ List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+ """
+ if self.model_type == 'i2v':
+ assert clip_fea is not None and y is not None
+ # params
+ device = self.patch_embedding.weight.device
+ dtype = x.dtype
+ if self.freqs.device != device and torch.device(type="meta") != device:
+ self.freqs = self.freqs.to(device)
+
+ if y is not None:
+ x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+ # embeddings
+ x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+ 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)
+ if self.sp_world_size > 1:
+ seq_len = int(math.ceil(seq_len / self.sp_world_size)) * self.sp_world_size
+ assert seq_lens.max() <= seq_len
+ x = torch.cat([torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x])
+
+ # time embeddings
+ with amp.autocast(dtype=torch.float32):
+ e = self.time_embedding(sinusoidal_embedding_1d(self.freq_dim, t).float())
+ e0 = self.time_projection(e).unflatten(1, (6, self.dim))
+ e0 = e0.to(dtype)
+ e = e.to(dtype)
+
+ # 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
+ ]))
+ context_clip = self.img_emb(clip_fea) # bs x 257 x dim
+
+ # print("-----------------------------------------")
+ # print(f"motion scale: {motion_scale}") # 0.2
+ # print(f"The size of face_masks: {face_masks.size()}") # [1, 81, 1, 512, 512]
+ # print(f"The size of context: {context.size()}") # [1, 512, 5120]
+ # print(f"The size of context_clip: {context_clip.size()}") # [1, 257, 5120]
+ # print(f"The size of e: {e.size()}") # [1, 5120]
+ # print(f"The size of e0: {e0.size()}") # [1, 6, 5120]
+ # print(f"The size of vocal_context: {vocal_context.size()}") # [1, 21, 32, 5120]
+ # print(f"The size of audio_context: {audio_context.size()}") # [1, 21, 32, 5120]
+ # print("-----------------------------------------")
+
+
+ context = torch.concat([context_clip, context], dim=1)
+ vocal_context, vocal_context_lens = self.vocal_projector(vocal_embeddings=vocal_embeddings, video_sample_n_frames=81, latents=x, e0=e0, e=e)
+ if is_clip_level_modeling:
+ vocal_context = rearrange(vocal_context, "b f n c -> b (f n) c", f=21)
+ print("You are in the clip-level audio modeling")
+
+ # Context Parallel
+ if self.sp_world_size > 1:
+ x = torch.chunk(x, self.sp_world_size, dim=1)[self.sp_world_rank]
+
+ # TeaCache
+ if self.teacache is not None:
+ if cond_flag:
+ modulated_inp = e0
+ skip_flag = self.teacache.cnt < self.teacache.num_skip_start_steps
+ if self.teacache.cnt == 0 or self.teacache.cnt == self.teacache.num_steps - 1 or skip_flag:
+ should_calc = True
+ self.teacache.accumulated_rel_l1_distance = 0
+ else:
+ if cond_flag:
+ rel_l1_distance = self.teacache.compute_rel_l1_distance(self.teacache.previous_modulated_input,
+ modulated_inp)
+ self.teacache.accumulated_rel_l1_distance += self.teacache.rescale_func(rel_l1_distance)
+ if self.teacache.accumulated_rel_l1_distance < self.teacache.rel_l1_thresh:
+ should_calc = False
+ else:
+ should_calc = True
+ self.teacache.accumulated_rel_l1_distance = 0
+ self.teacache.previous_modulated_input = modulated_inp
+ self.teacache.cnt += 1
+ if self.teacache.cnt == self.teacache.num_steps:
+ self.teacache.reset()
+ self.teacache.should_calc = should_calc
+ else:
+ should_calc = self.teacache.should_calc
+
+ # TeaCache
+ if self.teacache is not None:
+ if not should_calc:
+ previous_residual = self.teacache.previous_residual_cond if cond_flag else self.teacache.previous_residual_uncond
+ x = x + previous_residual.to(x.device)
+ else:
+ ori_x = x.clone().cpu() if self.teacache.offload else x.clone()
+
+ for block in self.blocks:
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=",
+ "1.11.0") else {}
+ x = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ x,
+ e0,
+ seq_lens,
+ grid_sizes,
+ self.freqs,
+ context,
+ context_lens,
+ dtype,
+ vocal_context,
+ vocal_context_lens,
+ **ckpt_kwargs,
+ )
+ else:
+ # arguments
+ kwargs = dict(
+ e=e0,
+ seq_lens=seq_lens,
+ grid_sizes=grid_sizes,
+ freqs=self.freqs,
+ context=context,
+ context_lens=context_lens,
+ dtype=dtype,
+ vocal_context=vocal_context,
+ vocal_context_lens=vocal_context_lens,
+ )
+ x = block(x, **kwargs)
+
+ if cond_flag:
+ self.teacache.previous_residual_cond = x.cpu() - ori_x if self.teacache.offload else x - ori_x
+ else:
+ self.teacache.previous_residual_uncond = x.cpu() - ori_x if self.teacache.offload else x - ori_x
+ else:
+ for block in self.blocks:
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+ x = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ x,
+ e0,
+ seq_lens,
+ grid_sizes,
+ self.freqs,
+ context,
+ context_lens,
+ dtype,
+ vocal_context,
+ vocal_context_lens,
+ **ckpt_kwargs,
+ )
+ else:
+ # arguments
+ kwargs = dict(
+ e=e0,
+ seq_lens=seq_lens,
+ grid_sizes=grid_sizes,
+ freqs=self.freqs,
+ context=context,
+ context_lens=context_lens,
+ dtype=dtype,
+ vocal_context=vocal_context,
+ vocal_context_lens=vocal_context_lens,
+ )
+ x = block(x, **kwargs)
+
+ if self.sp_world_size > 1:
+ x = get_sp_group().all_gather(x, dim=1)
+
+ # head
+ x = self.head(x, e)
+
+ # unpatchify
+ x = self.unpatchify(x, grid_sizes)
+ x = torch.stack(x)
+ return x
+
+ def unpatchify(self, x, grid_sizes):
+ r"""
+ 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)
+
+ @classmethod
+ def from_pretrained(
+ cls, pretrained_model_path, subfolder=None, transformer_additional_kwargs={},
+ low_cpu_mem_usage=False, torch_dtype=torch.bfloat16
+ ):
+ if subfolder is not None:
+ pretrained_model_path = os.path.join(pretrained_model_path, subfolder)
+ print(f"loaded 3D transformer's pretrained weights from {pretrained_model_path} ...")
+
+ config_file = os.path.join(pretrained_model_path, 'config.json')
+ if not os.path.isfile(config_file):
+ raise RuntimeError(f"{config_file} does not exist")
+ with open(config_file, "r") as f:
+ config = json.load(f)
+
+ from diffusers.utils import WEIGHTS_NAME
+ model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME)
+ model_file_safetensors = model_file.replace(".bin", ".safetensors")
+
+ if "dict_mapping" in transformer_additional_kwargs.keys():
+ for key in transformer_additional_kwargs["dict_mapping"]:
+ transformer_additional_kwargs[transformer_additional_kwargs["dict_mapping"][key]] = config[key]
+
+ # {'patch_size', 'qk_norm', 'window_size', 'cross_attn_norm', 'text_dim'} was not found in config. Values will be initialized to default values.
+ transformer_additional_kwargs["patch_size"] = (1, 2, 2)
+ transformer_additional_kwargs["qk_norm"] = True
+ transformer_additional_kwargs["window_size"] = (-1, -1)
+ transformer_additional_kwargs["cross_attn_norm"] = True
+
+ if low_cpu_mem_usage:
+ try:
+ import re
+
+ from diffusers.models.modeling_utils import \
+ load_model_dict_into_meta
+ from diffusers.utils import is_accelerate_available
+ if is_accelerate_available():
+ import accelerate
+
+ # Instantiate model with empty weights
+ with accelerate.init_empty_weights():
+ model = cls.from_config(config, **transformer_additional_kwargs)
+
+ param_device = "cpu"
+ if os.path.exists(model_file):
+ state_dict = torch.load(model_file, map_location="cpu")
+ elif os.path.exists(model_file_safetensors):
+ from safetensors.torch import load_file, safe_open
+ state_dict = load_file(model_file_safetensors)
+ else:
+ from safetensors.torch import load_file, safe_open
+ model_files_safetensors = glob.glob(os.path.join(pretrained_model_path, "*.safetensors"))
+ state_dict = {}
+ print(model_files_safetensors)
+ for _model_file_safetensors in model_files_safetensors:
+ _state_dict = load_file(_model_file_safetensors)
+ for key in _state_dict:
+ state_dict[key] = _state_dict[key]
+ model._convert_deprecated_attention_blocks(state_dict)
+ # move the params from meta device to cpu
+ missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())
+ if len(missing_keys) > 0:
+ raise ValueError(
+ f"Cannot load {cls} from {pretrained_model_path} because the following keys are"
+ f" missing: \n {', '.join(missing_keys)}. \n Please make sure to pass"
+ " `low_cpu_mem_usage=False` and `device_map=None` if you want to randomly initialize"
+ " those weights or else make sure your checkpoint file is correct."
+ )
+
+ unexpected_keys = load_model_dict_into_meta(
+ model,
+ state_dict,
+ device=param_device,
+ dtype=torch_dtype,
+ model_name_or_path=pretrained_model_path,
+ )
+
+ if cls._keys_to_ignore_on_load_unexpected is not None:
+ for pat in cls._keys_to_ignore_on_load_unexpected:
+ unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]
+
+ if len(unexpected_keys) > 0:
+ print(
+ f"Some weights of the model checkpoint were not used when initializing {cls.__name__}: \n {[', '.join(unexpected_keys)]}"
+ )
+ return model
+ except Exception as e:
+ print(
+ f"The low_cpu_mem_usage mode is not work because {e}. Use low_cpu_mem_usage=False instead."
+ )
+
+ model = cls.from_config(config, **transformer_additional_kwargs)
+ if os.path.exists(model_file):
+ state_dict = torch.load(model_file, map_location="cpu")
+ elif os.path.exists(model_file_safetensors):
+ from safetensors.torch import load_file, safe_open
+ state_dict = load_file(model_file_safetensors)
+ else:
+ from safetensors.torch import load_file, safe_open
+ model_files_safetensors = glob.glob(os.path.join(pretrained_model_path, "*.safetensors"))
+ state_dict = {}
+ for _model_file_safetensors in model_files_safetensors:
+ _state_dict = load_file(_model_file_safetensors)
+ for key in _state_dict:
+ state_dict[key] = _state_dict[key]
+
+ if model.state_dict()['patch_embedding.weight'].size() != state_dict['patch_embedding.weight'].size():
+ model.state_dict()['patch_embedding.weight'][:, :state_dict['patch_embedding.weight'].size()[1], :, :] = \
+ state_dict['patch_embedding.weight']
+ model.state_dict()['patch_embedding.weight'][:, state_dict['patch_embedding.weight'].size()[1]:, :, :] = 0
+ state_dict['patch_embedding.weight'] = model.state_dict()['patch_embedding.weight']
+
+ tmp_state_dict = {}
+ for key in state_dict:
+ if key in model.state_dict().keys() and model.state_dict()[key].size() == state_dict[key].size():
+ tmp_state_dict[key] = state_dict[key]
+ else:
+ print(key, "Size don't match, skip")
+
+ state_dict = tmp_state_dict
+
+ m, u = model.load_state_dict(state_dict, strict=False)
+ print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
+ print(m)
+
+ params = [p.numel() if "." in n else 0 for n, p in model.named_parameters()]
+ print(f"### All Parameters: {sum(params) / 1e6} M")
+
+ model = model.to(torch_dtype)
+ return model
\ No newline at end of file
diff --git a/wan/models/wan_fantasy_transformer3d_1B.py b/wan/models/wan_fantasy_transformer3d_1B.py
new file mode 100644
index 0000000000000000000000000000000000000000..ad715865d64bed9b02f14699a8f4db749f226237
--- /dev/null
+++ b/wan/models/wan_fantasy_transformer3d_1B.py
@@ -0,0 +1,1321 @@
+# Modified from https://github.com/Wan-Video/Wan2.1/blob/main/wan/modules/model.py
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+
+import glob
+import json
+import math
+import os
+import types
+import warnings
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+import torch
+import torch.amp as amp
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import is_torch_version, logging
+from torch import nn
+from einops import rearrange
+
+from .vocal_projector_fantasy_1B import FantasyTalkingVocalCondition1BModel
+from ..dist import (get_sequence_parallel_rank,
+ get_sequence_parallel_world_size, get_sp_group,
+ xFuserLongContextAttention)
+
+from .cache_utils import TeaCache
+from ..dist.wan_xfuser import usp_attn_forward
+
+try:
+ import flash_attn_interface
+
+ FLASH_ATTN_3_AVAILABLE = True
+except ModuleNotFoundError:
+ FLASH_ATTN_3_AVAILABLE = False
+
+try:
+ import flash_attn
+
+ FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+ FLASH_ATTN_2_AVAILABLE = False
+
+FLASH_ATTN_2_AVAILABLE = False
+FLASH_ATTN_3_AVAILABLE = False
+
+
+def flash_attention(
+ q,
+ k,
+ v,
+ q_lens=None,
+ k_lens=None,
+ dropout_p=0.,
+ softmax_scale=None,
+ q_scale=None,
+ causal=False,
+ window_size=(-1, -1),
+ deterministic=False,
+ dtype=torch.bfloat16,
+ version=None,
+):
+ """
+ q: [B, Lq, Nq, C1].
+ k: [B, Lk, Nk, C1].
+ v: [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+ q_lens: [B].
+ k_lens: [B].
+ dropout_p: float. Dropout probability.
+ softmax_scale: float. The scaling of QK^T before applying softmax.
+ causal: bool. Whether to apply causal attention mask.
+ window_size: (left right). If not (-1, -1), apply sliding window local attention.
+ deterministic: bool. If True, slightly slower and uses more memory.
+ dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+ """
+ half_dtypes = (torch.float16, torch.bfloat16)
+ assert dtype in half_dtypes
+ assert q.device.type == 'cuda' and q.size(-1) <= 256
+
+ # params
+ b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+ def half(x):
+ return x if x.dtype in half_dtypes else x.to(dtype)
+
+ # preprocess query
+ if q_lens is None:
+ q = half(q.flatten(0, 1))
+ q_lens = torch.tensor(
+ [lq] * b, dtype=torch.int32).to(
+ device=q.device, non_blocking=True)
+ else:
+ q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
+
+ # preprocess key, value
+ if k_lens is None:
+ k = half(k.flatten(0, 1))
+ v = half(v.flatten(0, 1))
+ k_lens = torch.tensor(
+ [lk] * b, dtype=torch.int32).to(
+ device=k.device, non_blocking=True)
+ else:
+ k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
+ v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
+
+ q = q.to(v.dtype)
+ k = k.to(v.dtype)
+
+ if q_scale is not None:
+ q = q * q_scale
+
+ if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+ warnings.warn(
+ 'Flash attention 3 is not available, use flash attention 2 instead.'
+ )
+
+ # apply attention
+ if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+ # Note: dropout_p, window_size are not supported in FA3 now.
+ x = flash_attn_interface.flash_attn_varlen_func(
+ q=q,
+ k=k,
+ v=v,
+ cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ seqused_q=None,
+ seqused_k=None,
+ max_seqlen_q=lq,
+ max_seqlen_k=lk,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ deterministic=deterministic)[0].unflatten(0, (b, lq))
+ else:
+ assert FLASH_ATTN_2_AVAILABLE
+ x = flash_attn.flash_attn_varlen_func(
+ q=q,
+ k=k,
+ v=v,
+ cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ max_seqlen_q=lq,
+ max_seqlen_k=lk,
+ dropout_p=dropout_p,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ window_size=window_size,
+ deterministic=deterministic).unflatten(0, (b, lq))
+
+ # output
+ return x.type(out_dtype)
+
+
+def attention(
+ q,
+ k,
+ v,
+ q_lens=None,
+ k_lens=None,
+ dropout_p=0.,
+ softmax_scale=None,
+ q_scale=None,
+ causal=False,
+ window_size=(-1, -1),
+ deterministic=False,
+ dtype=torch.bfloat16,
+ fa_version=None,
+):
+ if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+ return flash_attention(
+ q=q,
+ k=k,
+ v=v,
+ q_lens=q_lens,
+ k_lens=k_lens,
+ dropout_p=dropout_p,
+ softmax_scale=softmax_scale,
+ q_scale=q_scale,
+ causal=causal,
+ window_size=window_size,
+ deterministic=deterministic,
+ dtype=dtype,
+ version=fa_version,
+ )
+ else:
+ if q_lens is not None or k_lens is not None:
+ warnings.warn(
+ 'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.'
+ )
+ attn_mask = None
+
+ q = q.transpose(1, 2)
+ k = k.transpose(1, 2)
+ v = v.transpose(1, 2)
+
+ if torch.backends.cuda.flash_sdp_enabled() is False or torch.backends.cuda.enable_flash_sdp is False:
+ print(1/0)
+
+ out = torch.nn.functional.scaled_dot_product_attention(
+ q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p)
+
+ out = out.transpose(1, 2).contiguous()
+ return out
+
+
+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('cuda', 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
+
+
+# modified from https://github.com/thu-ml/RIFLEx/blob/main/riflex_utils.py
+@amp.autocast('cuda', enabled=False)
+def get_1d_rotary_pos_embed_riflex(
+ pos: Union[np.ndarray, int],
+ dim: int,
+ theta: float = 10000.0,
+ use_real=False,
+ k: Optional[int] = None,
+ L_test: Optional[int] = None,
+ L_test_scale: Optional[int] = None,
+):
+ """
+ RIFLEx: Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
+
+ This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' and the end
+ index 'end'. The 'theta' parameter scales the frequencies. The returned tensor contains complex values in complex64
+ data type.
+
+ Args:
+ dim (`int`): Dimension of the frequency tensor.
+ pos (`np.ndarray` or `int`): Position indices for the frequency tensor. [S] or scalar
+ theta (`float`, *optional*, defaults to 10000.0):
+ Scaling factor for frequency computation. Defaults to 10000.0.
+ use_real (`bool`, *optional*):
+ If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+ k (`int`, *optional*, defaults to None): the index for the intrinsic frequency in RoPE
+ L_test (`int`, *optional*, defaults to None): the number of frames for inference
+ Returns:
+ `torch.Tensor`: Precomputed frequency tensor with complex exponentials. [S, D/2]
+ """
+ assert dim % 2 == 0
+
+ if isinstance(pos, int):
+ pos = torch.arange(pos)
+ if isinstance(pos, np.ndarray):
+ pos = torch.from_numpy(pos) # type: ignore # [S]
+
+ freqs = 1.0 / torch.pow(theta,
+ torch.arange(0, dim, 2).to(torch.float64).div(dim))
+
+ # === Riflex modification start ===
+ # Reduce the intrinsic frequency to stay within a single period after extrapolation (see Eq. (8)).
+ # Empirical observations show that a few videos may exhibit repetition in the tail frames.
+ # To be conservative, we multiply by 0.9 to keep the extrapolated length below 90% of a single period.
+ if k is not None:
+ freqs[k - 1] = 0.9 * 2 * torch.pi / L_test
+ # === Riflex modification end ===
+ if L_test_scale is not None:
+ freqs[k - 1] = freqs[k - 1] / L_test_scale
+
+ freqs = torch.outer(pos, freqs) # type: ignore # [S, D/2]
+ if use_real:
+ freqs_cos = freqs.cos().repeat_interleave(2, dim=1).float() # [S, D]
+ freqs_sin = freqs.sin().repeat_interleave(2, dim=1).float() # [S, D]
+ return freqs_cos, freqs_sin
+ else:
+ # lumina
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 # [S, D/2]
+ return freqs_cis
+
+
+@amp.autocast('cuda', enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+ n, c = x.size(2), x.size(3) // 2
+
+ # split freqs
+ freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+ # loop over samples
+ output = []
+ for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+ seq_len = f * h * w
+
+ # precompute multipliers
+ x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float32).reshape(
+ seq_len, n, -1, 2))
+ freqs_i = torch.cat([
+ freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+ freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+ freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+ ],
+ dim=-1).reshape(seq_len, 1, -1)
+
+ # apply rotary embedding
+ x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+ x_i = torch.cat([x_i, x[i, seq_len:]])
+
+ # append to collection
+ output.append(x_i)
+ return torch.stack(output).float()
+
+
+class WanRMSNorm(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):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ """
+ 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 WanLayerNorm(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):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ """
+ return super().forward(x.float()).type_as(x)
+
+
+class WanSelfAttention(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 = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+ self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+ def forward(self, x, seq_lens, grid_sizes, freqs, dtype):
+ r"""
+ 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.to(dtype))).view(b, s, n, d)
+ k = self.norm_k(self.k(x.to(dtype))).view(b, s, n, d)
+ v = self.v(x.to(dtype)).view(b, s, n, d)
+ return q, k, v
+
+ q, k, v = qkv_fn(x)
+
+ x = attention(
+ q=rope_apply(q, grid_sizes, freqs).to(dtype),
+ k=rope_apply(k, grid_sizes, freqs).to(dtype),
+ v=v.to(dtype),
+ k_lens=seq_lens,
+ window_size=self.window_size)
+ x = x.to(dtype)
+
+ # output
+ x = x.flatten(2)
+ x = self.o(x)
+ return x
+
+
+class WanT2VCrossAttention(WanSelfAttention):
+
+ def forward(self, x, context, context_lens, dtype):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ context(Tensor): Shape [B, L2, C]
+ context_lens(Tensor): Shape [B]
+ """
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.norm_q(self.q(x.to(dtype))).view(b, -1, n, d)
+ k = self.norm_k(self.k(context.to(dtype))).view(b, -1, n, d)
+ v = self.v(context.to(dtype)).view(b, -1, n, d)
+
+ # compute attention
+ x = attention(
+ q.to(dtype),
+ k.to(dtype),
+ v.to(dtype),
+ k_lens=context_lens
+ )
+ x = x.to(dtype)
+
+ # output
+ x = x.flatten(2)
+ x = self.o(x)
+ return x
+
+
+class WanI2VCrossAttention(WanSelfAttention):
+
+ def __init__(self,
+ dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ eps=1e-6):
+ super().__init__(dim, num_heads, window_size, qk_norm, eps)
+
+ # self.alpha = nn.Parameter(torch.zeros((1, )))
+ self.k_img = nn.Linear(dim, dim)
+ self.v_img = nn.Linear(dim, dim)
+ self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+ def forward(self, x, context, context_lens, dtype):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ context(Tensor): Shape [B, L2, C]
+ context_lens(Tensor): Shape [B]
+ """
+ context_img = context[:, :257]
+ context = context[:, 257:]
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.norm_q(self.q(x.to(dtype))).view(b, -1, n, d)
+ k = self.norm_k(self.k(context.to(dtype))).view(b, -1, n, d)
+ v = self.v(context.to(dtype)).view(b, -1, n, d)
+ k_img = self.norm_k_img(self.k_img(context_img.to(dtype))).view(b, -1, n, d)
+ v_img = self.v_img(context_img.to(dtype)).view(b, -1, n, d)
+
+ img_x = attention(
+ q.to(dtype),
+ k_img.to(dtype),
+ v_img.to(dtype),
+ k_lens=None
+ )
+ img_x = img_x.to(dtype)
+ # compute attention
+ x = attention(
+ q.to(dtype),
+ k.to(dtype),
+ v.to(dtype),
+ k_lens=context_lens
+ )
+ x = x.to(dtype)
+
+ # output
+ x = x.flatten(2)
+ img_x = img_x.flatten(2)
+ x = x + img_x
+ x = self.o(x)
+ return x
+
+
+class WanI2VTalkingCrossAttention(WanSelfAttention):
+ def __init__(self,
+ dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ eps=1e-6,
+ audio_context_dim=1024
+ ):
+ super().__init__(dim, num_heads, window_size, qk_norm, eps)
+
+ # 14B should be audio_context_dim=2048
+ # 1.3B should be audio_context_dim=1024
+
+ # self.alpha = nn.Parameter(torch.zeros((1, )))
+ self.k_img = nn.Linear(dim, dim)
+ self.v_img = nn.Linear(dim, dim)
+ self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+ self.k_vocal = nn.Linear(dim, dim)
+ self.v_vocal = nn.Linear(dim, dim)
+
+ nn.init.zeros_(self.k_vocal.weight)
+ nn.init.zeros_(self.v_vocal.weight)
+ if self.k_vocal.bias is not None:
+ nn.init.zeros_(self.k_vocal.bias)
+ if self.v_vocal.bias is not None:
+ nn.init.zeros_(self.v_vocal.bias)
+
+
+ def forward(self, x, context, context_lens, dtype, vocal_context=None, vocal_context_lens=None):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ context(Tensor): Shape [B, L2, C]
+ context_lens(Tensor): Shape [B]
+ """
+
+ # print(f"The size of x in cross attention: {x.size()}") # [1, 21504, 5120]
+ # print(f"The size of context_clip in cross attention: {context_img.size()}") # [1, 257, 5120]
+ context_img = context[:, :257]
+ context = context[:, 257:]
+
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.norm_q(self.q(x.to(dtype))).view(b, -1, n, d)
+ k = self.norm_k(self.k(context.to(dtype))).view(b, -1, n, d)
+ v = self.v(context.to(dtype)).view(b, -1, n, d)
+ k_img = self.norm_k_img(self.k_img(context_img.to(dtype))).view(b, -1, n, d)
+ v_img = self.v_img(context_img.to(dtype)).view(b, -1, n, d)
+
+ img_x = attention(
+ q.to(dtype),
+ k_img.to(dtype),
+ v_img.to(dtype),
+ k_lens=None
+ )
+ img_x = img_x.to(dtype)
+ # compute attention
+ x = attention(
+ q.to(dtype),
+ k.to(dtype),
+ v.to(dtype),
+ k_lens=context_lens
+ )
+ x = x.to(dtype)
+
+ latents_num_frames = 21
+ if len(vocal_context.shape) == 4:
+ vocal_q = q.view(b * latents_num_frames, -1, n, d)
+ vocal_ip_key = self.k_vocal(vocal_context).view(b * latents_num_frames, -1, n, d)
+ vocal_ip_value = self.v_vocal(vocal_context).view(b * latents_num_frames, -1, n, d)
+ vocal_x = attention(
+ vocal_q.to(dtype),
+ vocal_ip_key.to(dtype),
+ vocal_ip_value.to(dtype),
+ k_lens=vocal_context_lens
+ )
+ vocal_x = vocal_x.view(b, q.size(1), n, d)
+ vocal_x = vocal_x.flatten(2)
+ else:
+ vocal_ip_key = self.k_vocal(vocal_context).view(b, -1, n, d)
+ vocal_ip_value = self.v_vocal(vocal_context).view(b, -1, n, d)
+ vocal_x = attention(
+ q.to(dtype),
+ vocal_ip_key.to(dtype),
+ vocal_ip_value.to(dtype),
+ k_lens=None,
+ )
+ vocal_x = vocal_x.flatten(2)
+
+
+ # output
+ x = x.flatten(2)
+ img_x = img_x.flatten(2)
+
+ x = x + img_x + vocal_x
+ x = self.o(x)
+ return x
+
+
+WAN_CROSSATTENTION_CLASSES = {
+ 't2v_cross_attn': WanT2VCrossAttention,
+ 'i2v_cross_attn': WanI2VCrossAttention,
+}
+
+
+class WanAttentionBlock(nn.Module):
+
+ def __init__(self,
+ cross_attn_type,
+ dim,
+ ffn_dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ cross_attn_norm=False,
+ eps=1e-6):
+ 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 = WanLayerNorm(dim, eps)
+ self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm,
+ eps)
+ self.norm3 = WanLayerNorm(
+ dim, eps,
+ elementwise_affine=True) if cross_attn_norm else nn.Identity()
+ # self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim, num_heads, (-1, -1), qk_norm, eps)
+ self.cross_attn = WanI2VTalkingCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps)
+ self.norm2 = WanLayerNorm(dim, eps)
+ self.ffn = nn.Sequential(
+ nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
+ nn.Linear(ffn_dim, dim))
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim ** 0.5)
+
+ def forward(
+ self,
+ x,
+ e,
+ seq_lens,
+ grid_sizes,
+ freqs,
+ context,
+ context_lens,
+ dtype=torch.float32,
+ vocal_context=None,
+ vocal_context_lens=None,
+ ):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ e(Tensor): Shape [B, 6, C]
+ seq_lens(Tensor): Shape [B], length of each sequence in batch
+ grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+ freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+ """
+ e = (self.modulation + e).chunk(6, dim=1)
+
+ # self-attention
+ temp_x = self.norm1(x) * (1 + e[1]) + e[0]
+ temp_x = temp_x.to(dtype)
+
+ y = self.self_attn(temp_x, seq_lens, grid_sizes, freqs, dtype)
+ x = x + y * e[2]
+
+ # cross-attention & ffn function
+ def cross_attn_ffn(x, context, context_lens, e, vocal_context, vocal_context_lens):
+ # cross-attention
+ x = x + self.cross_attn(self.norm3(x), context, context_lens, dtype, vocal_context, vocal_context_lens)
+
+ # ffn function
+ temp_x = self.norm2(x) * (1 + e[4]) + e[3]
+ temp_x = temp_x.to(dtype)
+
+ y = self.ffn(temp_x)
+ x = x + y * e[5]
+ return x
+
+ x = cross_attn_ffn(x, context, context_lens, e, vocal_context, vocal_context_lens)
+ 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 = WanLayerNorm(dim, eps)
+ self.head = nn.Linear(dim, out_dim)
+
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim ** 0.5)
+
+ def forward(self, x, e):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ e(Tensor): Shape [B, C]
+ """
+ e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
+ x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))
+ return x
+
+
+class MLPProj(torch.nn.Module):
+
+ def __init__(self, in_dim, out_dim):
+ super().__init__()
+
+ self.proj = torch.nn.Sequential(
+ torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim),
+ torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim),
+ torch.nn.LayerNorm(out_dim))
+
+ def forward(self, image_embeds):
+ clip_extra_context_tokens = self.proj(image_embeds)
+ return clip_extra_context_tokens
+
+
+class WanTransformer3DFantasyModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+ r"""
+ Wan diffusion backbone supporting both text-to-video and image-to-video.
+ """
+
+ # ignore_for_config = [
+ # 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size'
+ # ]
+ # _no_split_modules = ['WanAttentionBlock']
+ _supports_gradient_checkpointing = True
+
+ @register_to_config
+ def __init__(
+ self,
+ model_type='t2v',
+ 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,
+ in_channels=16,
+ hidden_size=2048,
+ ):
+ r"""
+ Initialize the diffusion model backbone.
+
+ Args:
+ model_type (`str`, *optional*, defaults to 't2v'):
+ Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+ patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+ 3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+ text_len (`int`, *optional*, defaults to 512):
+ Fixed length for text embeddings
+ in_dim (`int`, *optional*, defaults to 16):
+ Input video channels (C_in)
+ dim (`int`, *optional*, defaults to 2048):
+ Hidden dimension of the transformer
+ ffn_dim (`int`, *optional*, defaults to 8192):
+ Intermediate dimension in feed-forward network
+ freq_dim (`int`, *optional*, defaults to 256):
+ Dimension for sinusoidal time embeddings
+ text_dim (`int`, *optional*, defaults to 4096):
+ Input dimension for text embeddings
+ out_dim (`int`, *optional*, defaults to 16):
+ Output video channels (C_out)
+ num_heads (`int`, *optional*, defaults to 16):
+ Number of attention heads
+ num_layers (`int`, *optional*, defaults to 32):
+ Number of transformer blocks
+ window_size (`tuple`, *optional*, defaults to (-1, -1)):
+ Window size for local attention (-1 indicates global attention)
+ qk_norm (`bool`, *optional*, defaults to True):
+ Enable query/key normalization
+ cross_attn_norm (`bool`, *optional*, defaults to False):
+ Enable cross-attention normalization
+ eps (`float`, *optional*, defaults to 1e-6):
+ Epsilon value for normalization layers
+ """
+
+ super().__init__()
+
+ assert model_type in ['t2v', 'i2v']
+ 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
+ cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'
+ self.blocks = nn.ModuleList([
+ WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,
+ window_size, qk_norm, cross_attn_norm, eps)
+ for _ in range(num_layers)
+ ])
+
+ # head
+ self.head = Head(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.d = d
+ 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
+ )
+
+ if model_type == 'i2v':
+ self.img_emb = MLPProj(1280, dim)
+
+ self.teacache = None
+ self.gradient_checkpointing = False
+ self.sp_world_size = 1
+ self.sp_world_rank = 0
+
+ self.vocal_projector = FantasyTalkingVocalCondition1BModel(audio_in_dim=768, audio_proj_dim=1536, dit_dim=dim)
+
+ def enable_teacache(
+ self,
+ coefficients,
+ num_steps: int,
+ rel_l1_thresh: float,
+ num_skip_start_steps: int = 0,
+ offload: bool = True
+ ):
+ self.teacache = TeaCache(
+ coefficients, num_steps, rel_l1_thresh=rel_l1_thresh, num_skip_start_steps=num_skip_start_steps,
+ offload=offload
+ )
+
+ def disable_teacache(self):
+ self.teacache = None
+
+ def enable_riflex(
+ self,
+ k=6,
+ L_test=66,
+ L_test_scale=4.886,
+ ):
+ device = self.freqs.device
+ self.freqs = torch.cat(
+ [
+ get_1d_rotary_pos_embed_riflex(1024, self.d - 4 * (self.d // 6), use_real=False, k=k, L_test=L_test,
+ L_test_scale=L_test_scale),
+ rope_params(1024, 2 * (self.d // 6)),
+ rope_params(1024, 2 * (self.d // 6))
+ ],
+ dim=1
+ ).to(device)
+
+ def disable_riflex(self):
+ device = self.freqs.device
+ self.freqs = torch.cat(
+ [
+ rope_params(1024, self.d - 4 * (self.d // 6)),
+ rope_params(1024, 2 * (self.d // 6)),
+ rope_params(1024, 2 * (self.d // 6))
+ ],
+ dim=1
+ ).to(device)
+
+ def enable_multi_gpus_inference(self, ):
+ self.sp_world_size = get_sequence_parallel_world_size()
+ self.sp_world_rank = get_sequence_parallel_rank()
+ for block in self.blocks:
+ block.self_attn.forward = types.MethodType(
+ usp_attn_forward, block.self_attn)
+
+ def _set_gradient_checkpointing(self, module, value=False):
+ self.gradient_checkpointing = value
+
+ def forward(
+ self,
+ x,
+ t,
+ context,
+ seq_len,
+ clip_fea=None,
+ y=None,
+ cond_flag=True,
+ vocal_embeddings=None,
+ is_clip_level_modeling=False,
+ ):
+ r"""
+ Forward pass through the diffusion model
+
+ Args:
+ x (List[Tensor]):
+ List of input video tensors, each with shape [C_in, F, H, W]
+ t (Tensor):
+ Diffusion timesteps tensor of shape [B]
+ context (List[Tensor]):
+ List of text embeddings each with shape [L, C]
+ seq_len (`int`):
+ Maximum sequence length for positional encoding
+ clip_fea (Tensor, *optional*):
+ CLIP image features for image-to-video mode
+ y (List[Tensor], *optional*):
+ Conditional video inputs for image-to-video mode, same shape as x
+ cond_flag (`bool`, *optional*, defaults to True):
+ Flag to indicate whether to forward the condition input
+
+ Returns:
+ List[Tensor]:
+ List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+ """
+ if self.model_type == 'i2v':
+ assert clip_fea is not None and y is not None
+ # params
+ device = self.patch_embedding.weight.device
+ dtype = x.dtype
+ if self.freqs.device != device and torch.device(type="meta") != device:
+ self.freqs = self.freqs.to(device)
+
+ if y is not None:
+ x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+ # embeddings
+ x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+ 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)
+ if self.sp_world_size > 1:
+ seq_len = int(math.ceil(seq_len / self.sp_world_size)) * self.sp_world_size
+ assert seq_lens.max() <= seq_len
+ x = torch.cat([torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x])
+
+ # time embeddings
+ with amp.autocast('cuda', 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))
+ e0 = e0.to(dtype)
+ e = e.to(dtype)
+
+ # 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
+ ]))
+ context_clip = self.img_emb(clip_fea) # bs x 257 x dim
+
+ context = torch.concat([context_clip, context], dim=1)
+
+ if vocal_embeddings.size()[0] > 1:
+ vocal_embeddings = vocal_embeddings[-1:]
+ vocal_context, vocal_context_lens = self.vocal_projector(vocal_embeddings=vocal_embeddings, video_sample_n_frames=81, latents=x[-1:], e0=e0[-1:], e=e[-1:])
+ vocal_context = torch.cat([torch.zeros_like(vocal_context), vocal_context, vocal_context])
+ else:
+ vocal_context, vocal_context_lens = self.vocal_projector(vocal_embeddings=vocal_embeddings, video_sample_n_frames=81, latents=x, e0=e0, e=e)
+
+ if is_clip_level_modeling:
+ vocal_context = rearrange(vocal_context, "b f n c -> b (f n) c", f=21)
+ print("You are in the clip-level audio modeling")
+
+ # Context Parallel
+ if self.sp_world_size > 1:
+ x = torch.chunk(x, self.sp_world_size, dim=1)[self.sp_world_rank]
+
+ # TeaCache
+ if self.teacache is not None:
+ if cond_flag:
+ modulated_inp = e0
+ skip_flag = self.teacache.cnt < self.teacache.num_skip_start_steps
+ if self.teacache.cnt == 0 or self.teacache.cnt == self.teacache.num_steps - 1 or skip_flag:
+ should_calc = True
+ self.teacache.accumulated_rel_l1_distance = 0
+ else:
+ if cond_flag:
+ rel_l1_distance = self.teacache.compute_rel_l1_distance(self.teacache.previous_modulated_input, modulated_inp)
+ self.teacache.accumulated_rel_l1_distance += self.teacache.rescale_func(rel_l1_distance)
+ if self.teacache.accumulated_rel_l1_distance < self.teacache.rel_l1_thresh:
+ should_calc = False
+ else:
+ should_calc = True
+ self.teacache.accumulated_rel_l1_distance = 0
+ self.teacache.previous_modulated_input = modulated_inp
+ self.teacache.cnt += 1
+ if self.teacache.cnt == self.teacache.num_steps:
+ self.teacache.reset()
+ self.teacache.should_calc = should_calc
+ else:
+ should_calc = self.teacache.should_calc
+
+ # TeaCache
+ if self.teacache is not None:
+ if not should_calc:
+ previous_residual = self.teacache.previous_residual_cond if cond_flag else self.teacache.previous_residual_uncond
+ x = x + previous_residual.to(x.device)
+ else:
+ ori_x = x.clone().cpu() if self.teacache.offload else x.clone()
+
+ for block in self.blocks:
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=",
+ "1.11.0") else {}
+ x = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ x,
+ e0,
+ seq_lens,
+ grid_sizes,
+ self.freqs,
+ context,
+ context_lens,
+ dtype,
+ vocal_context,
+ vocal_context_lens,
+ **ckpt_kwargs,
+ )
+ else:
+ # arguments
+ kwargs = dict(
+ e=e0,
+ seq_lens=seq_lens,
+ grid_sizes=grid_sizes,
+ freqs=self.freqs,
+ context=context,
+ context_lens=context_lens,
+ dtype=dtype,
+ vocal_context=vocal_context,
+ vocal_context_lens=vocal_context_lens,
+ )
+ x = block(x, **kwargs)
+
+ if cond_flag:
+ self.teacache.previous_residual_cond = x.cpu() - ori_x if self.teacache.offload else x - ori_x
+ else:
+ self.teacache.previous_residual_uncond = x.cpu() - ori_x if self.teacache.offload else x - ori_x
+ else:
+ for block in self.blocks:
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+ x = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ x,
+ e0,
+ seq_lens,
+ grid_sizes,
+ self.freqs,
+ context,
+ context_lens,
+ dtype,
+ vocal_context,
+ vocal_context_lens,
+ **ckpt_kwargs,
+ )
+ else:
+ # arguments
+ kwargs = dict(
+ e=e0,
+ seq_lens=seq_lens,
+ grid_sizes=grid_sizes,
+ freqs=self.freqs,
+ context=context,
+ context_lens=context_lens,
+ dtype=dtype,
+ vocal_context=vocal_context,
+ vocal_context_lens=vocal_context_lens,
+ )
+ x = block(x, **kwargs)
+
+ if self.sp_world_size > 1:
+ x = get_sp_group().all_gather(x, dim=1)
+
+ # head
+ x = self.head(x, e)
+
+ # unpatchify
+ x = self.unpatchify(x, grid_sizes)
+ x = torch.stack(x)
+ return x
+
+ def unpatchify(self, x, grid_sizes):
+ r"""
+ 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)
+
+ @classmethod
+ def from_pretrained(
+ cls, pretrained_model_path, subfolder=None, transformer_additional_kwargs={},
+ low_cpu_mem_usage=False, torch_dtype=torch.bfloat16
+ ):
+ if subfolder is not None:
+ pretrained_model_path = os.path.join(pretrained_model_path, subfolder)
+ print(f"loaded 3D transformer's pretrained weights from {pretrained_model_path} ...")
+
+ config_file = os.path.join(pretrained_model_path, 'config.json')
+ if not os.path.isfile(config_file):
+ raise RuntimeError(f"{config_file} does not exist")
+ with open(config_file, "r") as f:
+ config = json.load(f)
+
+ from diffusers.utils import WEIGHTS_NAME
+ model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME)
+ model_file_safetensors = model_file.replace(".bin", ".safetensors")
+
+ if "dict_mapping" in transformer_additional_kwargs.keys():
+ for key in transformer_additional_kwargs["dict_mapping"]:
+ transformer_additional_kwargs[transformer_additional_kwargs["dict_mapping"][key]] = config[key]
+
+ # {'patch_size', 'qk_norm', 'window_size', 'cross_attn_norm', 'text_dim'} was not found in config. Values will be initialized to default values.
+ transformer_additional_kwargs["patch_size"] = (1, 2, 2)
+ transformer_additional_kwargs["qk_norm"] = True
+ transformer_additional_kwargs["window_size"] = (-1, -1)
+ transformer_additional_kwargs["cross_attn_norm"] = True
+
+ if low_cpu_mem_usage:
+ try:
+ import re
+
+ from diffusers.models.modeling_utils import \
+ load_model_dict_into_meta
+ from diffusers.utils import is_accelerate_available
+ if is_accelerate_available():
+ import accelerate
+
+ # Instantiate model with empty weights
+ with accelerate.init_empty_weights():
+ model = cls.from_config(config, **transformer_additional_kwargs)
+
+ param_device = "cpu"
+ if os.path.exists(model_file):
+ state_dict = torch.load(model_file, map_location="cpu")
+ elif os.path.exists(model_file_safetensors):
+ from safetensors.torch import load_file, safe_open
+ state_dict = load_file(model_file_safetensors)
+ else:
+ from safetensors.torch import load_file, safe_open
+ model_files_safetensors = glob.glob(os.path.join(pretrained_model_path, "*.safetensors"))
+ state_dict = {}
+ print(model_files_safetensors)
+ for _model_file_safetensors in model_files_safetensors:
+ _state_dict = load_file(_model_file_safetensors)
+ for key in _state_dict:
+ state_dict[key] = _state_dict[key]
+ model._convert_deprecated_attention_blocks(state_dict)
+ # move the params from meta device to cpu
+ missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())
+ if len(missing_keys) > 0:
+ raise ValueError(
+ f"Cannot load {cls} from {pretrained_model_path} because the following keys are"
+ f" missing: \n {', '.join(missing_keys)}. \n Please make sure to pass"
+ " `low_cpu_mem_usage=False` and `device_map=None` if you want to randomly initialize"
+ " those weights or else make sure your checkpoint file is correct."
+ )
+
+ unexpected_keys = load_model_dict_into_meta(
+ model,
+ state_dict,
+ device=param_device,
+ dtype=torch_dtype,
+ model_name_or_path=pretrained_model_path,
+ )
+
+ if cls._keys_to_ignore_on_load_unexpected is not None:
+ for pat in cls._keys_to_ignore_on_load_unexpected:
+ unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]
+
+ if len(unexpected_keys) > 0:
+ print(
+ f"Some weights of the model checkpoint were not used when initializing {cls.__name__}: \n {[', '.join(unexpected_keys)]}"
+ )
+ return model
+ except Exception as e:
+ print(
+ f"The low_cpu_mem_usage mode is not work because {e}. Use low_cpu_mem_usage=False instead."
+ )
+
+ model = cls.from_config(config, **transformer_additional_kwargs)
+ if os.path.exists(model_file):
+ state_dict = torch.load(model_file, map_location="cpu")
+ elif os.path.exists(model_file_safetensors):
+ from safetensors.torch import load_file, safe_open
+ state_dict = load_file(model_file_safetensors)
+ else:
+ from safetensors.torch import load_file, safe_open
+ model_files_safetensors = glob.glob(os.path.join(pretrained_model_path, "*.safetensors"))
+ state_dict = {}
+ for _model_file_safetensors in model_files_safetensors:
+ _state_dict = load_file(_model_file_safetensors)
+ for key in _state_dict:
+ state_dict[key] = _state_dict[key]
+
+ if model.state_dict()['patch_embedding.weight'].size() != state_dict['patch_embedding.weight'].size():
+ model.state_dict()['patch_embedding.weight'][:, :state_dict['patch_embedding.weight'].size()[1], :, :] = \
+ state_dict['patch_embedding.weight']
+ model.state_dict()['patch_embedding.weight'][:, state_dict['patch_embedding.weight'].size()[1]:, :, :] = 0
+ state_dict['patch_embedding.weight'] = model.state_dict()['patch_embedding.weight']
+
+ tmp_state_dict = {}
+ for key in state_dict:
+ if key in model.state_dict().keys() and model.state_dict()[key].size() == state_dict[key].size():
+ tmp_state_dict[key] = state_dict[key]
+ else:
+ print(key, "Size don't match, skip")
+
+ state_dict = tmp_state_dict
+
+ m, u = model.load_state_dict(state_dict, strict=False)
+ print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
+ print(m)
+
+ params = [p.numel() if "." in n else 0 for n, p in model.named_parameters()]
+ print(f"### All Parameters: {sum(params) / 1e6} M")
+
+ model = model.to(torch_dtype)
+ return model
\ No newline at end of file
diff --git a/wan/models/wan_image_encoder.py b/wan/models/wan_image_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..950b4cc49b8bc539cbf0d97c8bf68927d62d4d04
--- /dev/null
+++ b/wan/models/wan_image_encoder.py
@@ -0,0 +1,553 @@
+# Modified from ``https://github.com/openai/CLIP'' and ``https://github.com/mlfoundations/open_clip''
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms as T
+
+from .wan_transformer3d import attention
+from .wan_xlm_roberta import XLMRoberta
+from diffusers.configuration_utils import ConfigMixin
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.models.modeling_utils import ModelMixin
+
+
+__all__ = [
+ 'XLMRobertaCLIP',
+ 'clip_xlm_roberta_vit_h_14',
+ 'CLIPModel',
+]
+
+
+def pos_interpolate(pos, seq_len):
+ if pos.size(1) == seq_len:
+ return pos
+ else:
+ src_grid = int(math.sqrt(pos.size(1)))
+ tar_grid = int(math.sqrt(seq_len))
+ n = pos.size(1) - src_grid * src_grid
+ return torch.cat([
+ pos[:, :n],
+ F.interpolate(
+ pos[:, n:].float().reshape(1, src_grid, src_grid, -1).permute(
+ 0, 3, 1, 2),
+ size=(tar_grid, tar_grid),
+ mode='bicubic',
+ align_corners=False).flatten(2).transpose(1, 2)
+ ],
+ dim=1)
+
+
+class QuickGELU(nn.Module):
+
+ def forward(self, x):
+ return x * torch.sigmoid(1.702 * x)
+
+
+class LayerNorm(nn.LayerNorm):
+
+ def forward(self, x):
+ return super().forward(x.float()).type_as(x)
+
+
+class SelfAttention(nn.Module):
+
+ def __init__(self,
+ dim,
+ num_heads,
+ causal=False,
+ attn_dropout=0.0,
+ proj_dropout=0.0):
+ assert dim % num_heads == 0
+ super().__init__()
+ self.dim = dim
+ self.num_heads = num_heads
+ self.head_dim = dim // num_heads
+ self.causal = causal
+ self.attn_dropout = attn_dropout
+ self.proj_dropout = proj_dropout
+
+ # layers
+ self.to_qkv = nn.Linear(dim, dim * 3)
+ self.proj = nn.Linear(dim, dim)
+
+ def forward(self, x):
+ """
+ x: [B, L, C].
+ """
+ b, s, c, n, d = *x.size(), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q, k, v = self.to_qkv(x).view(b, s, 3, n, d).unbind(2)
+
+ # compute attention
+ p = self.attn_dropout if self.training else 0.0
+ x = attention(q, k, v, dropout_p=p, causal=self.causal)
+ x = x.reshape(b, s, c)
+
+ # output
+ x = self.proj(x)
+ x = F.dropout(x, self.proj_dropout, self.training)
+ return x
+
+
+class SwiGLU(nn.Module):
+
+ def __init__(self, dim, mid_dim):
+ super().__init__()
+ self.dim = dim
+ self.mid_dim = mid_dim
+
+ # layers
+ self.fc1 = nn.Linear(dim, mid_dim)
+ self.fc2 = nn.Linear(dim, mid_dim)
+ self.fc3 = nn.Linear(mid_dim, dim)
+
+ def forward(self, x):
+ x = F.silu(self.fc1(x)) * self.fc2(x)
+ x = self.fc3(x)
+ return x
+
+
+class AttentionBlock(nn.Module):
+
+ def __init__(self,
+ dim,
+ mlp_ratio,
+ num_heads,
+ post_norm=False,
+ causal=False,
+ activation='quick_gelu',
+ attn_dropout=0.0,
+ proj_dropout=0.0,
+ norm_eps=1e-5):
+ assert activation in ['quick_gelu', 'gelu', 'swi_glu']
+ super().__init__()
+ self.dim = dim
+ self.mlp_ratio = mlp_ratio
+ self.num_heads = num_heads
+ self.post_norm = post_norm
+ self.causal = causal
+ self.norm_eps = norm_eps
+
+ # layers
+ self.norm1 = LayerNorm(dim, eps=norm_eps)
+ self.attn = SelfAttention(dim, num_heads, causal, attn_dropout,
+ proj_dropout)
+ self.norm2 = LayerNorm(dim, eps=norm_eps)
+ if activation == 'swi_glu':
+ self.mlp = SwiGLU(dim, int(dim * mlp_ratio))
+ else:
+ self.mlp = nn.Sequential(
+ nn.Linear(dim, int(dim * mlp_ratio)),
+ QuickGELU() if activation == 'quick_gelu' else nn.GELU(),
+ nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout))
+
+ def forward(self, x):
+ if self.post_norm:
+ x = x + self.norm1(self.attn(x))
+ x = x + self.norm2(self.mlp(x))
+ else:
+ x = x + self.attn(self.norm1(x))
+ x = x + self.mlp(self.norm2(x))
+ return x
+
+
+class AttentionPool(nn.Module):
+
+ def __init__(self,
+ dim,
+ mlp_ratio,
+ num_heads,
+ activation='gelu',
+ proj_dropout=0.0,
+ norm_eps=1e-5):
+ assert dim % num_heads == 0
+ super().__init__()
+ self.dim = dim
+ self.mlp_ratio = mlp_ratio
+ self.num_heads = num_heads
+ self.head_dim = dim // num_heads
+ self.proj_dropout = proj_dropout
+ self.norm_eps = norm_eps
+
+ # layers
+ gain = 1.0 / math.sqrt(dim)
+ self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))
+ self.to_q = nn.Linear(dim, dim)
+ self.to_kv = nn.Linear(dim, dim * 2)
+ self.proj = nn.Linear(dim, dim)
+ self.norm = LayerNorm(dim, eps=norm_eps)
+ self.mlp = nn.Sequential(
+ nn.Linear(dim, int(dim * mlp_ratio)),
+ QuickGELU() if activation == 'quick_gelu' else nn.GELU(),
+ nn.Linear(int(dim * mlp_ratio), dim), nn.Dropout(proj_dropout))
+
+ def forward(self, x):
+ """
+ x: [B, L, C].
+ """
+ b, s, c, n, d = *x.size(), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.to_q(self.cls_embedding).view(1, 1, n, d).expand(b, -1, -1, -1)
+ k, v = self.to_kv(x).view(b, s, 2, n, d).unbind(2)
+
+ # compute attention
+ x = flash_attention(q, k, v, version=2)
+ x = x.reshape(b, 1, c)
+
+ # output
+ x = self.proj(x)
+ x = F.dropout(x, self.proj_dropout, self.training)
+
+ # mlp
+ x = x + self.mlp(self.norm(x))
+ return x[:, 0]
+
+
+class VisionTransformer(nn.Module):
+
+ def __init__(self,
+ image_size=224,
+ patch_size=16,
+ dim=768,
+ mlp_ratio=4,
+ out_dim=512,
+ num_heads=12,
+ num_layers=12,
+ pool_type='token',
+ pre_norm=True,
+ post_norm=False,
+ activation='quick_gelu',
+ attn_dropout=0.0,
+ proj_dropout=0.0,
+ embedding_dropout=0.0,
+ norm_eps=1e-5):
+ if image_size % patch_size != 0:
+ print(
+ '[WARNING] image_size is not divisible by patch_size',
+ flush=True)
+ assert pool_type in ('token', 'token_fc', 'attn_pool')
+ out_dim = out_dim or dim
+ super().__init__()
+ self.image_size = image_size
+ self.patch_size = patch_size
+ self.num_patches = (image_size // patch_size)**2
+ self.dim = dim
+ self.mlp_ratio = mlp_ratio
+ self.out_dim = out_dim
+ self.num_heads = num_heads
+ self.num_layers = num_layers
+ self.pool_type = pool_type
+ self.post_norm = post_norm
+ self.norm_eps = norm_eps
+
+ # embeddings
+ gain = 1.0 / math.sqrt(dim)
+ self.patch_embedding = nn.Conv2d(
+ 3,
+ dim,
+ kernel_size=patch_size,
+ stride=patch_size,
+ bias=not pre_norm)
+ if pool_type in ('token', 'token_fc'):
+ self.cls_embedding = nn.Parameter(gain * torch.randn(1, 1, dim))
+ self.pos_embedding = nn.Parameter(gain * torch.randn(
+ 1, self.num_patches +
+ (1 if pool_type in ('token', 'token_fc') else 0), dim))
+ self.dropout = nn.Dropout(embedding_dropout)
+
+ # transformer
+ self.pre_norm = LayerNorm(dim, eps=norm_eps) if pre_norm else None
+ self.transformer = nn.Sequential(*[
+ AttentionBlock(dim, mlp_ratio, num_heads, post_norm, False,
+ activation, attn_dropout, proj_dropout, norm_eps)
+ for _ in range(num_layers)
+ ])
+ self.post_norm = LayerNorm(dim, eps=norm_eps)
+
+ # head
+ if pool_type == 'token':
+ self.head = nn.Parameter(gain * torch.randn(dim, out_dim))
+ elif pool_type == 'token_fc':
+ self.head = nn.Linear(dim, out_dim)
+ elif pool_type == 'attn_pool':
+ self.head = AttentionPool(dim, mlp_ratio, num_heads, activation,
+ proj_dropout, norm_eps)
+
+ def forward(self, x, interpolation=False, use_31_block=False):
+ b = x.size(0)
+
+ # embeddings
+ x = self.patch_embedding(x).flatten(2).permute(0, 2, 1)
+ if self.pool_type in ('token', 'token_fc'):
+ x = torch.cat([self.cls_embedding.expand(b, -1, -1), x], dim=1)
+ if interpolation:
+ e = pos_interpolate(self.pos_embedding, x.size(1))
+ else:
+ e = self.pos_embedding
+ x = self.dropout(x + e)
+ if self.pre_norm is not None:
+ x = self.pre_norm(x)
+
+ # transformer
+ if use_31_block:
+ x = self.transformer[:-1](x)
+ return x
+ else:
+ x = self.transformer(x)
+ return x
+
+
+class XLMRobertaWithHead(XLMRoberta):
+
+ def __init__(self, **kwargs):
+ self.out_dim = kwargs.pop('out_dim')
+ super().__init__(**kwargs)
+
+ # head
+ mid_dim = (self.dim + self.out_dim) // 2
+ self.head = nn.Sequential(
+ nn.Linear(self.dim, mid_dim, bias=False), nn.GELU(),
+ nn.Linear(mid_dim, self.out_dim, bias=False))
+
+ def forward(self, ids):
+ # xlm-roberta
+ x = super().forward(ids)
+
+ # average pooling
+ mask = ids.ne(self.pad_id).unsqueeze(-1).to(x)
+ x = (x * mask).sum(dim=1) / mask.sum(dim=1)
+
+ # head
+ x = self.head(x)
+ return x
+
+
+class XLMRobertaCLIP(nn.Module):
+
+ def __init__(self,
+ embed_dim=1024,
+ image_size=224,
+ patch_size=14,
+ vision_dim=1280,
+ vision_mlp_ratio=4,
+ vision_heads=16,
+ vision_layers=32,
+ vision_pool='token',
+ vision_pre_norm=True,
+ vision_post_norm=False,
+ activation='gelu',
+ vocab_size=250002,
+ max_text_len=514,
+ type_size=1,
+ pad_id=1,
+ text_dim=1024,
+ text_heads=16,
+ text_layers=24,
+ text_post_norm=True,
+ text_dropout=0.1,
+ attn_dropout=0.0,
+ proj_dropout=0.0,
+ embedding_dropout=0.0,
+ norm_eps=1e-5):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.image_size = image_size
+ self.patch_size = patch_size
+ self.vision_dim = vision_dim
+ self.vision_mlp_ratio = vision_mlp_ratio
+ self.vision_heads = vision_heads
+ self.vision_layers = vision_layers
+ self.vision_pre_norm = vision_pre_norm
+ self.vision_post_norm = vision_post_norm
+ self.activation = activation
+ self.vocab_size = vocab_size
+ self.max_text_len = max_text_len
+ self.type_size = type_size
+ self.pad_id = pad_id
+ self.text_dim = text_dim
+ self.text_heads = text_heads
+ self.text_layers = text_layers
+ self.text_post_norm = text_post_norm
+ self.norm_eps = norm_eps
+
+ # models
+ self.visual = VisionTransformer(
+ image_size=image_size,
+ patch_size=patch_size,
+ dim=vision_dim,
+ mlp_ratio=vision_mlp_ratio,
+ out_dim=embed_dim,
+ num_heads=vision_heads,
+ num_layers=vision_layers,
+ pool_type=vision_pool,
+ pre_norm=vision_pre_norm,
+ post_norm=vision_post_norm,
+ activation=activation,
+ attn_dropout=attn_dropout,
+ proj_dropout=proj_dropout,
+ embedding_dropout=embedding_dropout,
+ norm_eps=norm_eps)
+ self.textual = XLMRobertaWithHead(
+ vocab_size=vocab_size,
+ max_seq_len=max_text_len,
+ type_size=type_size,
+ pad_id=pad_id,
+ dim=text_dim,
+ out_dim=embed_dim,
+ num_heads=text_heads,
+ num_layers=text_layers,
+ post_norm=text_post_norm,
+ dropout=text_dropout)
+ self.log_scale = nn.Parameter(math.log(1 / 0.07) * torch.ones([]))
+
+ def forward(self, imgs, txt_ids):
+ """
+ imgs: [B, 3, H, W] of torch.float32.
+ - mean: [0.48145466, 0.4578275, 0.40821073]
+ - std: [0.26862954, 0.26130258, 0.27577711]
+ txt_ids: [B, L] of torch.long.
+ Encoded by data.CLIPTokenizer.
+ """
+ xi = self.visual(imgs)
+ xt = self.textual(txt_ids)
+ return xi, xt
+
+ def param_groups(self):
+ groups = [{
+ 'params': [
+ p for n, p in self.named_parameters()
+ if 'norm' in n or n.endswith('bias')
+ ],
+ 'weight_decay': 0.0
+ }, {
+ 'params': [
+ p for n, p in self.named_parameters()
+ if not ('norm' in n or n.endswith('bias'))
+ ]
+ }]
+ return groups
+
+
+def _clip(pretrained=False,
+ pretrained_name=None,
+ model_cls=XLMRobertaCLIP,
+ return_transforms=False,
+ return_tokenizer=False,
+ tokenizer_padding='eos',
+ dtype=torch.float32,
+ device='cpu',
+ **kwargs):
+ # init a model on device
+ with torch.device(device):
+ model = model_cls(**kwargs)
+
+ # set device
+ model = model.to(dtype=dtype, device=device)
+ output = (model,)
+
+ # init transforms
+ if return_transforms:
+ # mean and std
+ if 'siglip' in pretrained_name.lower():
+ mean, std = [0.5, 0.5, 0.5], [0.5, 0.5, 0.5]
+ else:
+ mean = [0.48145466, 0.4578275, 0.40821073]
+ std = [0.26862954, 0.26130258, 0.27577711]
+
+ # transforms
+ transforms = T.Compose([
+ T.Resize((model.image_size, model.image_size),
+ interpolation=T.InterpolationMode.BICUBIC),
+ T.ToTensor(),
+ T.Normalize(mean=mean, std=std)
+ ])
+ output += (transforms,)
+ return output[0] if len(output) == 1 else output
+
+
+def clip_xlm_roberta_vit_h_14(
+ pretrained=False,
+ pretrained_name='open-clip-xlm-roberta-large-vit-huge-14',
+ **kwargs):
+ cfg = dict(
+ embed_dim=1024,
+ image_size=224,
+ patch_size=14,
+ vision_dim=1280,
+ vision_mlp_ratio=4,
+ vision_heads=16,
+ vision_layers=32,
+ vision_pool='token',
+ activation='gelu',
+ vocab_size=250002,
+ max_text_len=514,
+ type_size=1,
+ pad_id=1,
+ text_dim=1024,
+ text_heads=16,
+ text_layers=24,
+ text_post_norm=True,
+ text_dropout=0.1,
+ attn_dropout=0.0,
+ proj_dropout=0.0,
+ embedding_dropout=0.0)
+ cfg.update(**kwargs)
+ return _clip(pretrained, pretrained_name, XLMRobertaCLIP, **cfg)
+
+
+class CLIPModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+
+ def __init__(self):
+ super(CLIPModel, self).__init__()
+ # init model
+ self.model, self.transforms = clip_xlm_roberta_vit_h_14(
+ pretrained=False,
+ return_transforms=True,
+ return_tokenizer=False)
+
+ def forward(self, videos):
+ # preprocess
+ size = (self.model.image_size,) * 2
+ videos = torch.cat([
+ F.interpolate(
+ u.transpose(0, 1),
+ size=size,
+ mode='bicubic',
+ align_corners=False) for u in videos
+ ])
+ videos = self.transforms.transforms[-1](videos.mul_(0.5).add_(0.5))
+
+ # forward
+ with torch.amp.autocast('cuda', dtype=self.dtype):
+ out = self.model.visual(videos, use_31_block=True)
+ return out
+
+ @classmethod
+ def from_pretrained(cls, pretrained_model_path, transformer_additional_kwargs={}):
+ def filter_kwargs(cls, kwargs):
+ import inspect
+ sig = inspect.signature(cls.__init__)
+ valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
+ filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
+ return filtered_kwargs
+
+ model = cls(**filter_kwargs(cls, transformer_additional_kwargs))
+ if pretrained_model_path.endswith(".safetensors"):
+ from safetensors.torch import load_file, safe_open
+ state_dict = load_file(pretrained_model_path)
+ else:
+ state_dict = torch.load(pretrained_model_path, map_location="cpu")
+ tmp_state_dict = {}
+ for key in state_dict:
+ tmp_state_dict["model." + key] = state_dict[key]
+ state_dict = tmp_state_dict
+ m, u = model.load_state_dict(state_dict)
+
+ print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
+ print(m, u)
+ return model
\ No newline at end of file
diff --git a/wan/models/wan_text_encoder.py b/wan/models/wan_text_encoder.py
new file mode 100644
index 0000000000000000000000000000000000000000..34a032318cfba414cbce1ff7f0f7f4b406f452f3
--- /dev/null
+++ b/wan/models/wan_text_encoder.py
@@ -0,0 +1,365 @@
+# Modified from https://github.com/Wan-Video/Wan2.1/blob/main/wan/modules/t5.py
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import math
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from diffusers.configuration_utils import ConfigMixin
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.models.modeling_utils import ModelMixin
+
+
+def fp16_clamp(x):
+ if x.dtype == torch.float16 and torch.isinf(x).any():
+ clamp = torch.finfo(x.dtype).max - 1000
+ x = torch.clamp(x, min=-clamp, max=clamp)
+ return x
+
+
+def init_weights(m):
+ if isinstance(m, T5LayerNorm):
+ nn.init.ones_(m.weight)
+ elif isinstance(m, T5FeedForward):
+ nn.init.normal_(m.gate[0].weight, std=m.dim**-0.5)
+ nn.init.normal_(m.fc1.weight, std=m.dim**-0.5)
+ nn.init.normal_(m.fc2.weight, std=m.dim_ffn**-0.5)
+ elif isinstance(m, T5Attention):
+ nn.init.normal_(m.q.weight, std=(m.dim * m.dim_attn)**-0.5)
+ nn.init.normal_(m.k.weight, std=m.dim**-0.5)
+ nn.init.normal_(m.v.weight, std=m.dim**-0.5)
+ nn.init.normal_(m.o.weight, std=(m.num_heads * m.dim_attn)**-0.5)
+ elif isinstance(m, T5RelativeEmbedding):
+ nn.init.normal_(
+ m.embedding.weight, std=(2 * m.num_buckets * m.num_heads)**-0.5)
+
+
+class GELU(nn.Module):
+ def forward(self, x):
+ return 0.5 * x * (1.0 + torch.tanh(
+ math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
+
+
+class T5LayerNorm(nn.Module):
+ def __init__(self, dim, eps=1e-6):
+ super(T5LayerNorm, self).__init__()
+ self.dim = dim
+ self.eps = eps
+ self.weight = nn.Parameter(torch.ones(dim))
+
+ def forward(self, x):
+ x = x * torch.rsqrt(x.float().pow(2).mean(dim=-1, keepdim=True) +
+ self.eps)
+ if self.weight.dtype in [torch.float16, torch.bfloat16]:
+ x = x.type_as(self.weight)
+ return self.weight * x
+
+
+class T5Attention(nn.Module):
+ def __init__(self, dim, dim_attn, num_heads, dropout=0.1):
+ assert dim_attn % num_heads == 0
+ super(T5Attention, self).__init__()
+ self.dim = dim
+ self.dim_attn = dim_attn
+ self.num_heads = num_heads
+ self.head_dim = dim_attn // num_heads
+
+ # layers
+ self.q = nn.Linear(dim, dim_attn, bias=False)
+ self.k = nn.Linear(dim, dim_attn, bias=False)
+ self.v = nn.Linear(dim, dim_attn, bias=False)
+ self.o = nn.Linear(dim_attn, dim, bias=False)
+ self.dropout = nn.Dropout(dropout)
+
+ def forward(self, x, context=None, mask=None, pos_bias=None):
+ """
+ x: [B, L1, C].
+ context: [B, L2, C] or None.
+ mask: [B, L2] or [B, L1, L2] or None.
+ """
+ # check inputs
+ context = x if context is None else context
+ b, n, c = x.size(0), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.q(x).view(b, -1, n, c)
+ k = self.k(context).view(b, -1, n, c)
+ v = self.v(context).view(b, -1, n, c)
+
+ # attention bias
+ attn_bias = x.new_zeros(b, n, q.size(1), k.size(1))
+ if pos_bias is not None:
+ attn_bias += pos_bias
+ if mask is not None:
+ assert mask.ndim in [2, 3]
+ mask = mask.view(b, 1, 1,
+ -1) if mask.ndim == 2 else mask.unsqueeze(1)
+ attn_bias.masked_fill_(mask == 0, torch.finfo(x.dtype).min)
+
+ # compute attention (T5 does not use scaling)
+ attn = torch.einsum('binc,bjnc->bnij', q, k) + attn_bias
+ attn = F.softmax(attn.float(), dim=-1).type_as(attn)
+ x = torch.einsum('bnij,bjnc->binc', attn, v)
+
+ # output
+ x = x.reshape(b, -1, n * c)
+ x = self.o(x)
+ x = self.dropout(x)
+ return x
+
+
+class T5FeedForward(nn.Module):
+
+ def __init__(self, dim, dim_ffn, dropout=0.1):
+ super(T5FeedForward, self).__init__()
+ self.dim = dim
+ self.dim_ffn = dim_ffn
+
+ # layers
+ self.gate = nn.Sequential(nn.Linear(dim, dim_ffn, bias=False), GELU())
+ self.fc1 = nn.Linear(dim, dim_ffn, bias=False)
+ self.fc2 = nn.Linear(dim_ffn, dim, bias=False)
+ self.dropout = nn.Dropout(dropout)
+
+ def forward(self, x):
+ x = self.fc1(x) * self.gate(x)
+ x = self.dropout(x)
+ x = self.fc2(x)
+ x = self.dropout(x)
+ return x
+
+
+class T5SelfAttention(nn.Module):
+ def __init__(self,
+ dim,
+ dim_attn,
+ dim_ffn,
+ num_heads,
+ num_buckets,
+ shared_pos=True,
+ dropout=0.1):
+ super(T5SelfAttention, self).__init__()
+ self.dim = dim
+ self.dim_attn = dim_attn
+ self.dim_ffn = dim_ffn
+ self.num_heads = num_heads
+ self.num_buckets = num_buckets
+ self.shared_pos = shared_pos
+
+ # layers
+ self.norm1 = T5LayerNorm(dim)
+ self.attn = T5Attention(dim, dim_attn, num_heads, dropout)
+ self.norm2 = T5LayerNorm(dim)
+ self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+ self.pos_embedding = None if shared_pos else T5RelativeEmbedding(
+ num_buckets, num_heads, bidirectional=True)
+
+ def forward(self, x, mask=None, pos_bias=None):
+ e = pos_bias if self.shared_pos else self.pos_embedding(
+ x.size(1), x.size(1))
+ x = fp16_clamp(x + self.attn(self.norm1(x), mask=mask, pos_bias=e))
+ x = fp16_clamp(x + self.ffn(self.norm2(x)))
+ return x
+
+
+class T5CrossAttention(nn.Module):
+ def __init__(self,
+ dim,
+ dim_attn,
+ dim_ffn,
+ num_heads,
+ num_buckets,
+ shared_pos=True,
+ dropout=0.1):
+ super(T5CrossAttention, self).__init__()
+ self.dim = dim
+ self.dim_attn = dim_attn
+ self.dim_ffn = dim_ffn
+ self.num_heads = num_heads
+ self.num_buckets = num_buckets
+ self.shared_pos = shared_pos
+
+ # layers
+ self.norm1 = T5LayerNorm(dim)
+ self.self_attn = T5Attention(dim, dim_attn, num_heads, dropout)
+ self.norm2 = T5LayerNorm(dim)
+ self.cross_attn = T5Attention(dim, dim_attn, num_heads, dropout)
+ self.norm3 = T5LayerNorm(dim)
+ self.ffn = T5FeedForward(dim, dim_ffn, dropout)
+ self.pos_embedding = None if shared_pos else T5RelativeEmbedding(
+ num_buckets, num_heads, bidirectional=False)
+
+ def forward(self,
+ x,
+ mask=None,
+ encoder_states=None,
+ encoder_mask=None,
+ pos_bias=None):
+ e = pos_bias if self.shared_pos else self.pos_embedding(
+ x.size(1), x.size(1))
+ x = fp16_clamp(x + self.self_attn(self.norm1(x), mask=mask, pos_bias=e))
+ x = fp16_clamp(x + self.cross_attn(
+ self.norm2(x), context=encoder_states, mask=encoder_mask))
+ x = fp16_clamp(x + self.ffn(self.norm3(x)))
+ return x
+
+
+class T5RelativeEmbedding(nn.Module):
+ def __init__(self, num_buckets, num_heads, bidirectional, max_dist=128):
+ super(T5RelativeEmbedding, self).__init__()
+ self.num_buckets = num_buckets
+ self.num_heads = num_heads
+ self.bidirectional = bidirectional
+ self.max_dist = max_dist
+
+ # layers
+ self.embedding = nn.Embedding(num_buckets, num_heads)
+
+ def forward(self, lq, lk):
+ device = self.embedding.weight.device
+ # rel_pos = torch.arange(lk).unsqueeze(0).to(device) - \
+ # torch.arange(lq).unsqueeze(1).to(device)
+ if torch.device(type="meta") != device:
+ rel_pos = torch.arange(lk, device=device).unsqueeze(0) - \
+ torch.arange(lq, device=device).unsqueeze(1)
+ else:
+ rel_pos = torch.arange(lk).unsqueeze(0) - \
+ torch.arange(lq).unsqueeze(1)
+ rel_pos = self._relative_position_bucket(rel_pos)
+ rel_pos_embeds = self.embedding(rel_pos)
+ rel_pos_embeds = rel_pos_embeds.permute(2, 0, 1).unsqueeze(
+ 0) # [1, N, Lq, Lk]
+ return rel_pos_embeds.contiguous()
+
+ def _relative_position_bucket(self, rel_pos):
+ # preprocess
+ if self.bidirectional:
+ num_buckets = self.num_buckets // 2
+ rel_buckets = (rel_pos > 0).long() * num_buckets
+ rel_pos = torch.abs(rel_pos)
+ else:
+ num_buckets = self.num_buckets
+ rel_buckets = 0
+ rel_pos = -torch.min(rel_pos, torch.zeros_like(rel_pos))
+
+ # embeddings for small and large positions
+ max_exact = num_buckets // 2
+ rel_pos_large = max_exact + (torch.log(rel_pos.float() / max_exact) /
+ math.log(self.max_dist / max_exact) *
+ (num_buckets - max_exact)).long()
+ rel_pos_large = torch.min(
+ rel_pos_large, torch.full_like(rel_pos_large, num_buckets - 1))
+ rel_buckets += torch.where(rel_pos < max_exact, rel_pos, rel_pos_large)
+ return rel_buckets
+
+class WanT5EncoderModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+ def __init__(self,
+ vocab,
+ dim,
+ dim_attn,
+ dim_ffn,
+ num_heads,
+ num_layers,
+ num_buckets,
+ shared_pos=True,
+ dropout=0.1):
+ super(WanT5EncoderModel, self).__init__()
+ self.dim = dim
+ self.dim_attn = dim_attn
+ self.dim_ffn = dim_ffn
+ self.num_heads = num_heads
+ self.num_layers = num_layers
+ self.num_buckets = num_buckets
+ self.shared_pos = shared_pos
+
+ # layers
+ self.token_embedding = vocab if isinstance(vocab, nn.Embedding) \
+ else nn.Embedding(vocab, dim)
+ self.pos_embedding = T5RelativeEmbedding(
+ num_buckets, num_heads, bidirectional=True) if shared_pos else None
+ self.dropout = nn.Dropout(dropout)
+ self.blocks = nn.ModuleList([
+ T5SelfAttention(dim, dim_attn, dim_ffn, num_heads, num_buckets,
+ shared_pos, dropout) for _ in range(num_layers)
+ ])
+ self.norm = T5LayerNorm(dim)
+
+ # initialize weights
+ self.apply(init_weights)
+
+ def forward(
+ self,
+ input_ids: Optional[torch.LongTensor] = None,
+ attention_mask: Optional[torch.FloatTensor] = None,
+ ):
+ x = self.token_embedding(input_ids)
+ x = self.dropout(x)
+ e = self.pos_embedding(x.size(1),
+ x.size(1)) if self.shared_pos else None
+ for block in self.blocks:
+ x = block(x, attention_mask, pos_bias=e)
+ x = self.norm(x)
+ x = self.dropout(x)
+ return (x, )
+
+ @classmethod
+ def from_pretrained(cls, pretrained_model_path, additional_kwargs={}, low_cpu_mem_usage=False, torch_dtype=torch.bfloat16):
+ def filter_kwargs(cls, kwargs):
+ import inspect
+ sig = inspect.signature(cls.__init__)
+ valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
+ filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
+ return filtered_kwargs
+
+ if low_cpu_mem_usage:
+ try:
+ import re
+
+ from diffusers.models.modeling_utils import \
+ load_model_dict_into_meta
+ from diffusers.utils import is_accelerate_available
+ if is_accelerate_available():
+ import accelerate
+
+ # Instantiate model with empty weights
+ with accelerate.init_empty_weights():
+ model = cls(**filter_kwargs(cls, additional_kwargs))
+
+ param_device = "cpu"
+ if pretrained_model_path.endswith(".safetensors"):
+ from safetensors.torch import load_file
+ state_dict = load_file(pretrained_model_path)
+ else:
+ state_dict = torch.load(pretrained_model_path, map_location="cpu")
+ # move the params from meta device to cpu
+ missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())
+ if len(missing_keys) > 0:
+ raise ValueError(
+ f"Cannot load {cls} from {pretrained_model_path} because the following keys are"
+ f" missing: \n {', '.join(missing_keys)}. \n Please make sure to pass"
+ " `low_cpu_mem_usage=False` and `device_map=None` if you want to randomly initialize"
+ " those weights or else make sure your checkpoint file is correct."
+ )
+
+ unexpected_keys = load_model_dict_into_meta(
+ model,
+ state_dict,
+ device=param_device,
+ dtype=torch_dtype,
+ model_name_or_path=pretrained_model_path,
+ )
+
+ if cls._keys_to_ignore_on_load_unexpected is not None:
+ for pat in cls._keys_to_ignore_on_load_unexpected:
+ unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]
+
+ if len(unexpected_keys) > 0:
+ print(
+ f"Some weights of the model checkpoint were not used when initializing {cls.__name__}: \n {[', '.join(unexpected_keys)]}"
+ )
+ return model
+ except Exception as e:
+ print(
+ f"The low_cpu_mem_usage mode is not work because {e}. Use low_cpu_mem_usage=False instead."
+ )
\ No newline at end of file
diff --git a/wan/models/wan_transformer3d.py b/wan/models/wan_transformer3d.py
new file mode 100644
index 0000000000000000000000000000000000000000..3507488e991a4671a62f23056eb5caf4d888f0a5
--- /dev/null
+++ b/wan/models/wan_transformer3d.py
@@ -0,0 +1,1391 @@
+# Modified from https://github.com/Wan-Video/Wan2.1/blob/main/wan/modules/model.py
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+
+import glob
+import json
+import math
+import os
+import types
+import warnings
+from typing import Any, Dict, Optional, Union
+
+import numpy as np
+import torch
+import torch.amp as amp
+import torch.nn as nn
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils import is_torch_version, logging
+from torch import nn
+from einops import rearrange
+
+from ..dist import (get_sequence_parallel_rank,
+ get_sequence_parallel_world_size, get_sp_group,
+ xFuserLongContextAttention)
+
+from .cache_utils import TeaCache
+from ..dist.wan_xfuser import usp_attn_forward
+
+try:
+ import flash_attn_interface
+ FLASH_ATTN_3_AVAILABLE = True
+except ModuleNotFoundError:
+ FLASH_ATTN_3_AVAILABLE = False
+
+try:
+ import flash_attn
+ FLASH_ATTN_2_AVAILABLE = True
+except ModuleNotFoundError:
+ FLASH_ATTN_2_AVAILABLE = False
+
+# if FLASH_ATTN_2_AVAILABLE is False and FLASH_ATTN_3_AVAILABLE is False:
+# print(1/0)
+
+def flash_attention(
+ q,
+ k,
+ v,
+ q_lens=None,
+ k_lens=None,
+ dropout_p=0.,
+ softmax_scale=None,
+ q_scale=None,
+ causal=False,
+ window_size=(-1, -1),
+ deterministic=False,
+ dtype=torch.bfloat16,
+ version=None,
+):
+ """
+ q: [B, Lq, Nq, C1].
+ k: [B, Lk, Nk, C1].
+ v: [B, Lk, Nk, C2]. Nq must be divisible by Nk.
+ q_lens: [B].
+ k_lens: [B].
+ dropout_p: float. Dropout probability.
+ softmax_scale: float. The scaling of QK^T before applying softmax.
+ causal: bool. Whether to apply causal attention mask.
+ window_size: (left right). If not (-1, -1), apply sliding window local attention.
+ deterministic: bool. If True, slightly slower and uses more memory.
+ dtype: torch.dtype. Apply when dtype of q/k/v is not float16/bfloat16.
+ """
+ half_dtypes = (torch.float16, torch.bfloat16)
+ assert dtype in half_dtypes
+ assert q.device.type == 'cuda' and q.size(-1) <= 256
+
+ # params
+ b, lq, lk, out_dtype = q.size(0), q.size(1), k.size(1), q.dtype
+
+ def half(x):
+ return x if x.dtype in half_dtypes else x.to(dtype)
+
+ # preprocess query
+ if q_lens is None:
+ q = half(q.flatten(0, 1))
+ q_lens = torch.tensor(
+ [lq] * b, dtype=torch.int32).to(
+ device=q.device, non_blocking=True)
+ else:
+ q = half(torch.cat([u[:v] for u, v in zip(q, q_lens)]))
+
+ # preprocess key, value
+ if k_lens is None:
+ k = half(k.flatten(0, 1))
+ v = half(v.flatten(0, 1))
+ k_lens = torch.tensor(
+ [lk] * b, dtype=torch.int32).to(
+ device=k.device, non_blocking=True)
+ else:
+ k = half(torch.cat([u[:v] for u, v in zip(k, k_lens)]))
+ v = half(torch.cat([u[:v] for u, v in zip(v, k_lens)]))
+
+ q = q.to(v.dtype)
+ k = k.to(v.dtype)
+
+ if q_scale is not None:
+ q = q * q_scale
+
+ if version is not None and version == 3 and not FLASH_ATTN_3_AVAILABLE:
+ warnings.warn(
+ 'Flash attention 3 is not available, use flash attention 2 instead.'
+ )
+
+ # apply attention
+ if (version is None or version == 3) and FLASH_ATTN_3_AVAILABLE:
+ # Note: dropout_p, window_size are not supported in FA3 now.
+ x = flash_attn_interface.flash_attn_varlen_func(
+ q=q,
+ k=k,
+ v=v,
+ cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ seqused_q=None,
+ seqused_k=None,
+ max_seqlen_q=lq,
+ max_seqlen_k=lk,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ deterministic=deterministic)[0].unflatten(0, (b, lq))
+ else:
+ assert FLASH_ATTN_2_AVAILABLE
+ x = flash_attn.flash_attn_varlen_func(
+ q=q,
+ k=k,
+ v=v,
+ cu_seqlens_q=torch.cat([q_lens.new_zeros([1]), q_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ cu_seqlens_k=torch.cat([k_lens.new_zeros([1]), k_lens]).cumsum(
+ 0, dtype=torch.int32).to(q.device, non_blocking=True),
+ max_seqlen_q=lq,
+ max_seqlen_k=lk,
+ dropout_p=dropout_p,
+ softmax_scale=softmax_scale,
+ causal=causal,
+ window_size=window_size,
+ deterministic=deterministic).unflatten(0, (b, lq))
+
+ # output
+ return x.type(out_dtype)
+
+
+def attention(
+ q,
+ k,
+ v,
+ q_lens=None,
+ k_lens=None,
+ dropout_p=0.,
+ softmax_scale=None,
+ q_scale=None,
+ causal=False,
+ window_size=(-1, -1),
+ deterministic=False,
+ dtype=torch.bfloat16,
+ fa_version=None,
+):
+ if FLASH_ATTN_2_AVAILABLE or FLASH_ATTN_3_AVAILABLE:
+ return flash_attention(
+ q=q,
+ k=k,
+ v=v,
+ q_lens=q_lens,
+ k_lens=k_lens,
+ dropout_p=dropout_p,
+ softmax_scale=softmax_scale,
+ q_scale=q_scale,
+ causal=causal,
+ window_size=window_size,
+ deterministic=deterministic,
+ dtype=dtype,
+ version=fa_version,
+ )
+ else:
+ if q_lens is not None or k_lens is not None:
+ warnings.warn(
+ 'Padding mask is disabled when using scaled_dot_product_attention. It can have a significant impact on performance.'
+ )
+ attn_mask = None
+
+ q = q.transpose(1, 2)
+ k = k.transpose(1, 2)
+ v = v.transpose(1, 2)
+
+ if torch.backends.cuda.flash_sdp_enabled() is False or torch.backends.cuda.enable_flash_sdp is False:
+ print(1/0)
+
+ out = torch.nn.functional.scaled_dot_product_attention(
+ q, k, v, attn_mask=attn_mask, is_causal=causal, dropout_p=dropout_p)
+
+ out = out.transpose(1, 2).contiguous()
+ return out
+
+
+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('cuda', 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
+
+# modified from https://github.com/thu-ml/RIFLEx/blob/main/riflex_utils.py
+@amp.autocast('cuda', enabled=False)
+def get_1d_rotary_pos_embed_riflex(
+ pos: Union[np.ndarray, int],
+ dim: int,
+ theta: float = 10000.0,
+ use_real=False,
+ k: Optional[int] = None,
+ L_test: Optional[int] = None,
+ L_test_scale: Optional[int] = None,
+):
+ """
+ RIFLEx: Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
+
+ This function calculates a frequency tensor with complex exponentials using the given dimension 'dim' and the end
+ index 'end'. The 'theta' parameter scales the frequencies. The returned tensor contains complex values in complex64
+ data type.
+
+ Args:
+ dim (`int`): Dimension of the frequency tensor.
+ pos (`np.ndarray` or `int`): Position indices for the frequency tensor. [S] or scalar
+ theta (`float`, *optional*, defaults to 10000.0):
+ Scaling factor for frequency computation. Defaults to 10000.0.
+ use_real (`bool`, *optional*):
+ If True, return real part and imaginary part separately. Otherwise, return complex numbers.
+ k (`int`, *optional*, defaults to None): the index for the intrinsic frequency in RoPE
+ L_test (`int`, *optional*, defaults to None): the number of frames for inference
+ Returns:
+ `torch.Tensor`: Precomputed frequency tensor with complex exponentials. [S, D/2]
+ """
+ assert dim % 2 == 0
+
+ if isinstance(pos, int):
+ pos = torch.arange(pos)
+ if isinstance(pos, np.ndarray):
+ pos = torch.from_numpy(pos) # type: ignore # [S]
+
+ freqs = 1.0 / torch.pow(theta,
+ torch.arange(0, dim, 2).to(torch.float64).div(dim))
+
+ # === Riflex modification start ===
+ # Reduce the intrinsic frequency to stay within a single period after extrapolation (see Eq. (8)).
+ # Empirical observations show that a few videos may exhibit repetition in the tail frames.
+ # To be conservative, we multiply by 0.9 to keep the extrapolated length below 90% of a single period.
+ if k is not None:
+ freqs[k-1] = 0.9 * 2 * torch.pi / L_test
+ # === Riflex modification end ===
+ if L_test_scale is not None:
+ freqs[k-1] = freqs[k-1] / L_test_scale
+
+ freqs = torch.outer(pos, freqs) # type: ignore # [S, D/2]
+ if use_real:
+ freqs_cos = freqs.cos().repeat_interleave(2, dim=1).float() # [S, D]
+ freqs_sin = freqs.sin().repeat_interleave(2, dim=1).float() # [S, D]
+ return freqs_cos, freqs_sin
+ else:
+ # lumina
+ freqs_cis = torch.polar(torch.ones_like(freqs), freqs) # complex64 # [S, D/2]
+ return freqs_cis
+
+@amp.autocast('cuda', enabled=False)
+def rope_apply(x, grid_sizes, freqs):
+ n, c = x.size(2), x.size(3) // 2
+
+ # split freqs
+ freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)
+
+ # loop over samples
+ output = []
+ for i, (f, h, w) in enumerate(grid_sizes.tolist()):
+ seq_len = f * h * w
+
+ # precompute multipliers
+ x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float32).reshape(
+ seq_len, n, -1, 2))
+ freqs_i = torch.cat([
+ freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
+ freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
+ freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
+ ],
+ dim=-1).reshape(seq_len, 1, -1)
+
+ # apply rotary embedding
+ x_i = torch.view_as_real(x_i * freqs_i).flatten(2)
+ x_i = torch.cat([x_i, x[i, seq_len:]])
+
+ # append to collection
+ output.append(x_i)
+ return torch.stack(output).float()
+
+
+class WanRMSNorm(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):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ """
+ 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 WanLayerNorm(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):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ """
+ return super().forward(x.float()).type_as(x)
+
+
+class WanSelfAttention(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 = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+ self.norm_k = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+ def forward(self, x, seq_lens, grid_sizes, freqs, dtype):
+ r"""
+ 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.to(dtype))).view(b, s, n, d)
+ k = self.norm_k(self.k(x.to(dtype))).view(b, s, n, d)
+ v = self.v(x.to(dtype)).view(b, s, n, d)
+ return q, k, v
+
+ q, k, v = qkv_fn(x)
+
+ x = attention(
+ q=rope_apply(q, grid_sizes, freqs).to(dtype),
+ k=rope_apply(k, grid_sizes, freqs).to(dtype),
+ v=v.to(dtype),
+ k_lens=seq_lens,
+ window_size=self.window_size)
+ x = x.to(dtype)
+
+ # output
+ x = x.flatten(2)
+ x = self.o(x)
+ return x
+
+
+class WanT2VCrossAttention(WanSelfAttention):
+
+ def forward(self, x, context, context_lens, dtype):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ context(Tensor): Shape [B, L2, C]
+ context_lens(Tensor): Shape [B]
+ """
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.norm_q(self.q(x.to(dtype))).view(b, -1, n, d)
+ k = self.norm_k(self.k(context.to(dtype))).view(b, -1, n, d)
+ v = self.v(context.to(dtype)).view(b, -1, n, d)
+
+ # compute attention
+ x = attention(
+ q.to(dtype),
+ k.to(dtype),
+ v.to(dtype),
+ k_lens=context_lens
+ )
+ x = x.to(dtype)
+
+ # output
+ x = x.flatten(2)
+ x = self.o(x)
+ return x
+
+
+class WanI2VCrossAttention(WanSelfAttention):
+
+ def __init__(self,
+ dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ eps=1e-6):
+ super().__init__(dim, num_heads, window_size, qk_norm, eps)
+
+ # self.alpha = nn.Parameter(torch.zeros((1, )))
+ self.k_img = nn.Linear(dim, dim)
+ self.v_img = nn.Linear(dim, dim)
+ self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+ def forward(self, x, context, context_lens, dtype):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ context(Tensor): Shape [B, L2, C]
+ context_lens(Tensor): Shape [B]
+ """
+ context_img = context[:, :257]
+ context = context[:, 257:]
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.norm_q(self.q(x.to(dtype))).view(b, -1, n, d)
+ k = self.norm_k(self.k(context.to(dtype))).view(b, -1, n, d)
+ v = self.v(context.to(dtype)).view(b, -1, n, d)
+ k_img = self.norm_k_img(self.k_img(context_img.to(dtype))).view(b, -1, n, d)
+ v_img = self.v_img(context_img.to(dtype)).view(b, -1, n, d)
+
+ img_x = attention(
+ q.to(dtype),
+ k_img.to(dtype),
+ v_img.to(dtype),
+ k_lens=None
+ )
+ img_x = img_x.to(dtype)
+ # compute attention
+ x = attention(
+ q.to(dtype),
+ k.to(dtype),
+ v.to(dtype),
+ k_lens=context_lens
+ )
+ x = x.to(dtype)
+
+ # output
+ x = x.flatten(2)
+ img_x = img_x.flatten(2)
+ x = x + img_x
+ x = self.o(x)
+ return x
+
+
+class WanI2VTalkingCrossAttention(WanSelfAttention):
+ def __init__(self,
+ dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ eps=1e-6):
+ super().__init__(dim, num_heads, window_size, qk_norm, eps)
+
+ # self.alpha = nn.Parameter(torch.zeros((1, )))
+ self.k_img = nn.Linear(dim, dim)
+ self.v_img = nn.Linear(dim, dim)
+ self.norm_k_img = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+ self.k_vocal = nn.Linear(dim, dim)
+ self.v_vocal = nn.Linear(dim, dim)
+ self.norm_k_vocal = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+ self.k_audio = nn.Linear(dim, dim)
+ self.v_audio = nn.Linear(dim, dim)
+ self.norm_k_audio = WanRMSNorm(dim, eps=eps) if qk_norm else nn.Identity()
+
+
+ def forward(self, x, context, context_lens, dtype):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ context(Tensor): Shape [B, L2, C]
+ context_lens(Tensor): Shape [B]
+ """
+
+
+ # print(f"The size of x in cross attention: {x.size()}") # [1, 21504, 5120]
+ # print(f"The size of context_clip in cross attention: {context_img.size()}") # [1, 257, 5120]
+ context_img = context[:, :257]
+ context_vocal = context[:, 257:(257+21*32)]
+ context_audio = context[:, (257+21*32):(257+21*32*2)]
+ context = context[:, (257+21*32*2):]
+
+ b, n, d = x.size(0), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.norm_q(self.q(x.to(dtype))).view(b, -1, n, d)
+ k = self.norm_k(self.k(context.to(dtype))).view(b, -1, n, d)
+ v = self.v(context.to(dtype)).view(b, -1, n, d)
+ k_img = self.norm_k_img(self.k_img(context_img.to(dtype))).view(b, -1, n, d)
+ v_img = self.v_img(context_img.to(dtype)).view(b, -1, n, d)
+
+ img_x = attention(
+ q.to(dtype),
+ k_img.to(dtype),
+ v_img.to(dtype),
+ k_lens=None
+ )
+ img_x = img_x.to(dtype)
+ # compute attention
+ x = attention(
+ q.to(dtype),
+ k.to(dtype),
+ v.to(dtype),
+ k_lens=context_lens
+ )
+ x = x.to(dtype)
+
+ k_vocal = self.norm_k_vocal(self.k_vocal(context_vocal.to(dtype))).view(b, -1, n, d)
+ v_vocal = self.v_vocal(context_vocal.to(dtype)).view(b, -1, n, d)
+ vocal_x = attention(
+ q.to(dtype),
+ k_vocal.to(dtype),
+ v_vocal.to(dtype),
+ k_lens=context_lens
+ )
+ vocal_x = vocal_x.to(dtype)
+
+ compressed_frame_number = 21
+ q = rearrange(q, "b (f n) h d -> (b f) n h d", f=compressed_frame_number)
+ context_audio = rearrange(context_audio, "b (f n) d -> (b f) n d", f=compressed_frame_number)
+ k_audio = self.norm_k_audio(self.k_audio(context_audio.to(dtype))).view(b*compressed_frame_number, -1, self.num_heads, self.head_dim)
+ v_audio = self.v_audio(context_audio.to(dtype)).view(b*compressed_frame_number, -1, self.num_heads, self.head_dim)
+
+ # k_audio = self.norm_k_audio(self.k_audio(context_audio.to(dtype))).view(b, -1, n, d)
+ # v_audio = self.v_audio(context_audio.to(dtype)).view(b, -1, n, d)
+ audio_x = attention(
+ q.to(dtype),
+ k_audio.to(dtype),
+ v_audio.to(dtype),
+ k_lens=context_lens
+ )
+ audio_x = audio_x.to(dtype)
+
+ audio_x = rearrange(audio_x, "(b f) n h d -> b (f n) h d", f=compressed_frame_number)
+
+ # print(f"The size of context: {context.size()}") # [1, 512, 5120]
+ # print(f"The size of context_vocal: {context_vocal.size()}") # [1, 672, 5120]
+ # print(f"The size of context_audio: {context_audio.size()}") # [1, 672, 5120]
+ # print(f"The size of vocal_x: {vocal_x.size()}") # [1, 21504, 40, 128]
+ # print(f"The size of audio_x: {audio_x.size()}") # [1, 21504, 40, 128]
+ # print(f"The size of face_masks: {face_masks.size()}") # [1, 21504, 5120]
+ # print(f"The size of negative_face_masks: {negative_face_masks.size()}") # [1, 21504, 5120]
+
+ # output
+ x = x.flatten(2)
+ img_x = img_x.flatten(2)
+
+ vocal_x = vocal_x.flatten(2)
+ audio_x = audio_x.flatten(2)
+
+ x = x + img_x + vocal_x + audio_x
+ x = self.o(x)
+ return x
+
+
+WAN_CROSSATTENTION_CLASSES = {
+ 't2v_cross_attn': WanT2VCrossAttention,
+ 'i2v_cross_attn': WanI2VCrossAttention,
+}
+
+
+class WanAttentionBlock(nn.Module):
+
+ def __init__(self,
+ cross_attn_type,
+ dim,
+ ffn_dim,
+ num_heads,
+ window_size=(-1, -1),
+ qk_norm=True,
+ cross_attn_norm=False,
+ eps=1e-6):
+ 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 = WanLayerNorm(dim, eps)
+ self.self_attn = WanSelfAttention(dim, num_heads, window_size, qk_norm,
+ eps)
+ self.norm3 = WanLayerNorm(
+ dim, eps,
+ elementwise_affine=True) if cross_attn_norm else nn.Identity()
+ # self.cross_attn = WAN_CROSSATTENTION_CLASSES[cross_attn_type](dim, num_heads, (-1, -1), qk_norm, eps)
+ self.cross_attn = WanI2VTalkingCrossAttention(dim, num_heads, (-1, -1), qk_norm, eps)
+ self.norm2 = WanLayerNorm(dim, eps)
+ self.ffn = nn.Sequential(
+ nn.Linear(dim, ffn_dim), nn.GELU(approximate='tanh'),
+ nn.Linear(ffn_dim, dim))
+
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5)
+
+ def forward(
+ self,
+ x,
+ e,
+ seq_lens,
+ grid_sizes,
+ freqs,
+ context,
+ context_lens,
+ dtype=torch.float32,
+ ):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L, C]
+ e(Tensor): Shape [B, 6, C]
+ seq_lens(Tensor): Shape [B], length of each sequence in batch
+ grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
+ freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
+ """
+ e = (self.modulation + e).chunk(6, dim=1)
+
+ # self-attention
+ temp_x = self.norm1(x) * (1 + e[1]) + e[0]
+ temp_x = temp_x.to(dtype)
+
+ y = self.self_attn(temp_x, seq_lens, grid_sizes, freqs, dtype)
+ x = x + y * e[2]
+
+ # cross-attention & ffn function
+ def cross_attn_ffn(x, context, context_lens, e):
+ # cross-attention
+ x = x + self.cross_attn(self.norm3(x), context, context_lens, dtype)
+
+ # ffn function
+ temp_x = self.norm2(x) * (1 + e[4]) + e[3]
+ temp_x = temp_x.to(dtype)
+
+ y = self.ffn(temp_x)
+ x = x + y * e[5]
+ return x
+
+ x = cross_attn_ffn(x, context, context_lens, e)
+ 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 = WanLayerNorm(dim, eps)
+ self.head = nn.Linear(dim, out_dim)
+
+ # modulation
+ self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5)
+
+ def forward(self, x, e):
+ r"""
+ Args:
+ x(Tensor): Shape [B, L1, C]
+ e(Tensor): Shape [B, C]
+ """
+ e = (self.modulation + e.unsqueeze(1)).chunk(2, dim=1)
+ x = (self.head(self.norm(x) * (1 + e[1]) + e[0]))
+ return x
+
+
+class MLPProj(torch.nn.Module):
+
+ def __init__(self, in_dim, out_dim):
+ super().__init__()
+
+ self.proj = torch.nn.Sequential(
+ torch.nn.LayerNorm(in_dim), torch.nn.Linear(in_dim, in_dim),
+ torch.nn.GELU(), torch.nn.Linear(in_dim, out_dim),
+ torch.nn.LayerNorm(out_dim))
+
+ def forward(self, image_embeds):
+ clip_extra_context_tokens = self.proj(image_embeds)
+ return clip_extra_context_tokens
+
+
+
+class WanTransformer3DModel(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+ r"""
+ Wan diffusion backbone supporting both text-to-video and image-to-video.
+ """
+
+ # ignore_for_config = [
+ # 'patch_size', 'cross_attn_norm', 'qk_norm', 'text_dim', 'window_size'
+ # ]
+ # _no_split_modules = ['WanAttentionBlock']
+ _supports_gradient_checkpointing = True
+
+ @register_to_config
+ def __init__(
+ self,
+ model_type='t2v',
+ 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,
+ in_channels=16,
+ hidden_size=2048,
+ ):
+ r"""
+ Initialize the diffusion model backbone.
+
+ Args:
+ model_type (`str`, *optional*, defaults to 't2v'):
+ Model variant - 't2v' (text-to-video) or 'i2v' (image-to-video)
+ patch_size (`tuple`, *optional*, defaults to (1, 2, 2)):
+ 3D patch dimensions for video embedding (t_patch, h_patch, w_patch)
+ text_len (`int`, *optional*, defaults to 512):
+ Fixed length for text embeddings
+ in_dim (`int`, *optional*, defaults to 16):
+ Input video channels (C_in)
+ dim (`int`, *optional*, defaults to 2048):
+ Hidden dimension of the transformer
+ ffn_dim (`int`, *optional*, defaults to 8192):
+ Intermediate dimension in feed-forward network
+ freq_dim (`int`, *optional*, defaults to 256):
+ Dimension for sinusoidal time embeddings
+ text_dim (`int`, *optional*, defaults to 4096):
+ Input dimension for text embeddings
+ out_dim (`int`, *optional*, defaults to 16):
+ Output video channels (C_out)
+ num_heads (`int`, *optional*, defaults to 16):
+ Number of attention heads
+ num_layers (`int`, *optional*, defaults to 32):
+ Number of transformer blocks
+ window_size (`tuple`, *optional*, defaults to (-1, -1)):
+ Window size for local attention (-1 indicates global attention)
+ qk_norm (`bool`, *optional*, defaults to True):
+ Enable query/key normalization
+ cross_attn_norm (`bool`, *optional*, defaults to False):
+ Enable cross-attention normalization
+ eps (`float`, *optional*, defaults to 1e-6):
+ Epsilon value for normalization layers
+ """
+
+ super().__init__()
+
+ assert model_type in ['t2v', 'i2v']
+ 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
+ cross_attn_type = 't2v_cross_attn' if model_type == 't2v' else 'i2v_cross_attn'
+ self.blocks = nn.ModuleList([
+ WanAttentionBlock(cross_attn_type, dim, ffn_dim, num_heads,
+ window_size, qk_norm, cross_attn_norm, eps)
+ for _ in range(num_layers)
+ ])
+
+ # head
+ self.head = Head(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.d = d
+ 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
+ )
+
+ if model_type == 'i2v':
+ self.img_emb = MLPProj(1280, dim)
+
+ self.teacache = None
+ self.gradient_checkpointing = False
+ self.sp_world_size = 1
+ self.sp_world_rank = 0
+
+ def enable_teacache(
+ self,
+ coefficients,
+ num_steps: int,
+ rel_l1_thresh: float,
+ num_skip_start_steps: int = 0,
+ offload: bool = True
+ ):
+ self.teacache = TeaCache(
+ coefficients, num_steps, rel_l1_thresh=rel_l1_thresh, num_skip_start_steps=num_skip_start_steps, offload=offload
+ )
+
+ def disable_teacache(self):
+ self.teacache = None
+
+ def enable_riflex(
+ self,
+ k = 6,
+ L_test = 66,
+ L_test_scale = 4.886,
+ ):
+ device = self.freqs.device
+ self.freqs = torch.cat(
+ [
+ get_1d_rotary_pos_embed_riflex(1024, self.d - 4 * (self.d // 6), use_real=False, k=k, L_test=L_test, L_test_scale=L_test_scale),
+ rope_params(1024, 2 * (self.d // 6)),
+ rope_params(1024, 2 * (self.d // 6))
+ ],
+ dim=1
+ ).to(device)
+
+ def disable_riflex(self):
+ device = self.freqs.device
+ self.freqs = torch.cat(
+ [
+ rope_params(1024, self.d - 4 * (self.d // 6)),
+ rope_params(1024, 2 * (self.d // 6)),
+ rope_params(1024, 2 * (self.d // 6))
+ ],
+ dim=1
+ ).to(device)
+
+ def enable_multi_gpus_inference(self,):
+ self.sp_world_size = get_sequence_parallel_world_size()
+ self.sp_world_rank = get_sequence_parallel_rank()
+ for block in self.blocks:
+ block.self_attn.forward = types.MethodType(
+ usp_attn_forward, block.self_attn)
+
+ def _set_gradient_checkpointing(self, module, value=False):
+ self.gradient_checkpointing = value
+
+ def forward(
+ self,
+ x,
+ t,
+ context,
+ seq_len,
+ clip_fea=None,
+ y=None,
+ cond_flag=True,
+ vocal_context=None,
+ audio_context=None,
+ motion_bucket=None,
+ motion_embeddings=None,
+ ):
+ r"""
+ Forward pass through the diffusion model
+
+ Args:
+ x (List[Tensor]):
+ List of input video tensors, each with shape [C_in, F, H, W]
+ t (Tensor):
+ Diffusion timesteps tensor of shape [B]
+ context (List[Tensor]):
+ List of text embeddings each with shape [L, C]
+ seq_len (`int`):
+ Maximum sequence length for positional encoding
+ clip_fea (Tensor, *optional*):
+ CLIP image features for image-to-video mode
+ y (List[Tensor], *optional*):
+ Conditional video inputs for image-to-video mode, same shape as x
+ cond_flag (`bool`, *optional*, defaults to True):
+ Flag to indicate whether to forward the condition input
+
+ Returns:
+ List[Tensor]:
+ List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]
+ """
+ if self.model_type == 'i2v':
+ assert clip_fea is not None and y is not None
+ # params
+ device = self.patch_embedding.weight.device
+ dtype = x.dtype
+ if self.freqs.device != device and torch.device(type="meta") != device:
+ self.freqs = self.freqs.to(device)
+
+ if y is not None:
+ x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]
+
+ # embeddings
+ x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
+ 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)
+ if self.sp_world_size > 1:
+ seq_len = int(math.ceil(seq_len / self.sp_world_size)) * self.sp_world_size
+ assert seq_lens.max() <= seq_len
+ x = torch.cat([torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in x])
+
+ # face_masks = [torch.cat([u, v], dim=0) for u, v in zip(face_masks, face_masks_y)]
+ # face_masks = [self.patch_embedding(u.unsqueeze(0)) for u in face_masks]
+ # face_masks = [u.flatten(2).transpose(1, 2) for u in face_masks]
+ # face_masks_seq_lens = torch.tensor([u.size(1) for u in face_masks], dtype=torch.long)
+ # if self.sp_world_size > 1:
+ # seq_len = int(math.ceil(seq_len / self.sp_world_size)) * self.sp_world_size
+ # assert face_masks_seq_lens.max() <= seq_len
+ # face_masks = torch.cat([torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in face_masks])
+ # negative_face_masks = [torch.cat([u, v], dim=0) for u, v in zip(negative_face_masks, negative_face_masks_y)]
+ # negative_face_masks = [self.patch_embedding(u.unsqueeze(0)) for u in negative_face_masks]
+ # negative_face_masks = [u.flatten(2).transpose(1, 2) for u in negative_face_masks]
+ # negative_face_masks_seq_lens = torch.tensor([u.size(1) for u in negative_face_masks], dtype=torch.long)
+ # if self.sp_world_size > 1:
+ # seq_len = int(math.ceil(seq_len / self.sp_world_size)) * self.sp_world_size
+ # assert negative_face_masks_seq_lens.max() <= seq_len
+ # negative_face_masks = torch.cat([torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1) for u in negative_face_masks])
+
+ # time embeddings
+ with amp.autocast('cuda', 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))
+ if motion_embeddings is not None:
+ e0 = e0 + motion_embeddings.unflatten(1, (6, self.dim))
+ # to bfloat16 for saving memeory
+ # assert e.dtype == torch.float32 and e0.dtype == torch.float32
+ e0 = e0.to(dtype)
+ e = e.to(dtype)
+
+ # 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
+ ]))
+ context_clip = self.img_emb(clip_fea) # bs x 257 x dim
+
+ # print("-----------------------------------------")
+ # print(f"motion scale: {motion_scale}") # 0.2
+ # print(f"The size of face_masks: {face_masks.size()}") # [1, 81, 1, 512, 512]
+ # print(f"The size of context: {context.size()}") # [1, 512, 5120]
+ # print(f"The size of context_clip: {context_clip.size()}") # [1, 257, 5120]
+ # print(f"The size of e: {e.size()}") # [1, 5120]
+ # print(f"The size of e0: {e0.size()}") # [1, 6, 5120]
+ # print(f"The size of vocal_context: {vocal_context.size()}") # [1, 21, 32, 5120]
+ # print(f"The size of audio_context: {audio_context.size()}") # [1, 21, 32, 5120]
+ # print("-----------------------------------------")
+
+ if vocal_context is None:
+ dim = context_clip.size()[-1]
+ batch_size = context_clip.size()[0]
+ vocal_context = torch.zeros(batch_size, 21, 32, dim).to(device).to(dtype)
+ if audio_context is None:
+ dim = context_clip.size()[-1]
+ batch_size = context_clip.size()[0]
+ audio_context = torch.zeros(batch_size, 21, 32, dim).to(device).to(dtype)
+
+ vocal_context = rearrange(vocal_context, "b f n d -> b (f n) d")
+ audio_context = rearrange(audio_context, "b f n d -> b (f n) d")
+
+ context = torch.concat([context_clip, vocal_context, audio_context, context], dim=1)
+
+ # Context Parallel
+ if self.sp_world_size > 1:
+ x = torch.chunk(x, self.sp_world_size, dim=1)[self.sp_world_rank]
+
+ # TeaCache
+ if self.teacache is not None:
+ if cond_flag:
+ modulated_inp = e0
+ skip_flag = self.teacache.cnt < self.teacache.num_skip_start_steps
+ if self.teacache.cnt == 0 or self.teacache.cnt == self.teacache.num_steps - 1 or skip_flag:
+ should_calc = True
+ self.teacache.accumulated_rel_l1_distance = 0
+ else:
+ if cond_flag:
+ rel_l1_distance = self.teacache.compute_rel_l1_distance(self.teacache.previous_modulated_input, modulated_inp)
+ self.teacache.accumulated_rel_l1_distance += self.teacache.rescale_func(rel_l1_distance)
+ if self.teacache.accumulated_rel_l1_distance < self.teacache.rel_l1_thresh:
+ should_calc = False
+ else:
+ should_calc = True
+ self.teacache.accumulated_rel_l1_distance = 0
+ self.teacache.previous_modulated_input = modulated_inp
+ self.teacache.cnt += 1
+ if self.teacache.cnt == self.teacache.num_steps:
+ self.teacache.reset()
+ self.teacache.should_calc = should_calc
+ else:
+ should_calc = self.teacache.should_calc
+
+ # TeaCache
+ if self.teacache is not None:
+ if not should_calc:
+ previous_residual = self.teacache.previous_residual_cond if cond_flag else self.teacache.previous_residual_uncond
+ x = x + previous_residual.to(x.device)
+ else:
+ ori_x = x.clone().cpu() if self.teacache.offload else x.clone()
+
+ for block in self.blocks:
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+ x = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ x,
+ e0,
+ seq_lens,
+ grid_sizes,
+ self.freqs,
+ context,
+ context_lens,
+ dtype,
+ **ckpt_kwargs,
+ )
+ else:
+ # arguments
+ kwargs = dict(
+ e=e0,
+ seq_lens=seq_lens,
+ grid_sizes=grid_sizes,
+ freqs=self.freqs,
+ context=context,
+ context_lens=context_lens,
+ dtype=dtype,
+ )
+ x = block(x, **kwargs)
+
+ if cond_flag:
+ self.teacache.previous_residual_cond = x.cpu() - ori_x if self.teacache.offload else x - ori_x
+ else:
+ self.teacache.previous_residual_uncond = x.cpu() - ori_x if self.teacache.offload else x - ori_x
+ else:
+ for block in self.blocks:
+ if torch.is_grad_enabled() and self.gradient_checkpointing:
+
+ def create_custom_forward(module):
+ def custom_forward(*inputs):
+ return module(*inputs)
+
+ return custom_forward
+ ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+ x = torch.utils.checkpoint.checkpoint(
+ create_custom_forward(block),
+ x,
+ e0,
+ seq_lens,
+ grid_sizes,
+ self.freqs,
+ context,
+ context_lens,
+ dtype,
+ **ckpt_kwargs,
+ )
+ else:
+ # arguments
+ kwargs = dict(
+ e=e0,
+ seq_lens=seq_lens,
+ grid_sizes=grid_sizes,
+ freqs=self.freqs,
+ context=context,
+ context_lens=context_lens,
+ dtype=dtype,
+ )
+ x = block(x, **kwargs)
+
+ if self.sp_world_size > 1:
+ x = get_sp_group().all_gather(x, dim=1)
+
+ # head
+ x = self.head(x, e)
+
+ # unpatchify
+ x = self.unpatchify(x, grid_sizes)
+ x = torch.stack(x)
+ return x
+
+
+ def unpatchify(self, x, grid_sizes):
+ r"""
+ 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)
+
+ @classmethod
+ def from_pretrained(
+ cls, pretrained_model_path, subfolder=None, transformer_additional_kwargs={},
+ low_cpu_mem_usage=False, torch_dtype=torch.bfloat16
+ ):
+ if subfolder is not None:
+ pretrained_model_path = os.path.join(pretrained_model_path, subfolder)
+ print(f"loaded 3D transformer's pretrained weights from {pretrained_model_path} ...")
+
+ config_file = os.path.join(pretrained_model_path, 'config.json')
+ if not os.path.isfile(config_file):
+ raise RuntimeError(f"{config_file} does not exist")
+ with open(config_file, "r") as f:
+ config = json.load(f)
+
+ from diffusers.utils import WEIGHTS_NAME
+ model_file = os.path.join(pretrained_model_path, WEIGHTS_NAME)
+ model_file_safetensors = model_file.replace(".bin", ".safetensors")
+
+ if "dict_mapping" in transformer_additional_kwargs.keys():
+ for key in transformer_additional_kwargs["dict_mapping"]:
+ transformer_additional_kwargs[transformer_additional_kwargs["dict_mapping"][key]] = config[key]
+
+ # {'patch_size', 'qk_norm', 'window_size', 'cross_attn_norm', 'text_dim'} was not found in config. Values will be initialized to default values.
+ transformer_additional_kwargs["patch_size"] = (1, 2, 2)
+ transformer_additional_kwargs["qk_norm"] = True
+ transformer_additional_kwargs["window_size"] = (-1, -1)
+ transformer_additional_kwargs["cross_attn_norm"] = True
+
+ if low_cpu_mem_usage:
+ try:
+ import re
+
+ from diffusers.models.modeling_utils import \
+ load_model_dict_into_meta
+ from diffusers.utils import is_accelerate_available
+ if is_accelerate_available():
+ import accelerate
+
+ # Instantiate model with empty weights
+ with accelerate.init_empty_weights():
+ model = cls.from_config(config, **transformer_additional_kwargs)
+
+ param_device = "cpu"
+ if os.path.exists(model_file):
+ state_dict = torch.load(model_file, map_location="cpu")
+ elif os.path.exists(model_file_safetensors):
+ from safetensors.torch import load_file, safe_open
+ state_dict = load_file(model_file_safetensors)
+ else:
+ from safetensors.torch import load_file, safe_open
+ model_files_safetensors = glob.glob(os.path.join(pretrained_model_path, "*.safetensors"))
+ state_dict = {}
+ print(model_files_safetensors)
+ for _model_file_safetensors in model_files_safetensors:
+ _state_dict = load_file(_model_file_safetensors)
+ for key in _state_dict:
+ state_dict[key] = _state_dict[key]
+ model._convert_deprecated_attention_blocks(state_dict)
+ # move the params from meta device to cpu
+ missing_keys = set(model.state_dict().keys()) - set(state_dict.keys())
+ if len(missing_keys) > 0:
+ raise ValueError(
+ f"Cannot load {cls} from {pretrained_model_path} because the following keys are"
+ f" missing: \n {', '.join(missing_keys)}. \n Please make sure to pass"
+ " `low_cpu_mem_usage=False` and `device_map=None` if you want to randomly initialize"
+ " those weights or else make sure your checkpoint file is correct."
+ )
+
+ unexpected_keys = load_model_dict_into_meta(
+ model,
+ state_dict,
+ device=param_device,
+ dtype=torch_dtype,
+ model_name_or_path=pretrained_model_path,
+ )
+
+ if cls._keys_to_ignore_on_load_unexpected is not None:
+ for pat in cls._keys_to_ignore_on_load_unexpected:
+ unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]
+
+ if len(unexpected_keys) > 0:
+ print(
+ f"Some weights of the model checkpoint were not used when initializing {cls.__name__}: \n {[', '.join(unexpected_keys)]}"
+ )
+ return model
+ except Exception as e:
+ print(
+ f"The low_cpu_mem_usage mode is not work because {e}. Use low_cpu_mem_usage=False instead."
+ )
+
+ model = cls.from_config(config, **transformer_additional_kwargs)
+ if os.path.exists(model_file):
+ state_dict = torch.load(model_file, map_location="cpu")
+ elif os.path.exists(model_file_safetensors):
+ from safetensors.torch import load_file, safe_open
+ state_dict = load_file(model_file_safetensors)
+ else:
+ from safetensors.torch import load_file, safe_open
+ model_files_safetensors = glob.glob(os.path.join(pretrained_model_path, "*.safetensors"))
+ state_dict = {}
+ for _model_file_safetensors in model_files_safetensors:
+ _state_dict = load_file(_model_file_safetensors)
+ for key in _state_dict:
+ state_dict[key] = _state_dict[key]
+
+ if model.state_dict()['patch_embedding.weight'].size() != state_dict['patch_embedding.weight'].size():
+ model.state_dict()['patch_embedding.weight'][:, :state_dict['patch_embedding.weight'].size()[1], :, :] = state_dict['patch_embedding.weight']
+ model.state_dict()['patch_embedding.weight'][:, state_dict['patch_embedding.weight'].size()[1]:, :, :] = 0
+ state_dict['patch_embedding.weight'] = model.state_dict()['patch_embedding.weight']
+
+
+ # txt_path = "/home/qid/v-shuyuantu/mycontainer/v-shuyuantu/MyTalking/wan_state_dict.txt"
+ # with open(txt_path, 'w') as f:
+ # for name in model.state_dict().keys():
+ # f.write(f"{name}\n")
+ # print("===============This is Transformer from pretrained function=======================")
+
+ tmp_state_dict = {}
+ for key in state_dict:
+ if key in model.state_dict().keys() and model.state_dict()[key].size() == state_dict[key].size():
+ tmp_state_dict[key] = state_dict[key]
+ if "cross_attn.k_img.weight" in key:
+ talking_vocal_key = key.replace("k_img", "k_vocal")
+ talking_audio_key = key.replace("k_img", "k_audio")
+ # tmp_state_dict[talking_vocal_key] = state_dict[key]
+ # tmp_state_dict[talking_audio_key] = state_dict[key]
+ tmp_state_dict[talking_vocal_key] = torch.zeros_like(state_dict[key])
+ tmp_state_dict[talking_audio_key] = torch.zeros_like(state_dict[key])
+ elif "cross_attn.k_img.bias" in key:
+ talking_vocal_key = key.replace("k_img", "k_vocal")
+ talking_audio_key = key.replace("k_img", "k_audio")
+ # tmp_state_dict[talking_vocal_key] = state_dict[key]
+ # tmp_state_dict[talking_audio_key] = state_dict[key]
+ tmp_state_dict[talking_vocal_key] = torch.zeros_like(state_dict[key])
+ tmp_state_dict[talking_audio_key] = torch.zeros_like(state_dict[key])
+ elif "cross_attn.v_img.weight" in key:
+ talking_vocal_key = key.replace("v_img", "v_vocal")
+ talking_audio_key = key.replace("v_img", "v_audio")
+ # tmp_state_dict[talking_vocal_key] = state_dict[key]
+ # tmp_state_dict[talking_audio_key] = state_dict[key]
+ tmp_state_dict[talking_vocal_key] = torch.zeros_like(state_dict[key])
+ tmp_state_dict[talking_audio_key] = torch.zeros_like(state_dict[key])
+ elif "cross_attn.v_img.bias" in key:
+ talking_vocal_key = key.replace("v_img", "v_vocal")
+ talking_audio_key = key.replace("v_img", "v_audio")
+ # tmp_state_dict[talking_vocal_key] = state_dict[key]
+ # tmp_state_dict[talking_audio_key] = state_dict[key]
+ tmp_state_dict[talking_vocal_key] = torch.zeros_like(state_dict[key])
+ tmp_state_dict[talking_audio_key] = torch.zeros_like(state_dict[key])
+ elif "cross_attn.norm_k_img.weight" in key:
+ talking_vocal_key = key.replace("norm_k_img", "norm_k_vocal")
+ talking_audio_key = key.replace("norm_k_img", "norm_k_audio")
+ # tmp_state_dict[talking_vocal_key] = state_dict[key]
+ # tmp_state_dict[talking_audio_key] = state_dict[key]
+ tmp_state_dict[talking_vocal_key] = torch.zeros_like(state_dict[key])
+ tmp_state_dict[talking_audio_key] = torch.zeros_like(state_dict[key])
+
+ else:
+ print(key, "Size don't match, skip")
+
+ state_dict = tmp_state_dict
+
+ m, u = model.load_state_dict(state_dict, strict=False)
+ print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
+ print(m)
+
+ params = [p.numel() if "." in n else 0 for n, p in model.named_parameters()]
+ print(f"### All Parameters: {sum(params) / 1e6} M")
+
+ params = [p.numel() if "attn1." in n else 0 for n, p in model.named_parameters()]
+ print(f"### attn1 Parameters: {sum(params) / 1e6} M")
+
+ model = model.to(torch_dtype)
+ return model
\ No newline at end of file
diff --git a/wan/models/wan_vae.py b/wan/models/wan_vae.py
new file mode 100644
index 0000000000000000000000000000000000000000..4afb122c65271e44a2f1947aed7bf7e4331fbc29
--- /dev/null
+++ b/wan/models/wan_vae.py
@@ -0,0 +1,705 @@
+# Modified from https://github.com/Wan-Video/Wan2.1/blob/main/wan/modules/vae.py
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from typing import 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.loaders.single_file_model import FromOriginalModelMixin
+from diffusers.models.autoencoders.vae import (DecoderOutput,
+ DiagonalGaussianDistribution)
+from diffusers.models.modeling_outputs import AutoencoderKLOutput
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.utils.accelerate_utils import apply_forward_hook
+from einops import rearrange
+
+CACHE_T = 2
+
+
+class CausalConv3d(nn.Conv3d):
+ """
+ Causal 3d convolusion.
+ """
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ self._padding = (self.padding[2], self.padding[2], self.padding[1],
+ self.padding[1], 2 * self.padding[0], 0)
+ self.padding = (0, 0, 0)
+
+ def forward(self, x, cache_x=None):
+ padding = list(self._padding)
+ if cache_x is not None and self._padding[4] > 0:
+ cache_x = cache_x.to(x.device)
+ x = torch.cat([cache_x, x], dim=2)
+ padding[4] -= cache_x.shape[2]
+ x = F.pad(x, padding)
+
+ return super().forward(x)
+
+
+class RMS_norm(nn.Module):
+
+ def __init__(self, dim, channel_first=True, images=True, bias=False):
+ super().__init__()
+ broadcastable_dims = (1, 1, 1) if not images else (1, 1)
+ shape = (dim, *broadcastable_dims) if channel_first else (dim,)
+
+ self.channel_first = channel_first
+ self.scale = dim**0.5
+ self.gamma = nn.Parameter(torch.ones(shape))
+ self.bias = nn.Parameter(torch.zeros(shape)) if bias else 0.
+
+ def forward(self, x):
+ return F.normalize(
+ x, dim=(1 if self.channel_first else
+ -1)) * self.scale * self.gamma + self.bias
+
+
+class Upsample(nn.Upsample):
+
+ def forward(self, x):
+ """
+ Fix bfloat16 support for nearest neighbor interpolation.
+ """
+ return super().forward(x.float()).type_as(x)
+
+
+class Resample(nn.Module):
+
+ def __init__(self, dim, mode):
+ assert mode in ('none', 'upsample2d', 'upsample3d', 'downsample2d',
+ 'downsample3d')
+ super().__init__()
+ self.dim = dim
+ self.mode = mode
+
+ # layers
+ if mode == 'upsample2d':
+ self.resample = nn.Sequential(
+ Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
+ nn.Conv2d(dim, dim // 2, 3, padding=1))
+ elif mode == 'upsample3d':
+ self.resample = nn.Sequential(
+ Upsample(scale_factor=(2., 2.), mode='nearest-exact'),
+ nn.Conv2d(dim, dim // 2, 3, padding=1))
+ self.time_conv = CausalConv3d(
+ dim, dim * 2, (3, 1, 1), padding=(1, 0, 0))
+
+ elif mode == 'downsample2d':
+ self.resample = nn.Sequential(
+ nn.ZeroPad2d((0, 1, 0, 1)),
+ nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+ elif mode == 'downsample3d':
+ self.resample = nn.Sequential(
+ nn.ZeroPad2d((0, 1, 0, 1)),
+ nn.Conv2d(dim, dim, 3, stride=(2, 2)))
+ self.time_conv = CausalConv3d(
+ dim, dim, (3, 1, 1), stride=(2, 1, 1), padding=(0, 0, 0))
+
+ else:
+ self.resample = nn.Identity()
+
+ def forward(self, x, feat_cache=None, feat_idx=[0]):
+ b, c, t, h, w = x.size()
+ if self.mode == 'upsample3d':
+ if feat_cache is not None:
+ idx = feat_idx[0]
+ if feat_cache[idx] is None:
+ feat_cache[idx] = 'Rep'
+ feat_idx[0] += 1
+ else:
+
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[
+ idx] is not None and feat_cache[idx] != 'Rep':
+ # cache last frame of last two chunk
+ cache_x = torch.cat([
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device), cache_x
+ ],
+ dim=2)
+ if cache_x.shape[2] < 2 and feat_cache[
+ idx] is not None and feat_cache[idx] == 'Rep':
+ cache_x = torch.cat([
+ torch.zeros_like(cache_x).to(cache_x.device),
+ cache_x
+ ],
+ dim=2)
+ if feat_cache[idx] == 'Rep':
+ x = self.time_conv(x)
+ else:
+ x = self.time_conv(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+
+ x = x.reshape(b, 2, c, t, h, w)
+ x = torch.stack((x[:, 0, :, :, :, :], x[:, 1, :, :, :, :]),
+ 3)
+ x = x.reshape(b, c, t * 2, h, w)
+ t = x.shape[2]
+ x = rearrange(x, 'b c t h w -> (b t) c h w')
+ x = self.resample(x)
+ x = rearrange(x, '(b t) c h w -> b c t h w', t=t)
+
+ if self.mode == 'downsample3d':
+ if feat_cache is not None:
+ idx = feat_idx[0]
+ if feat_cache[idx] is None:
+ feat_cache[idx] = x.clone()
+ feat_idx[0] += 1
+ else:
+
+ cache_x = x[:, :, -1:, :, :].clone()
+ # if cache_x.shape[2] < 2 and feat_cache[idx] is not None and feat_cache[idx]!='Rep':
+ # # cache last frame of last two chunk
+ # cache_x = torch.cat([feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(cache_x.device), cache_x], dim=2)
+
+ x = self.time_conv(
+ torch.cat([feat_cache[idx][:, :, -1:, :, :], x], 2))
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ return x
+
+ def init_weight(self, conv):
+ conv_weight = conv.weight
+ nn.init.zeros_(conv_weight)
+ c1, c2, t, h, w = conv_weight.size()
+ one_matrix = torch.eye(c1, c2)
+ init_matrix = one_matrix
+ nn.init.zeros_(conv_weight)
+ #conv_weight.data[:,:,-1,1,1] = init_matrix * 0.5
+ conv_weight.data[:, :, 1, 0, 0] = init_matrix #* 0.5
+ conv.weight.data.copy_(conv_weight)
+ nn.init.zeros_(conv.bias.data)
+
+ def init_weight2(self, conv):
+ conv_weight = conv.weight.data
+ nn.init.zeros_(conv_weight)
+ c1, c2, t, h, w = conv_weight.size()
+ init_matrix = torch.eye(c1 // 2, c2)
+ #init_matrix = repeat(init_matrix, 'o ... -> (o 2) ...').permute(1,0,2).contiguous().reshape(c1,c2)
+ conv_weight[:c1 // 2, :, -1, 0, 0] = init_matrix
+ conv_weight[c1 // 2:, :, -1, 0, 0] = init_matrix
+ conv.weight.data.copy_(conv_weight)
+ nn.init.zeros_(conv.bias.data)
+
+
+class ResidualBlock(nn.Module):
+
+ def __init__(self, in_dim, out_dim, dropout=0.0):
+ super().__init__()
+ self.in_dim = in_dim
+ self.out_dim = out_dim
+
+ # layers
+ self.residual = nn.Sequential(
+ RMS_norm(in_dim, images=False), nn.SiLU(),
+ CausalConv3d(in_dim, out_dim, 3, padding=1),
+ RMS_norm(out_dim, images=False), nn.SiLU(), nn.Dropout(dropout),
+ CausalConv3d(out_dim, out_dim, 3, padding=1))
+ self.shortcut = CausalConv3d(in_dim, out_dim, 1) \
+ if in_dim != out_dim else nn.Identity()
+
+ def forward(self, x, feat_cache=None, feat_idx=[0]):
+ h = self.shortcut(x)
+ for layer in self.residual:
+ if isinstance(layer, CausalConv3d) and feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ # cache last frame of last two chunk
+ cache_x = torch.cat([
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device), cache_x
+ ],
+ dim=2)
+ x = layer(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = layer(x)
+ return x + h
+
+
+class AttentionBlock(nn.Module):
+ """
+ Causal self-attention with a single head.
+ """
+
+ def __init__(self, dim):
+ super().__init__()
+ self.dim = dim
+
+ # layers
+ self.norm = RMS_norm(dim)
+ self.to_qkv = nn.Conv2d(dim, dim * 3, 1)
+ self.proj = nn.Conv2d(dim, dim, 1)
+
+ # zero out the last layer params
+ nn.init.zeros_(self.proj.weight)
+
+ def forward(self, x):
+ identity = x
+ b, c, t, h, w = x.size()
+ x = rearrange(x, 'b c t h w -> (b t) c h w')
+ x = self.norm(x)
+ # compute query, key, value
+ q, k, v = self.to_qkv(x).reshape(b * t, 1, c * 3,
+ -1).permute(0, 1, 3,
+ 2).contiguous().chunk(
+ 3, dim=-1)
+
+ # apply attention
+ x = F.scaled_dot_product_attention(
+ q,
+ k,
+ v,
+ )
+ x = x.squeeze(1).permute(0, 2, 1).reshape(b * t, c, h, w)
+
+ # output
+ x = self.proj(x)
+ x = rearrange(x, '(b t) c h w-> b c t h w', t=t)
+ return x + identity
+
+
+class Encoder3d(nn.Module):
+
+ def __init__(self,
+ dim=128,
+ z_dim=4,
+ dim_mult=[1, 2, 4, 4],
+ num_res_blocks=2,
+ attn_scales=[],
+ temperal_downsample=[True, True, False],
+ dropout=0.0):
+ super().__init__()
+ self.dim = dim
+ self.z_dim = z_dim
+ self.dim_mult = dim_mult
+ self.num_res_blocks = num_res_blocks
+ self.attn_scales = attn_scales
+ self.temperal_downsample = temperal_downsample
+
+ # dimensions
+ dims = [dim * u for u in [1] + dim_mult]
+ scale = 1.0
+
+ # init block
+ self.conv1 = CausalConv3d(3, dims[0], 3, padding=1)
+
+ # downsample blocks
+ downsamples = []
+ for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+ # residual (+attention) blocks
+ for _ in range(num_res_blocks):
+ downsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+ if scale in attn_scales:
+ downsamples.append(AttentionBlock(out_dim))
+ in_dim = out_dim
+
+ # downsample block
+ if i != len(dim_mult) - 1:
+ mode = 'downsample3d' if temperal_downsample[
+ i] else 'downsample2d'
+ downsamples.append(Resample(out_dim, mode=mode))
+ scale /= 2.0
+ self.downsamples = nn.Sequential(*downsamples)
+
+ # middle blocks
+ self.middle = nn.Sequential(
+ ResidualBlock(out_dim, out_dim, dropout), AttentionBlock(out_dim),
+ ResidualBlock(out_dim, out_dim, dropout))
+
+ # output blocks
+ self.head = nn.Sequential(
+ RMS_norm(out_dim, images=False), nn.SiLU(),
+ CausalConv3d(out_dim, z_dim, 3, padding=1))
+
+ def forward(self, x, feat_cache=None, feat_idx=[0]):
+ if feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ # cache last frame of last two chunk
+ cache_x = torch.cat([
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device), cache_x
+ ],
+ dim=2)
+ x = self.conv1(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = self.conv1(x)
+
+ ## downsamples
+ for layer in self.downsamples:
+ if feat_cache is not None:
+ x = layer(x, feat_cache, feat_idx)
+ else:
+ x = layer(x)
+
+ ## middle
+ for layer in self.middle:
+ if isinstance(layer, ResidualBlock) and feat_cache is not None:
+ x = layer(x, feat_cache, feat_idx)
+ else:
+ x = layer(x)
+
+ ## head
+ for layer in self.head:
+ if isinstance(layer, CausalConv3d) and feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ # cache last frame of last two chunk
+ cache_x = torch.cat([
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device), cache_x
+ ],
+ dim=2)
+ x = layer(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = layer(x)
+ return x
+
+
+class Decoder3d(nn.Module):
+
+ def __init__(self,
+ dim=128,
+ z_dim=4,
+ dim_mult=[1, 2, 4, 4],
+ num_res_blocks=2,
+ attn_scales=[],
+ temperal_upsample=[False, True, True],
+ dropout=0.0):
+ super().__init__()
+ self.dim = dim
+ self.z_dim = z_dim
+ self.dim_mult = dim_mult
+ self.num_res_blocks = num_res_blocks
+ self.attn_scales = attn_scales
+ self.temperal_upsample = temperal_upsample
+
+ # dimensions
+ dims = [dim * u for u in [dim_mult[-1]] + dim_mult[::-1]]
+ scale = 1.0 / 2**(len(dim_mult) - 2)
+
+ # init block
+ self.conv1 = CausalConv3d(z_dim, dims[0], 3, padding=1)
+
+ # middle blocks
+ self.middle = nn.Sequential(
+ ResidualBlock(dims[0], dims[0], dropout), AttentionBlock(dims[0]),
+ ResidualBlock(dims[0], dims[0], dropout))
+
+ # upsample blocks
+ upsamples = []
+ for i, (in_dim, out_dim) in enumerate(zip(dims[:-1], dims[1:])):
+ # residual (+attention) blocks
+ if i == 1 or i == 2 or i == 3:
+ in_dim = in_dim // 2
+ for _ in range(num_res_blocks + 1):
+ upsamples.append(ResidualBlock(in_dim, out_dim, dropout))
+ if scale in attn_scales:
+ upsamples.append(AttentionBlock(out_dim))
+ in_dim = out_dim
+
+ # upsample block
+ if i != len(dim_mult) - 1:
+ mode = 'upsample3d' if temperal_upsample[i] else 'upsample2d'
+ upsamples.append(Resample(out_dim, mode=mode))
+ scale *= 2.0
+ self.upsamples = nn.Sequential(*upsamples)
+
+ # output blocks
+ self.head = nn.Sequential(
+ RMS_norm(out_dim, images=False), nn.SiLU(),
+ CausalConv3d(out_dim, 3, 3, padding=1))
+
+ def forward(self, x, feat_cache=None, feat_idx=[0]):
+ ## conv1
+ if feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ # cache last frame of last two chunk
+ cache_x = torch.cat([
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device), cache_x
+ ],
+ dim=2)
+ x = self.conv1(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = self.conv1(x)
+
+ ## middle
+ for layer in self.middle:
+ if isinstance(layer, ResidualBlock) and feat_cache is not None:
+ x = layer(x, feat_cache, feat_idx)
+ else:
+ x = layer(x)
+
+ ## upsamples
+ for layer in self.upsamples:
+ if feat_cache is not None:
+ x = layer(x, feat_cache, feat_idx)
+ else:
+ x = layer(x)
+
+ ## head
+ for layer in self.head:
+ if isinstance(layer, CausalConv3d) and feat_cache is not None:
+ idx = feat_idx[0]
+ cache_x = x[:, :, -CACHE_T:, :, :].clone()
+ if cache_x.shape[2] < 2 and feat_cache[idx] is not None:
+ # cache last frame of last two chunk
+ cache_x = torch.cat([
+ feat_cache[idx][:, :, -1, :, :].unsqueeze(2).to(
+ cache_x.device), cache_x
+ ],
+ dim=2)
+ x = layer(x, feat_cache[idx])
+ feat_cache[idx] = cache_x
+ feat_idx[0] += 1
+ else:
+ x = layer(x)
+ return x
+
+
+def count_conv3d(model):
+ count = 0
+ for m in model.modules():
+ if isinstance(m, CausalConv3d):
+ count += 1
+ return count
+
+
+class AutoencoderKLWan_(nn.Module):
+
+ def __init__(self,
+ dim=128,
+ z_dim=4,
+ dim_mult=[1, 2, 4, 4],
+ num_res_blocks=2,
+ attn_scales=[],
+ temperal_downsample=[True, True, False],
+ dropout=0.0):
+ super().__init__()
+ self.dim = dim
+ self.z_dim = z_dim
+ self.dim_mult = dim_mult
+ self.num_res_blocks = num_res_blocks
+ self.attn_scales = attn_scales
+ self.temperal_downsample = temperal_downsample
+ self.temperal_upsample = temperal_downsample[::-1]
+
+ # modules
+ self.encoder = Encoder3d(dim, z_dim * 2, dim_mult, num_res_blocks,
+ attn_scales, self.temperal_downsample, dropout)
+ self.conv1 = CausalConv3d(z_dim * 2, z_dim * 2, 1)
+ self.conv2 = CausalConv3d(z_dim, z_dim, 1)
+ self.decoder = Decoder3d(dim, z_dim, dim_mult, num_res_blocks,
+ attn_scales, self.temperal_upsample, dropout)
+
+ def forward(self, x):
+ mu, log_var = self.encode(x)
+ z = self.reparameterize(mu, log_var)
+ x_recon = self.decode(z)
+ return x_recon, mu, log_var
+
+ def encode(self, x, scale):
+ self.clear_cache()
+ ## cache
+ t = x.shape[2]
+ iter_ = 1 + (t - 1) // 4
+ scale = [item.to(x.device, x.dtype) for item in scale]
+ ## 对encode输入的x，按时间拆分为1、4、4、4....
+ for i in range(iter_):
+ self._enc_conv_idx = [0]
+ if i == 0:
+ out = self.encoder(
+ x[:, :, :1, :, :],
+ feat_cache=self._enc_feat_map,
+ feat_idx=self._enc_conv_idx)
+ else:
+ out_ = self.encoder(
+ x[:, :, 1 + 4 * (i - 1):1 + 4 * i, :, :],
+ feat_cache=self._enc_feat_map,
+ feat_idx=self._enc_conv_idx)
+ out = torch.cat([out, out_], 2)
+ mu, log_var = self.conv1(out).chunk(2, dim=1)
+ if isinstance(scale[0], torch.Tensor):
+ mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(
+ 1, self.z_dim, 1, 1, 1)
+ else:
+ mu = (mu - scale[0]) * scale[1]
+ x = torch.cat([mu, log_var], dim = 1)
+ self.clear_cache()
+ return x
+
+ def decode(self, z, scale):
+ self.clear_cache()
+ # z: [b,c,t,h,w]
+ scale = [item.to(z.device, z.dtype) for item in scale]
+ if isinstance(scale[0], torch.Tensor):
+ z = z / scale[1].view(1, self.z_dim, 1, 1, 1) + scale[0].view(
+ 1, self.z_dim, 1, 1, 1)
+ else:
+ z = z / scale[1] + scale[0]
+ iter_ = z.shape[2]
+ x = self.conv2(z)
+ for i in range(iter_):
+ self._conv_idx = [0]
+ if i == 0:
+ out = self.decoder(
+ x[:, :, i:i + 1, :, :],
+ feat_cache=self._feat_map,
+ feat_idx=self._conv_idx)
+ else:
+ out_ = self.decoder(
+ x[:, :, i:i + 1, :, :],
+ feat_cache=self._feat_map,
+ feat_idx=self._conv_idx)
+ out = torch.cat([out, out_], 2)
+ self.clear_cache()
+ return out
+
+ def reparameterize(self, mu, log_var):
+ std = torch.exp(0.5 * log_var)
+ eps = torch.randn_like(std)
+ return eps * std + mu
+
+ def sample(self, imgs, deterministic=False):
+ mu, log_var = self.encode(imgs)
+ if deterministic:
+ return mu
+ std = torch.exp(0.5 * log_var.clamp(-30.0, 20.0))
+ return mu + std * torch.randn_like(std)
+
+ def clear_cache(self):
+ self._conv_num = count_conv3d(self.decoder)
+ self._conv_idx = [0]
+ self._feat_map = [None] * self._conv_num
+ #cache encode
+ self._enc_conv_num = count_conv3d(self.encoder)
+ self._enc_conv_idx = [0]
+ self._enc_feat_map = [None] * self._enc_conv_num
+
+
+def _video_vae(z_dim=None, **kwargs):
+ """
+ Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.
+ """
+ # params
+ cfg = dict(
+ dim=96,
+ z_dim=z_dim,
+ dim_mult=[1, 2, 4, 4],
+ num_res_blocks=2,
+ attn_scales=[],
+ temperal_downsample=[False, True, True],
+ dropout=0.0)
+ cfg.update(**kwargs)
+
+ # init model
+ model = AutoencoderKLWan_(**cfg)
+
+ return model
+
+
+class AutoencoderKLWan(ModelMixin, ConfigMixin, FromOriginalModelMixin):
+
+ @register_to_config
+ def __init__(
+ self,
+ latent_channels=16,
+ temporal_compression_ratio=4,
+ spacial_compression_ratio=8
+ ):
+ super().__init__()
+ mean = [
+ -0.7571, -0.7089, -0.9113, 0.1075, -0.1745, 0.9653, -0.1517, 1.5508,
+ 0.4134, -0.0715, 0.5517, -0.3632, -0.1922, -0.9497, 0.2503, -0.2921
+ ]
+ std = [
+ 2.8184, 1.4541, 2.3275, 2.6558, 1.2196, 1.7708, 2.6052, 2.0743,
+ 3.2687, 2.1526, 2.8652, 1.5579, 1.6382, 1.1253, 2.8251, 1.9160
+ ]
+ self.mean = torch.tensor(mean, dtype=torch.float32)
+ self.std = torch.tensor(std, dtype=torch.float32)
+ self.scale = [self.mean, 1.0 / self.std]
+
+ # init model
+ self.model = _video_vae(
+ z_dim=latent_channels,
+ )
+
+ def _encode(self, x: torch.Tensor) -> torch.Tensor:
+ x = [
+ self.model.encode(u.unsqueeze(0), self.scale).squeeze(0)
+ for u in x
+ ]
+ x = torch.stack(x)
+ return x
+
+ @apply_forward_hook
+ def encode(
+ self, x: torch.Tensor, return_dict: bool = True
+ ) -> Union[AutoencoderKLOutput, Tuple[DiagonalGaussianDistribution]]:
+ h = self._encode(x)
+
+ posterior = DiagonalGaussianDistribution(h)
+
+ if not return_dict:
+ return (posterior,)
+ return AutoencoderKLOutput(latent_dist=posterior)
+
+ def _decode(self, zs):
+ dec = [
+ self.model.decode(u.unsqueeze(0), self.scale).clamp_(-1, 1).squeeze(0)
+ for u in zs
+ ]
+ dec = torch.stack(dec)
+
+ return DecoderOutput(sample=dec)
+
+ @apply_forward_hook
+ def decode(self, z: torch.Tensor, return_dict: bool = True) -> Union[DecoderOutput, torch.Tensor]:
+ decoded = self._decode(z).sample
+
+ if not return_dict:
+ return (decoded,)
+ return DecoderOutput(sample=decoded)
+
+ @classmethod
+ def from_pretrained(cls, pretrained_model_path, additional_kwargs={}):
+ def filter_kwargs(cls, kwargs):
+ import inspect
+ sig = inspect.signature(cls.__init__)
+ valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
+ filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
+ return filtered_kwargs
+
+ model = cls(**filter_kwargs(cls, additional_kwargs))
+ if pretrained_model_path.endswith(".safetensors"):
+ from safetensors.torch import load_file, safe_open
+ state_dict = load_file(pretrained_model_path)
+ else:
+ state_dict = torch.load(pretrained_model_path, map_location="cpu")
+ tmp_state_dict = {}
+ for key in state_dict:
+ tmp_state_dict["model." + key] = state_dict[key]
+ state_dict = tmp_state_dict
+ m, u = model.load_state_dict(state_dict, strict=False)
+ print(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")
+ print(m, u)
+ return model
\ No newline at end of file
diff --git a/wan/models/wan_xlm_roberta.py b/wan/models/wan_xlm_roberta.py
new file mode 100644
index 0000000000000000000000000000000000000000..755baf394431bee95e1eac835b5dafe6ed37c5b9
--- /dev/null
+++ b/wan/models/wan_xlm_roberta.py
@@ -0,0 +1,170 @@
+# Modified from transformers.models.xlm_roberta.modeling_xlm_roberta
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+__all__ = ['XLMRoberta', 'xlm_roberta_large']
+
+
+class SelfAttention(nn.Module):
+
+ def __init__(self, dim, num_heads, dropout=0.1, eps=1e-5):
+ assert dim % num_heads == 0
+ super().__init__()
+ self.dim = dim
+ self.num_heads = num_heads
+ self.head_dim = dim // num_heads
+ 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.dropout = nn.Dropout(dropout)
+
+ def forward(self, x, mask):
+ """
+ x: [B, L, C].
+ """
+ b, s, c, n, d = *x.size(), self.num_heads, self.head_dim
+
+ # compute query, key, value
+ q = self.q(x).reshape(b, s, n, d).permute(0, 2, 1, 3)
+ k = self.k(x).reshape(b, s, n, d).permute(0, 2, 1, 3)
+ v = self.v(x).reshape(b, s, n, d).permute(0, 2, 1, 3)
+
+ # compute attention
+ p = self.dropout.p if self.training else 0.0
+ x = F.scaled_dot_product_attention(q, k, v, mask, p)
+ x = x.permute(0, 2, 1, 3).reshape(b, s, c)
+
+ # output
+ x = self.o(x)
+ x = self.dropout(x)
+ return x
+
+
+class AttentionBlock(nn.Module):
+
+ def __init__(self, dim, num_heads, post_norm, dropout=0.1, eps=1e-5):
+ super().__init__()
+ self.dim = dim
+ self.num_heads = num_heads
+ self.post_norm = post_norm
+ self.eps = eps
+
+ # layers
+ self.attn = SelfAttention(dim, num_heads, dropout, eps)
+ self.norm1 = nn.LayerNorm(dim, eps=eps)
+ self.ffn = nn.Sequential(
+ nn.Linear(dim, dim * 4), nn.GELU(), nn.Linear(dim * 4, dim),
+ nn.Dropout(dropout))
+ self.norm2 = nn.LayerNorm(dim, eps=eps)
+
+ def forward(self, x, mask):
+ if self.post_norm:
+ x = self.norm1(x + self.attn(x, mask))
+ x = self.norm2(x + self.ffn(x))
+ else:
+ x = x + self.attn(self.norm1(x), mask)
+ x = x + self.ffn(self.norm2(x))
+ return x
+
+
+class XLMRoberta(nn.Module):
+ """
+ XLMRobertaModel with no pooler and no LM head.
+ """
+
+ def __init__(self,
+ vocab_size=250002,
+ max_seq_len=514,
+ type_size=1,
+ pad_id=1,
+ dim=1024,
+ num_heads=16,
+ num_layers=24,
+ post_norm=True,
+ dropout=0.1,
+ eps=1e-5):
+ super().__init__()
+ self.vocab_size = vocab_size
+ self.max_seq_len = max_seq_len
+ self.type_size = type_size
+ self.pad_id = pad_id
+ self.dim = dim
+ self.num_heads = num_heads
+ self.num_layers = num_layers
+ self.post_norm = post_norm
+ self.eps = eps
+
+ # embeddings
+ self.token_embedding = nn.Embedding(vocab_size, dim, padding_idx=pad_id)
+ self.type_embedding = nn.Embedding(type_size, dim)
+ self.pos_embedding = nn.Embedding(max_seq_len, dim, padding_idx=pad_id)
+ self.dropout = nn.Dropout(dropout)
+
+ # blocks
+ self.blocks = nn.ModuleList([
+ AttentionBlock(dim, num_heads, post_norm, dropout, eps)
+ for _ in range(num_layers)
+ ])
+
+ # norm layer
+ self.norm = nn.LayerNorm(dim, eps=eps)
+
+ def forward(self, ids):
+ """
+ ids: [B, L] of torch.LongTensor.
+ """
+ b, s = ids.shape
+ mask = ids.ne(self.pad_id).long()
+
+ # embeddings
+ x = self.token_embedding(ids) + \
+ self.type_embedding(torch.zeros_like(ids)) + \
+ self.pos_embedding(self.pad_id + torch.cumsum(mask, dim=1) * mask)
+ if self.post_norm:
+ x = self.norm(x)
+ x = self.dropout(x)
+
+ # blocks
+ mask = torch.where(
+ mask.view(b, 1, 1, s).gt(0), 0.0,
+ torch.finfo(x.dtype).min)
+ for block in self.blocks:
+ x = block(x, mask)
+
+ # output
+ if not self.post_norm:
+ x = self.norm(x)
+ return x
+
+
+def xlm_roberta_large(pretrained=False,
+ return_tokenizer=False,
+ device='cpu',
+ **kwargs):
+ """
+ XLMRobertaLarge adapted from Huggingface.
+ """
+ # params
+ cfg = dict(
+ vocab_size=250002,
+ max_seq_len=514,
+ type_size=1,
+ pad_id=1,
+ dim=1024,
+ num_heads=16,
+ num_layers=24,
+ post_norm=True,
+ dropout=0.1,
+ eps=1e-5)
+ cfg.update(**kwargs)
+
+ # init a model on device
+ with torch.device(device):
+ model = XLMRoberta(**cfg)
+ return model
\ No newline at end of file
diff --git a/wan/models/wav2vec.py b/wan/models/wav2vec.py
new file mode 100644
index 0000000000000000000000000000000000000000..5c2fad8bbb300738e3807da27b311cd2a39a2812
--- /dev/null
+++ b/wan/models/wav2vec.py
@@ -0,0 +1,206 @@
+"""
+This module defines the Wav2Vec model, which is a pre-trained model for speech recognition and understanding.
+It inherits from the Wav2Vec2Model class in the transformers library and provides additional functionalities
+such as feature extraction and encoding.
+
+Classes:
+ Wav2VecModel: Inherits from Wav2Vec2Model and adds additional methods for feature extraction and encoding.
+
+Functions:
+ linear_interpolation: Interpolates the features based on the sequence length.
+"""
+
+import torch.nn.functional as F
+from transformers import Wav2Vec2Model
+from transformers.modeling_outputs import BaseModelOutput
+
+
+class Wav2VecModel(Wav2Vec2Model):
+ """
+ Wav2VecModel is a custom model class that extends the Wav2Vec2Model class from the transformers library.
+ It inherits all the functionality of the Wav2Vec2Model and adds additional methods for feature extraction and encoding.
+ ...
+
+ Attributes:
+ base_model (Wav2Vec2Model): The base Wav2Vec2Model object.
+
+ Methods:
+ forward(input_values, seq_len, attention_mask=None, mask_time_indices=None
+ , output_attentions=None, output_hidden_states=None, return_dict=None):
+ Forward pass of the Wav2VecModel.
+ It takes input_values, seq_len, and other optional parameters as input and returns the output of the base model.
+
+ feature_extract(input_values, seq_len):
+ Extracts features from the input_values using the base model.
+
+ encode(extract_features, attention_mask=None, mask_time_indices=None, output_attentions=None, output_hidden_states=None, return_dict=None):
+ Encodes the extracted features using the base model and returns the encoded features.
+ """
+ def forward(
+ self,
+ input_values,
+ seq_len,
+ attention_mask=None,
+ mask_time_indices=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ ):
+ """
+ Forward pass of the Wav2Vec model.
+
+ Args:
+ self: The instance of the model.
+ input_values: The input values (waveform) to the model.
+ seq_len: The sequence length of the input values.
+ attention_mask: Attention mask to be used for the model.
+ mask_time_indices: Mask indices to be used for the model.
+ output_attentions: If set to True, returns attentions.
+ output_hidden_states: If set to True, returns hidden states.
+ return_dict: If set to True, returns a BaseModelOutput instead of a tuple.
+
+ Returns:
+ The output of the Wav2Vec model.
+ """
+ self.config.output_attentions = True
+
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ extract_features = self.feature_extractor(input_values)
+ extract_features = extract_features.transpose(1, 2)
+ extract_features = linear_interpolation(extract_features, seq_len=seq_len)
+
+ if attention_mask is not None:
+ # compute reduced attention_mask corresponding to feature vectors
+ attention_mask = self._get_feature_vector_attention_mask(
+ extract_features.shape[1], attention_mask, add_adapter=False
+ )
+
+ hidden_states, extract_features = self.feature_projection(extract_features)
+ hidden_states = self._mask_hidden_states(
+ hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
+ )
+
+ encoder_outputs = self.encoder(
+ hidden_states,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ hidden_states = encoder_outputs[0]
+
+ if self.adapter is not None:
+ hidden_states = self.adapter(hidden_states)
+
+ if not return_dict:
+ return (hidden_states, ) + encoder_outputs[1:]
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ )
+
+
+ def feature_extract(
+ self,
+ input_values,
+ seq_len,
+ ):
+ """
+ Extracts features from the input values and returns the extracted features.
+
+ Parameters:
+ input_values (torch.Tensor): The input values to be processed.
+ seq_len (torch.Tensor): The sequence lengths of the input values.
+
+ Returns:
+ extracted_features (torch.Tensor): The extracted features from the input values.
+ """
+ extract_features = self.feature_extractor(input_values)
+ extract_features = extract_features.transpose(1, 2)
+ extract_features = linear_interpolation(extract_features, seq_len=seq_len)
+
+ return extract_features
+
+ def encode(
+ self,
+ extract_features,
+ attention_mask=None,
+ mask_time_indices=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ ):
+ """
+ Encodes the input features into the output space.
+
+ Args:
+ extract_features (torch.Tensor): The extracted features from the audio signal.
+ attention_mask (torch.Tensor, optional): Attention mask to be used for padding.
+ mask_time_indices (torch.Tensor, optional): Masked indices for the time dimension.
+ output_attentions (bool, optional): If set to True, returns the attention weights.
+ output_hidden_states (bool, optional): If set to True, returns all hidden states.
+ return_dict (bool, optional): If set to True, returns a BaseModelOutput instead of the tuple.
+
+ Returns:
+ The encoded output features.
+ """
+ self.config.output_attentions = True
+
+ output_hidden_states = (
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+ )
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+ if attention_mask is not None:
+ # compute reduced attention_mask corresponding to feature vectors
+ attention_mask = self._get_feature_vector_attention_mask(
+ extract_features.shape[1], attention_mask, add_adapter=False
+ )
+
+ hidden_states, extract_features = self.feature_projection(extract_features)
+ hidden_states = self._mask_hidden_states(
+ hidden_states, mask_time_indices=mask_time_indices, attention_mask=attention_mask
+ )
+
+ encoder_outputs = self.encoder(
+ hidden_states,
+ attention_mask=attention_mask,
+ output_attentions=output_attentions,
+ output_hidden_states=output_hidden_states,
+ return_dict=return_dict,
+ )
+
+ hidden_states = encoder_outputs[0]
+
+ if self.adapter is not None:
+ hidden_states = self.adapter(hidden_states)
+
+ if not return_dict:
+ return (hidden_states, ) + encoder_outputs[1:]
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ hidden_states=encoder_outputs.hidden_states,
+ attentions=encoder_outputs.attentions,
+ )
+
+
+def linear_interpolation(features, seq_len):
+ """
+ Transpose the features to interpolate linearly.
+
+ Args:
+ features (torch.Tensor): The extracted features to be interpolated.
+ seq_len (torch.Tensor): The sequence lengths of the features.
+
+ Returns:
+ torch.Tensor: The interpolated features.
+ """
+ features = features.transpose(1, 2)
+ output_features = F.interpolate(features, size=seq_len, align_corners=True, mode='linear')
+ return output_features.transpose(1, 2)
diff --git a/wan/pipeline/__init__.py b/wan/pipeline/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/wan/pipeline/__pycache__/__init__.cpython-311.pyc b/wan/pipeline/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..1c6a7cf58fb333024d820a5a8eabb4470a74cddf
Binary files /dev/null and b/wan/pipeline/__pycache__/__init__.cpython-311.pyc differ
diff --git a/wan/pipeline/__pycache__/wan_inference_long_pipeline.cpython-311.pyc b/wan/pipeline/__pycache__/wan_inference_long_pipeline.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..eb067e3d3631fbb55e9b6f0d3b8aa7afafc531a3
Binary files /dev/null and b/wan/pipeline/__pycache__/wan_inference_long_pipeline.cpython-311.pyc differ
diff --git a/wan/pipeline/pipeline_wan_fun_inpaint.py b/wan/pipeline/pipeline_wan_fun_inpaint.py
new file mode 100644
index 0000000000000000000000000000000000000000..70dd114fee28141ca79e0516a066b5d22f595231
--- /dev/null
+++ b/wan/pipeline/pipeline_wan_fun_inpaint.py
@@ -0,0 +1,732 @@
+import inspect
+import math
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.embeddings import get_1d_rotary_pos_embed
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import BaseOutput, logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.video_processor import VideoProcessor
+from einops import rearrange
+from PIL import Image
+from transformers import AutoTokenizer
+
+from wan.models.wan_image_encoder import CLIPModel
+from wan.models.wan_text_encoder import WanT5EncoderModel
+from wan.models.wan_transformer3d import WanTransformer3DModel
+from wan.models.wan_vae import AutoencoderKLWan
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+EXAMPLE_DOC_STRING = """
+ Examples:
+ ```python
+ pass
+ ```
+"""
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: Optional[int] = None,
+ device: Optional[Union[str, torch.device]] = None,
+ timesteps: Optional[List[int]] = None,
+ sigmas: Optional[List[float]] = None,
+ **kwargs,
+):
+ """
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`List[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`List[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+def resize_mask(mask, latent, process_first_frame_only=True):
+ latent_size = latent.size()
+ batch_size, channels, num_frames, height, width = mask.shape
+
+ if process_first_frame_only:
+ target_size = list(latent_size[2:])
+ target_size[0] = 1
+ first_frame_resized = F.interpolate(
+ mask[:, :, 0:1, :, :],
+ size=target_size,
+ mode='trilinear',
+ align_corners=False
+ )
+
+ target_size = list(latent_size[2:])
+ target_size[0] = target_size[0] - 1
+ if target_size[0] != 0:
+ remaining_frames_resized = F.interpolate(
+ mask[:, :, 1:, :, :],
+ size=target_size,
+ mode='trilinear',
+ align_corners=False
+ )
+ resized_mask = torch.cat([first_frame_resized, remaining_frames_resized], dim=2)
+ else:
+ resized_mask = first_frame_resized
+ else:
+ target_size = list(latent_size[2:])
+ resized_mask = F.interpolate(
+ mask,
+ size=target_size,
+ mode='trilinear',
+ align_corners=False
+ )
+ return resized_mask
+
+
+@dataclass
+class WanPipelineOutput(BaseOutput):
+ r"""
+ Output class for CogVideo pipelines.
+
+ Args:
+ video (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
+ List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
+ `(batch_size, num_frames, channels, height, width)`.
+ """
+
+ videos: torch.Tensor
+
+
+class WanFunInpaintPipeline(DiffusionPipeline):
+ r"""
+ Pipeline for text-to-video generation using Wan.
+
+ This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+ library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+ """
+
+ _optional_components = []
+ model_cpu_offload_seq = "text_encoder->clip_image_encoder->transformer->vae"
+
+ _callback_tensor_inputs = [
+ "latents",
+ "prompt_embeds",
+ "negative_prompt_embeds",
+ ]
+
+ def __init__(
+ self,
+ tokenizer: AutoTokenizer,
+ text_encoder: WanT5EncoderModel,
+ vae: AutoencoderKLWan,
+ transformer: WanTransformer3DModel,
+ clip_image_encoder: CLIPModel,
+ scheduler: FlowMatchEulerDiscreteScheduler,
+ ):
+ super().__init__()
+
+ self.register_modules(
+ tokenizer=tokenizer,
+ text_encoder=text_encoder,
+ vae=vae,
+ transformer=transformer,
+ clip_image_encoder=clip_image_encoder,
+ scheduler=scheduler
+ )
+
+ self.video_processor = VideoProcessor(vae_scale_factor=self.vae.spacial_compression_ratio)
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae.spacial_compression_ratio)
+ self.mask_processor = VaeImageProcessor(
+ vae_scale_factor=self.vae.spacial_compression_ratio, do_normalize=False, do_binarize=True, do_convert_grayscale=True
+ )
+
+ def _get_t5_prompt_embeds(
+ self,
+ prompt: Union[str, List[str]] = None,
+ num_videos_per_prompt: int = 1,
+ max_sequence_length: int = 512,
+ device: Optional[torch.device] = None,
+ dtype: Optional[torch.dtype] = None,
+ ):
+ device = device or self._execution_device
+ dtype = dtype or self.text_encoder.dtype
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ batch_size = len(prompt)
+
+ text_inputs = self.tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=max_sequence_length,
+ truncation=True,
+ add_special_tokens=True,
+ return_tensors="pt",
+ )
+ text_input_ids = text_inputs.input_ids
+ prompt_attention_mask = text_inputs.attention_mask
+ untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+ removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1 : -1])
+ logger.warning(
+ "The following part of your input was truncated because `max_sequence_length` is set to "
+ f" {max_sequence_length} tokens: {removed_text}"
+ )
+
+ seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long()
+ prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask.to(device))[0]
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
+ _, seq_len, _ = prompt_embeds.shape
+ prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+ prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
+
+ return [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
+
+ def encode_prompt(
+ self,
+ prompt: Union[str, List[str]],
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ do_classifier_free_guidance: bool = True,
+ num_videos_per_prompt: int = 1,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ max_sequence_length: int = 512,
+ device: Optional[torch.device] = None,
+ dtype: Optional[torch.dtype] = None,
+ ):
+ r"""
+ Encodes the prompt into text encoder hidden states.
+
+ Args:
+ prompt (`str` or `List[str]`, *optional*):
+ prompt to be encoded
+ negative_prompt (`str` or `List[str]`, *optional*):
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+ less than `1`).
+ do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
+ Whether to use classifier free guidance or not.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+ argument.
+ device: (`torch.device`, *optional*):
+ torch device
+ dtype: (`torch.dtype`, *optional*):
+ torch dtype
+ """
+ device = device or self._execution_device
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ if prompt is not None:
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ if prompt_embeds is None:
+ prompt_embeds = self._get_t5_prompt_embeds(
+ prompt=prompt,
+ num_videos_per_prompt=num_videos_per_prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ dtype=dtype,
+ )
+
+ if do_classifier_free_guidance and negative_prompt_embeds is None:
+ negative_prompt = negative_prompt or ""
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+ if prompt is not None and type(prompt) is not type(negative_prompt):
+ raise TypeError(
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+ f" {type(prompt)}."
+ )
+ elif batch_size != len(negative_prompt):
+ raise ValueError(
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+ " the batch size of `prompt`."
+ )
+
+ negative_prompt_embeds = self._get_t5_prompt_embeds(
+ prompt=negative_prompt,
+ num_videos_per_prompt=num_videos_per_prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ dtype=dtype,
+ )
+
+ return prompt_embeds, negative_prompt_embeds
+
+ def prepare_latents(
+ self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None
+ ):
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ shape = (
+ batch_size,
+ num_channels_latents,
+ (num_frames - 1) // self.vae.temporal_compression_ratio + 1,
+ height // self.vae.spacial_compression_ratio,
+ width // self.vae.spacial_compression_ratio,
+ )
+
+ if latents is None:
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ else:
+ latents = latents.to(device)
+
+ # scale the initial noise by the standard deviation required by the scheduler
+ if hasattr(self.scheduler, "init_noise_sigma"):
+ latents = latents * self.scheduler.init_noise_sigma
+ return latents
+
+ def prepare_mask_latents(
+ self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance, noise_aug_strength
+ ):
+ # resize the mask to latents shape as we concatenate the mask to the latents
+ # we do that before converting to dtype to avoid breaking in case we're using cpu_offload
+ # and half precision
+
+ if mask is not None:
+ mask = mask.to(device=device, dtype=self.vae.dtype)
+ bs = 1
+ new_mask = []
+ for i in range(0, mask.shape[0], bs):
+ mask_bs = mask[i : i + bs]
+ mask_bs = self.vae.encode(mask_bs)[0]
+ mask_bs = mask_bs.mode()
+ new_mask.append(mask_bs)
+ mask = torch.cat(new_mask, dim = 0)
+ # mask = mask * self.vae.config.scaling_factor
+
+ if masked_image is not None:
+ masked_image = masked_image.to(device=device, dtype=self.vae.dtype)
+ bs = 1
+ new_mask_pixel_values = []
+ for i in range(0, masked_image.shape[0], bs):
+ mask_pixel_values_bs = masked_image[i : i + bs]
+ mask_pixel_values_bs = self.vae.encode(mask_pixel_values_bs)[0]
+ mask_pixel_values_bs = mask_pixel_values_bs.mode()
+ new_mask_pixel_values.append(mask_pixel_values_bs)
+ masked_image_latents = torch.cat(new_mask_pixel_values, dim = 0)
+ # masked_image_latents = masked_image_latents * self.vae.config.scaling_factor
+ else:
+ masked_image_latents = None
+
+ return mask, masked_image_latents
+
+ def decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
+ frames = self.vae.decode(latents.to(self.vae.dtype)).sample
+ frames = (frames / 2 + 0.5).clamp(0, 1)
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+ frames = frames.cpu().float().numpy()
+ return frames
+
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
+ def prepare_extra_step_kwargs(self, generator, eta):
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+ # and should be between [0, 1]
+
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+ extra_step_kwargs = {}
+ if accepts_eta:
+ extra_step_kwargs["eta"] = eta
+
+ # check if the scheduler accepts generator
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+ if accepts_generator:
+ extra_step_kwargs["generator"] = generator
+ return extra_step_kwargs
+
+ # Copied from diffusers.pipelines.latte.pipeline_latte.LattePipeline.check_inputs
+ def check_inputs(
+ self,
+ prompt,
+ height,
+ width,
+ negative_prompt,
+ callback_on_step_end_tensor_inputs,
+ prompt_embeds=None,
+ negative_prompt_embeds=None,
+ ):
+ if height % 8 != 0 or width % 8 != 0:
+ raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+ if callback_on_step_end_tensor_inputs is not None and not all(
+ k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+ ):
+ raise ValueError(
+ f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+ )
+ if prompt is not None and prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+ " only forward one of the two."
+ )
+ elif prompt is None and prompt_embeds is None:
+ raise ValueError(
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+ )
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+ if prompt is not None and negative_prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+ )
+
+ if negative_prompt is not None and negative_prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+ )
+
+ if prompt_embeds is not None and negative_prompt_embeds is not None:
+ if prompt_embeds.shape != negative_prompt_embeds.shape:
+ raise ValueError(
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+ f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+ f" {negative_prompt_embeds.shape}."
+ )
+
+ @property
+ def guidance_scale(self):
+ return self._guidance_scale
+
+ @property
+ def num_timesteps(self):
+ return self._num_timesteps
+
+ @property
+ def attention_kwargs(self):
+ return self._attention_kwargs
+
+ @property
+ def interrupt(self):
+ return self._interrupt
+
+ @torch.no_grad()
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
+ def __call__(
+ self,
+ prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ height: int = 480,
+ width: int = 720,
+ video: Union[torch.FloatTensor] = None,
+ mask_video: Union[torch.FloatTensor] = None,
+ num_frames: int = 49,
+ num_inference_steps: int = 50,
+ timesteps: Optional[List[int]] = None,
+ guidance_scale: float = 6,
+ num_videos_per_prompt: int = 1,
+ eta: float = 0.0,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.FloatTensor] = None,
+ prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ output_type: str = "numpy",
+ return_dict: bool = False,
+ callback_on_step_end: Optional[
+ Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+ ] = None,
+ attention_kwargs: Optional[Dict[str, Any]] = None,
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_image: Image = None,
+ max_sequence_length: int = 512,
+ comfyui_progressbar: bool = False,
+ ) -> Union[WanPipelineOutput, Tuple]:
+ """
+ Function invoked when calling the pipeline for generation.
+ Args:
+
+ Examples:
+
+ Returns:
+
+ """
+
+ if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+ callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+ num_videos_per_prompt = 1
+
+ # 1. Check inputs. Raise error if not correct
+ self.check_inputs(
+ prompt,
+ height,
+ width,
+ negative_prompt,
+ callback_on_step_end_tensor_inputs,
+ prompt_embeds,
+ negative_prompt_embeds,
+ )
+ self._guidance_scale = guidance_scale
+ self._attention_kwargs = attention_kwargs
+ self._interrupt = False
+
+ # 2. Default call parameters
+ if prompt is not None and isinstance(prompt, str):
+ batch_size = 1
+ elif prompt is not None and isinstance(prompt, list):
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ device = self._execution_device
+ weight_dtype = self.text_encoder.dtype
+
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+ # corresponds to doing no classifier free guidance.
+ do_classifier_free_guidance = guidance_scale > 1.0
+
+ # 3. Encode input prompt
+ prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+ prompt,
+ negative_prompt,
+ do_classifier_free_guidance,
+ num_videos_per_prompt=num_videos_per_prompt,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ )
+ if do_classifier_free_guidance:
+ prompt_embeds = negative_prompt_embeds + prompt_embeds
+
+ # 4. Prepare timesteps
+ if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
+ timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps, mu=1)
+ else:
+ timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+ self._num_timesteps = len(timesteps)
+ if comfyui_progressbar:
+ from comfy.utils import ProgressBar
+ pbar = ProgressBar(num_inference_steps + 2)
+
+ # 5. Prepare latents.
+ if video is not None:
+ video_length = video.shape[2]
+ init_video = self.image_processor.preprocess(rearrange(video, "b c f h w -> (b f) c h w"), height=height, width=width)
+ init_video = init_video.to(dtype=torch.float32)
+ init_video = rearrange(init_video, "(b f) c h w -> b c f h w", f=video_length)
+ else:
+ init_video = None
+
+ latent_channels = self.vae.config.latent_channels
+ latents = self.prepare_latents(
+ batch_size * num_videos_per_prompt,
+ latent_channels,
+ num_frames,
+ height,
+ width,
+ weight_dtype,
+ device,
+ generator,
+ latents,
+ )
+ if comfyui_progressbar:
+ pbar.update(1)
+
+ # Prepare mask latent variables
+ if init_video is not None:
+ if (mask_video == 255).all():
+ mask_latents = torch.tile(
+ torch.zeros_like(latents)[:, :1].to(device, weight_dtype), [1, 4, 1, 1, 1]
+ )
+ masked_video_latents = torch.zeros_like(latents).to(device, weight_dtype)
+
+ mask_input = torch.cat([mask_latents] * 2) if do_classifier_free_guidance else mask_latents
+ masked_video_latents_input = (
+ torch.cat([masked_video_latents] * 2) if do_classifier_free_guidance else masked_video_latents
+ )
+ y = torch.cat([mask_input, masked_video_latents_input], dim=1).to(device, weight_dtype)
+ else:
+ bs, _, video_length, height, width = video.size()
+ mask_condition = self.mask_processor.preprocess(rearrange(mask_video, "b c f h w -> (b f) c h w"), height=height, width=width)
+ mask_condition = mask_condition.to(dtype=torch.float32)
+ mask_condition = rearrange(mask_condition, "(b f) c h w -> b c f h w", f=video_length)
+
+ masked_video = init_video * (torch.tile(mask_condition, [1, 3, 1, 1, 1]) < 0.5)
+ _, masked_video_latents = self.prepare_mask_latents(
+ None,
+ masked_video,
+ batch_size,
+ height,
+ width,
+ weight_dtype,
+ device,
+ generator,
+ do_classifier_free_guidance,
+ noise_aug_strength=None,
+ )
+
+ mask_condition = torch.concat(
+ [
+ torch.repeat_interleave(mask_condition[:, :, 0:1], repeats=4, dim=2),
+ mask_condition[:, :, 1:]
+ ], dim=2
+ )
+ mask_condition = mask_condition.view(bs, mask_condition.shape[2] // 4, 4, height, width)
+ mask_condition = mask_condition.transpose(1, 2)
+ mask_latents = resize_mask(1 - mask_condition, masked_video_latents, True).to(device, weight_dtype)
+
+ mask_input = torch.cat([mask_latents] * 2) if do_classifier_free_guidance else mask_latents
+ masked_video_latents_input = (
+ torch.cat([masked_video_latents] * 2) if do_classifier_free_guidance else masked_video_latents
+ )
+
+ y = torch.cat([mask_input, masked_video_latents_input], dim=1).to(device, weight_dtype)
+
+ # Prepare clip latent variables
+ if clip_image is not None:
+ clip_image = TF.to_tensor(clip_image).sub_(0.5).div_(0.5).to(device, weight_dtype)
+ clip_context = self.clip_image_encoder([clip_image[:, None, :, :]])
+ clip_context = (
+ torch.cat([clip_context] * 2) if do_classifier_free_guidance else clip_context
+ )
+ else:
+ clip_image = Image.new("RGB", (512, 512), color=(0, 0, 0))
+ clip_image = TF.to_tensor(clip_image).sub_(0.5).div_(0.5).to(device, weight_dtype)
+ clip_context = self.clip_image_encoder([clip_image[:, None, :, :]])
+ clip_context = (
+ torch.cat([clip_context] * 2) if do_classifier_free_guidance else clip_context
+ )
+ clip_context = torch.zeros_like(clip_context)
+ if comfyui_progressbar:
+ pbar.update(1)
+
+ # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+ target_shape = (self.vae.latent_channels, (num_frames - 1) // self.vae.temporal_compression_ratio + 1, width // self.vae.spacial_compression_ratio, height // self.vae.spacial_compression_ratio)
+ seq_len = math.ceil((target_shape[2] * target_shape[3]) / (self.transformer.config.patch_size[1] * self.transformer.config.patch_size[2]) * target_shape[1])
+ # 7. Denoising loop
+ num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+ if self.interrupt:
+ continue
+
+ latent_model_input = torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+ if hasattr(self.scheduler, "scale_model_input"):
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+ timestep = t.expand(latent_model_input.shape[0])
+
+ # predict noise model_output
+ with torch.cuda.amp.autocast(dtype=weight_dtype):
+ noise_pred = self.transformer(
+ x=latent_model_input,
+ context=prompt_embeds,
+ t=timestep,
+ seq_len=seq_len,
+ y=y,
+ clip_fea=clip_context,
+ )
+
+ # perform guidance
+ if do_classifier_free_guidance:
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+ noise_pred = noise_pred_uncond + self.guidance_scale * (noise_pred_text - noise_pred_uncond)
+
+ # compute the previous noisy sample x_t -> x_t-1
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+ if callback_on_step_end is not None:
+ callback_kwargs = {}
+ for k in callback_on_step_end_tensor_inputs:
+ callback_kwargs[k] = locals()[k]
+ callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+ latents = callback_outputs.pop("latents", latents)
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+ negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+ progress_bar.update()
+ if comfyui_progressbar:
+ pbar.update(1)
+
+ if output_type == "numpy":
+ video = self.decode_latents(latents)
+ elif not output_type == "latent":
+ video = self.decode_latents(latents)
+ video = self.video_processor.postprocess_video(video=video, output_type=output_type)
+ else:
+ video = latents
+
+ # Offload all models
+ self.maybe_free_model_hooks()
+
+ if not return_dict:
+ video = torch.from_numpy(video)
+
+ return WanPipelineOutput(videos=video)
diff --git a/wan/pipeline/wan_inference_long_pipeline.py b/wan/pipeline/wan_inference_long_pipeline.py
new file mode 100644
index 0000000000000000000000000000000000000000..c871153e024c80150815b5a3033225f1e0cabca2
--- /dev/null
+++ b/wan/pipeline/wan_inference_long_pipeline.py
@@ -0,0 +1,804 @@
+import inspect
+import math
+import random
+from dataclasses import dataclass
+from functools import partial
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+from tqdm import tqdm
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.embeddings import get_1d_rotary_pos_embed
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import BaseOutput, logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.video_processor import VideoProcessor
+from einops import rearrange
+from PIL import Image
+from transformers import AutoTokenizer
+from transformers import Wav2Vec2Model, Wav2Vec2Processor
+
+from wan.models.vocal_projector_fantasy import split_audio_sequence, split_tensor_with_padding
+from wan.models.wan_fantasy_transformer3d_1B import WanTransformer3DFantasyModel
+from wan.models.wan_image_encoder import CLIPModel
+from wan.models.wan_text_encoder import WanT5EncoderModel
+from wan.models.wan_vae import AutoencoderKLWan
+from wan.utils.color_correction import match_and_blend_colors
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+ Examples:
+ ```python
+ pass
+ ```
+"""
+
+def enable_vram_management_recursively(model: torch.nn.Module, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None, total_num_param=0):
+ for name, module in model.named_children():
+ for source_module, target_module in module_map.items():
+ if isinstance(module, source_module):
+ num_param = sum(p.numel() for p in module.parameters())
+ if max_num_param is not None and total_num_param + num_param > max_num_param:
+ module_config_ = overflow_module_config
+ else:
+ module_config_ = module_config
+ module_ = target_module(module, **module_config_)
+ setattr(model, name, module_)
+ total_num_param += num_param
+ break
+ else:
+ total_num_param = enable_vram_management_recursively(module, module_map, module_config, max_num_param, overflow_module_config, total_num_param)
+ return total_num_param
+
+def enable_vram_management(model, module_map: dict, module_config: dict, max_num_param=None, overflow_module_config: dict = None):
+ enable_vram_management_recursively(model, module_map, module_config, max_num_param, overflow_module_config, total_num_param=0)
+ model.vram_management_enabled = True
+
+def timestep_transform(
+ t,
+ shift=5.0,
+ num_timesteps=1000,
+):
+ t = t / num_timesteps
+ # shift the timestep based on ratio
+ new_t = shift * t / (1 + (shift - 1) * t)
+ new_t = new_t * num_timesteps
+ return new_t
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: Optional[int] = None,
+ device: Optional[Union[str, torch.device]] = None,
+ timesteps: Optional[List[int]] = None,
+ sigmas: Optional[List[float]] = None,
+ **kwargs,
+):
+ """
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`List[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`List[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+def resize_mask(mask, latent, process_first_frame_only=True):
+ latent_size = latent.size()
+ batch_size, channels, num_frames, height, width = mask.shape
+
+ if process_first_frame_only:
+ target_size = list(latent_size[2:])
+ target_size[0] = 1
+ first_frame_resized = F.interpolate(
+ mask[:, :, 0:1, :, :],
+ size=target_size,
+ mode='trilinear',
+ align_corners=False
+ )
+
+ target_size = list(latent_size[2:])
+ target_size[0] = target_size[0] - 1
+ if target_size[0] != 0:
+ remaining_frames_resized = F.interpolate(
+ mask[:, :, 1:, :, :],
+ size=target_size,
+ mode='trilinear',
+ align_corners=False
+ )
+ resized_mask = torch.cat([first_frame_resized, remaining_frames_resized], dim=2)
+ else:
+ resized_mask = first_frame_resized
+ else:
+ target_size = list(latent_size[2:])
+ resized_mask = F.interpolate(
+ mask,
+ size=target_size,
+ mode='trilinear',
+ align_corners=False
+ )
+ return resized_mask
+
+
+@dataclass
+class WanI2VPipelineTalkingInferenceLongOutput(BaseOutput):
+ r"""
+ Output class for CogVideo pipelines.
+
+ Args:
+ video (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
+ List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
+ `(batch_size, num_frames, channels, height, width)`.
+ """
+
+ videos: torch.Tensor
+
+
+class WanI2VTalkingInferenceLongPipeline(DiffusionPipeline):
+ r"""
+ Pipeline for text-to-video generation using Wan.
+
+ This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+ library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+ """
+
+ _optional_components = []
+ model_cpu_offload_seq = "text_encoder->clip_image_encoder->transformer->vae"
+
+ _callback_tensor_inputs = [
+ "latents",
+ "prompt_embeds",
+ "negative_prompt_embeds",
+ ]
+
+ def __init__(
+ self,
+ tokenizer: AutoTokenizer,
+ text_encoder: WanT5EncoderModel,
+ vae: AutoencoderKLWan,
+ transformer: WanTransformer3DFantasyModel,
+ clip_image_encoder: CLIPModel,
+ scheduler: FlowMatchEulerDiscreteScheduler,
+ wav2vec_processor: Wav2Vec2Processor,
+ wav2vec: Wav2Vec2Model,
+ ):
+ super().__init__()
+
+ self.register_modules(
+ tokenizer=tokenizer,
+ text_encoder=text_encoder,
+ vae=vae,
+ transformer=transformer,
+ clip_image_encoder=clip_image_encoder,
+ scheduler=scheduler,
+ wav2vec_processor=wav2vec_processor,
+ wav2vec=wav2vec,
+ )
+
+ self.video_processor = VideoProcessor(vae_scale_factor=self.vae.config.spacial_compression_ratio)
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae.config.spacial_compression_ratio)
+ self.mask_processor = VaeImageProcessor(
+ vae_scale_factor=self.vae.spacial_compression_ratio, do_normalize=False, do_binarize=True,
+ do_convert_grayscale=True
+ )
+
+
+
+ def _get_t5_prompt_embeds(
+ self,
+ prompt: Union[str, List[str]] = None,
+ num_videos_per_prompt: int = 1,
+ max_sequence_length: int = 512,
+ device: Optional[torch.device] = None,
+ dtype: Optional[torch.dtype] = None,
+ ):
+ device = device or self._execution_device
+ dtype = dtype or self.text_encoder.dtype
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ batch_size = len(prompt)
+
+ text_inputs = self.tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=max_sequence_length,
+ truncation=True,
+ add_special_tokens=True,
+ return_tensors="pt",
+ )
+ text_input_ids = text_inputs.input_ids
+ prompt_attention_mask = text_inputs.attention_mask
+ untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+ removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1: -1])
+ logger.warning(
+ "The following part of your input was truncated because `max_sequence_length` is set to "
+ f" {max_sequence_length} tokens: {removed_text}"
+ )
+
+ seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long()
+ prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask.to(device))[0]
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
+ _, seq_len, _ = prompt_embeds.shape
+ prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+ prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
+
+ return [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
+
+ def encode_prompt(
+ self,
+ prompt: Union[str, List[str]],
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ do_classifier_free_guidance: bool = True,
+ num_videos_per_prompt: int = 1,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ max_sequence_length: int = 512,
+ device: Optional[torch.device] = None,
+ dtype: Optional[torch.dtype] = None,
+ ):
+ r"""
+ Encodes the prompt into text encoder hidden states.
+
+ Args:
+ prompt (`str` or `List[str]`, *optional*):
+ prompt to be encoded
+ negative_prompt (`str` or `List[str]`, *optional*):
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+ less than `1`).
+ do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
+ Whether to use classifier free guidance or not.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+ argument.
+ device: (`torch.device`, *optional*):
+ torch device
+ dtype: (`torch.dtype`, *optional*):
+ torch dtype
+ """
+ device = device or self._execution_device
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ if prompt is not None:
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ if prompt_embeds is None:
+ prompt_embeds = self._get_t5_prompt_embeds(
+ prompt=prompt,
+ num_videos_per_prompt=num_videos_per_prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ dtype=dtype,
+ )
+
+ if do_classifier_free_guidance and negative_prompt_embeds is None:
+ negative_prompt = negative_prompt or ""
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+ if prompt is not None and type(prompt) is not type(negative_prompt):
+ raise TypeError(
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+ f" {type(prompt)}."
+ )
+ elif batch_size != len(negative_prompt):
+ raise ValueError(
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+ " the batch size of `prompt`."
+ )
+
+ negative_prompt_embeds = self._get_t5_prompt_embeds(
+ prompt=negative_prompt,
+ num_videos_per_prompt=num_videos_per_prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ dtype=dtype,
+ )
+
+ return prompt_embeds, negative_prompt_embeds
+
+ def prepare_latents(
+ self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None
+ ):
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ shape = (
+ batch_size,
+ num_channels_latents,
+ (num_frames - 1) // self.vae.config.temporal_compression_ratio + 1,
+ height // self.vae.config.spacial_compression_ratio,
+ width // self.vae.config.spacial_compression_ratio,
+ )
+
+ if latents is None:
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ else:
+ latents = latents.to(device)
+
+ # scale the initial noise by the standard deviation required by the scheduler
+ if hasattr(self.scheduler, "init_noise_sigma"):
+ latents = latents * self.scheduler.init_noise_sigma
+ return latents
+
+ def prepare_mask_latents(
+ self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance,
+ noise_aug_strength
+ ):
+ # resize the mask to latents shape as we concatenate the mask to the latents
+ # we do that before converting to dtype to avoid breaking in case we're using cpu_offload
+ # and half precision
+
+ if mask is not None:
+ mask = mask.to(device=device, dtype=self.vae.dtype)
+ bs = 1
+ new_mask = []
+ for i in range(0, mask.shape[0], bs):
+ mask_bs = mask[i: i + bs]
+ mask_bs = self.vae.encode(mask_bs)[0]
+ mask_bs = mask_bs.mode()
+ new_mask.append(mask_bs)
+ mask = torch.cat(new_mask, dim=0)
+ # mask = mask * self.vae.config.scaling_factor
+
+ if masked_image is not None:
+ masked_image = masked_image.to(device=device, dtype=self.vae.dtype)
+ bs = 1
+ new_mask_pixel_values = []
+ for i in range(0, masked_image.shape[0], bs):
+ mask_pixel_values_bs = masked_image[i: i + bs]
+ mask_pixel_values_bs = self.vae.encode(mask_pixel_values_bs)[0]
+ mask_pixel_values_bs = mask_pixel_values_bs.mode()
+ new_mask_pixel_values.append(mask_pixel_values_bs)
+ masked_image_latents = torch.cat(new_mask_pixel_values, dim=0)
+ # masked_image_latents = masked_image_latents * self.vae.config.scaling_factor
+ else:
+ masked_image_latents = None
+
+ return mask, masked_image_latents
+
+ def decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
+ frames = self.vae.decode(latents.to(self.vae.dtype)).sample
+ frames = frames.cpu()
+ frames = (frames / 2 + 0.5).clamp(0, 1)
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+ frames = frames.float().numpy()
+ return frames
+
+ def decode_latents_audio_video(self, latents: torch.Tensor) -> torch.Tensor:
+ frames = self.vae.decode(latents.to(self.vae.dtype)).sample
+ # frames = (frames / 2 + 0.5).clamp(0, 1)
+ frames = frames.cpu().float()
+ return frames
+
+
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
+ def prepare_extra_step_kwargs(self, generator, eta):
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+ # and should be between [0, 1]
+
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+ extra_step_kwargs = {}
+ if accepts_eta:
+ extra_step_kwargs["eta"] = eta
+
+ # check if the scheduler accepts generator
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+ if accepts_generator:
+ extra_step_kwargs["generator"] = generator
+ return extra_step_kwargs
+
+ # Copied from diffusers.pipelines.latte.pipeline_latte.LattePipeline.check_inputs
+ def check_inputs(
+ self,
+ prompt,
+ height,
+ width,
+ negative_prompt,
+ callback_on_step_end_tensor_inputs,
+ prompt_embeds=None,
+ negative_prompt_embeds=None,
+ ):
+ if height % 8 != 0 or width % 8 != 0:
+ raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+ if callback_on_step_end_tensor_inputs is not None and not all(
+ k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+ ):
+ raise ValueError(
+ f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+ )
+ if prompt is not None and prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+ " only forward one of the two."
+ )
+ elif prompt is None and prompt_embeds is None:
+ raise ValueError(
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+ )
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+ if prompt is not None and negative_prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+ )
+
+ if negative_prompt is not None and negative_prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+ )
+
+ if prompt_embeds is not None and negative_prompt_embeds is not None:
+ if prompt_embeds.shape != negative_prompt_embeds.shape:
+ raise ValueError(
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+ f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+ f" {negative_prompt_embeds.shape}."
+ )
+
+ def infer_add_noise(
+ self,
+ original_samples: torch.FloatTensor,
+ noise: torch.FloatTensor,
+ timesteps: torch.IntTensor,
+ ) -> torch.FloatTensor:
+ """
+ compatible with diffusers add_noise()
+ """
+ timesteps = timesteps.float() / self.num_timesteps
+ timesteps = timesteps.view(timesteps.shape + (1,) * (len(noise.shape)-1))
+ return (1 - timesteps) * original_samples + timesteps * noise
+
+ @property
+ def guidance_scale(self):
+ return self._guidance_scale
+
+ # @property
+ # def num_timesteps(self):
+ # return self._num_timesteps
+
+ @property
+ def attention_kwargs(self):
+ return self._attention_kwargs
+
+ @property
+ def interrupt(self):
+ return self._interrupt
+
+ @torch.no_grad()
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
+ def __call__(
+ self,
+ prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ height: int = 480,
+ width: int = 720,
+ video: Union[torch.FloatTensor] = None,
+ mask_video: Union[torch.FloatTensor] = None,
+ num_frames: int = 81,
+ num_inference_steps: int = 50,
+ timesteps: Optional[List[int]] = None,
+ guidance_scale: float = 6,
+ num_videos_per_prompt: int = 1,
+ eta: float = 0.0,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.FloatTensor] = None,
+ prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ output_type: str = "numpy",
+ return_dict: bool = False,
+ callback_on_step_end: Optional[
+ Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+ ] = None,
+ attention_kwargs: Optional[Dict[str, Any]] = None,
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_image: Image = None,
+ max_sequence_length: int = 512,
+ text_guide_scale=None,
+ audio_guide_scale=None,
+ vocal_input_values=None,
+ motion_frame=None,
+ fps=None,
+ sr=None,
+ cond_file_path=None,
+ seed=None,
+ overlap_window_length=None,
+ overlapping_weight_scheme="uniform",
+ ) -> Union[WanI2VPipelineTalkingInferenceLongOutput, Tuple]:
+ """
+ Function invoked when calling the pipeline for generation.
+ Args:
+
+ Examples:
+
+ Returns:
+
+ """
+
+ if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+ callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+ num_videos_per_prompt = 1
+
+ # 1. Check inputs. Raise error if not correct
+ self.check_inputs(
+ prompt,
+ height,
+ width,
+ negative_prompt,
+ callback_on_step_end_tensor_inputs,
+ prompt_embeds,
+ negative_prompt_embeds,
+ )
+ self._guidance_scale = guidance_scale
+ self._attention_kwargs = attention_kwargs
+ self._interrupt = False
+
+ # 2. Default call parameters
+ if prompt is not None and isinstance(prompt, str):
+ batch_size = 1
+ elif prompt is not None and isinstance(prompt, list):
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ device = self._execution_device
+ weight_dtype = self.text_encoder.dtype
+
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+ # corresponds to doing no classifier free guidance.
+ do_classifier_free_guidance = guidance_scale > 1.0
+
+ # 3. Encode input prompt
+ prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+ prompt,
+ negative_prompt,
+ do_classifier_free_guidance,
+ num_videos_per_prompt=num_videos_per_prompt,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ )
+ if do_classifier_free_guidance:
+ # prompt_embeds = negative_prompt_embeds + prompt_embeds + prompt_embeds
+ prompt_embeds = negative_prompt_embeds + negative_prompt_embeds + prompt_embeds
+
+ clip_length = 81
+ audio_token_per_frame = int(sr / fps)
+ max_audio_index = vocal_input_values.shape[0]
+ total_frames = int(max_audio_index / audio_token_per_frame)
+ frames_per_batch = 21
+ if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
+ timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps, mu=1)
+ else:
+ timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+ self._num_timesteps = len(timesteps)
+
+ latent_channels = self.vae.config.latent_channels
+ latents = self.prepare_latents(
+ batch_size * num_videos_per_prompt,
+ latent_channels,
+ total_frames,
+ height,
+ width,
+ weight_dtype,
+ device,
+ generator,
+ latents,
+ )
+ infer_length = latents.size()[2]
+ num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+ latents_all = latents.clone()
+ clip_image = cond_image = Image.open(cond_file_path).convert('RGB')
+ cond_image = cond_image.resize([width, height])
+ clip_image = clip_image.resize([width, height])
+ clip_image = torch.from_numpy(np.array(clip_image)).permute(2, 0, 1)
+ clip_image = clip_image / 255
+ clip_image = (clip_image - 0.5) * 2 # C H W
+ cond_image = torch.from_numpy(np.array(cond_image)).permute(2, 0, 1).unsqueeze(1).unsqueeze(0)
+ cond_image = cond_image / 255
+ cond_image = (cond_image - 0.5) * 2 # normalization
+ cond_image = cond_image.to(device) # 1 C 1 H W
+
+ clip_image = clip_image.to(device, weight_dtype)
+ clip_context = self.clip_image_encoder([clip_image[:, None, :, :]])
+ clip_context = (torch.cat([clip_context, clip_context, clip_context], dim=0) if do_classifier_free_guidance else clip_context)
+ video_frames = torch.zeros(1, cond_image.shape[1], clip_length - cond_image.shape[2], height, width).to(device)
+ padding_frames_pixels_values = torch.concat([cond_image, video_frames], dim=2).to(dtype=torch.float32)
+ _, masked_video_latents = self.prepare_mask_latents(
+ None,
+ padding_frames_pixels_values,
+ batch_size,
+ height,
+ width,
+ weight_dtype,
+ device,
+ generator,
+ do_classifier_free_guidance,
+ noise_aug_strength=None,
+ ) # [1, 16, 21, 64, 64]
+ msk = torch.ones(1, clip_length, masked_video_latents.size()[-2], masked_video_latents.size()[-1], device=device)
+ msk[:, 1:] = 0
+ msk = torch.concat([torch.repeat_interleave(msk[:, 0:1], repeats=4, dim=1), msk[:, 1:]], dim=1)
+ msk = msk.view(1, msk.shape[1] // 4, 4, masked_video_latents.size()[-2], masked_video_latents.size()[-1])
+ msk = msk.transpose(1, 2).to(dtype=torch.float32)
+ mask_input = torch.cat([msk] * 3) if do_classifier_free_guidance else msk
+ masked_video_latents_input = (torch.cat([masked_video_latents] * 3) if do_classifier_free_guidance else masked_video_latents)
+ y = torch.cat([mask_input.to(device), masked_video_latents_input.to(device)], dim=1).to(device, weight_dtype)
+
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+ if self.interrupt:
+ continue
+
+ pred_latents = torch.zeros_like(latents_all, dtype=latents_all.dtype, )
+ arrive_last_frame = False
+ index_start = 0
+ overlap_window_length = overlap_window_length # [5, 7, 10] longer length --> higher quality
+ index_end = index_start + frames_per_batch
+ index_previous_end = index_end
+ while index_end <= infer_length:
+ self.scheduler._step_index = None
+ idx_list = [ii % latents_all.shape[2] for ii in range(index_start, index_end)]
+
+ if index_end == infer_length:
+ idx_list_audio = [ii % max_audio_index for ii in range(index_start * 4 * audio_token_per_frame, max_audio_index)]
+ else:
+ idx_list_audio = [ii % max_audio_index for ii in range(index_start * 4 * audio_token_per_frame, index_start * 4 * audio_token_per_frame + clip_length * audio_token_per_frame)]
+
+ # idx_list_audio = [ii % max_audio_index for ii in range(index_start * 4 * audio_token_per_frame, index_end * 4 * audio_token_per_frame)]
+ latents = latents_all[:, :, idx_list].clone()
+ sub_vocal_input_values = vocal_input_values[idx_list_audio]
+ sub_vocal_input_values = self.wav2vec_processor(sub_vocal_input_values, sampling_rate=sr, return_tensors="pt").input_values.to(device)
+ sub_vocal_embeddings = self.wav2vec(sub_vocal_input_values).last_hidden_state
+ latent_model_input = torch.cat([latents] * 3) if do_classifier_free_guidance else latents
+ if hasattr(self.scheduler, "scale_model_input"):
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+ timestep = t.expand(latent_model_input.shape[0])
+ target_shape = (self.vae.config.latent_channels, (num_frames - 1) // self.vae.config.temporal_compression_ratio + 1, width // self.vae.config.spacial_compression_ratio, height // self.vae.config.spacial_compression_ratio)
+ seq_len = math.ceil((target_shape[2] * target_shape[3]) / (self.transformer.config.patch_size[1] * self.transformer.config.patch_size[2]) * target_shape[1])
+ if text_guide_scale is not None and audio_guide_scale is not None:
+ sub_vocal_embeddings = torch.cat([torch.zeros_like(sub_vocal_embeddings), sub_vocal_embeddings, sub_vocal_embeddings], dim=0)
+ with torch.amp.autocast('cuda', dtype=weight_dtype):
+ legal_compressed_frames_num = latents.size()[2]
+ noise_pred = self.transformer(
+ x=latent_model_input,
+ context=prompt_embeds,
+ t=timestep,
+ seq_len=seq_len,
+ y=y[:, :, :legal_compressed_frames_num],
+ clip_fea=clip_context,
+ vocal_embeddings=sub_vocal_embeddings,
+ is_clip_level_modeling=False,
+ )
+ if do_classifier_free_guidance:
+ noise_pred_uncond, noise_pred_drop_audio, noise_pred_cond = noise_pred.chunk(3)
+ noise_pred = noise_pred_uncond + audio_guide_scale * (noise_pred_drop_audio - noise_pred_uncond) + text_guide_scale * (noise_pred_cond - noise_pred_drop_audio)
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+ torch.cuda.empty_cache()
+ if index_start != 0 and i != 0:
+ overlap_window_weight = torch.zeros(1, 1, overlap_window_length, 1, 1).to(device=latents.device, dtype=latents.dtype)
+ if overlapping_weight_scheme == "uniform":
+ for j in range(overlap_window_length):
+ overlap_window_weight[:, :, j] = j / (overlap_window_length-1)
+ elif overlapping_weight_scheme == "log":
+ init_weight = torch.linspace(0, 1, overlap_window_length)
+ init_weight = torch.log1p(init_weight * (torch.exp(torch.tensor(1.0)) - 1))
+ norm_weights = (init_weight - init_weight.min()) / (init_weight.max() - init_weight.min())
+ for j in range(overlap_window_length):
+ overlap_window_weight[:, :, j] = norm_weights[j]
+
+ overlap_idx_list_start = [ii % latents.shape[2] for ii in range(0, overlap_window_length)]
+ overlap_idx_list_end = [ii % latents_all.shape[2] for ii in range(index_previous_end-overlap_window_length, index_previous_end)]
+ latents[:, :, overlap_idx_list_start] = latents[:, :, overlap_idx_list_start] * overlap_window_weight + pred_latents[:, :, overlap_idx_list_end] * (1-overlap_window_weight)
+ latents = latents.to(torch.bfloat16)
+ for iii in range(legal_compressed_frames_num):
+ p = (index_start + iii) % pred_latents.shape[2]
+ pred_latents[:, :, p] = latents[:, :, iii]
+ else:
+ latents = latents.to(torch.bfloat16)
+ for iii in range(legal_compressed_frames_num):
+ p = (index_start + iii) % pred_latents.shape[2]
+ pred_latents[:, :, p] = latents[:, :, iii]
+ if arrive_last_frame:
+ break
+ if index_end != infer_length:
+ index_previous_end = index_end
+ index_start = index_start + (frames_per_batch-overlap_window_length)
+ if (index_start + frames_per_batch) < infer_length:
+ index_end = index_start + frames_per_batch
+ else:
+ index_end = infer_length
+ arrive_last_frame = True
+ latents_all = pred_latents
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+ progress_bar.update()
+ latents = latents_all.float()[:, :, :infer_length]
+ torch.cuda.empty_cache()
+ if output_type == "numpy":
+ video = self.decode_latents(latents)
+ elif not output_type == "latent":
+ video = self.decode_latents(latents)
+ video = self.video_processor.postprocess_video(video=video, output_type=output_type)
+ else:
+ video = latents
+ # Offload all models
+ self.maybe_free_model_hooks()
+ if not return_dict:
+ video = torch.from_numpy(video)
+ return WanI2VPipelineTalkingInferenceLongOutput(videos=video)
+
diff --git a/wan/pipeline/wan_inference_pipeline_fantasy.py b/wan/pipeline/wan_inference_pipeline_fantasy.py
new file mode 100644
index 0000000000000000000000000000000000000000..b89d4c6453be6cbb1a131903b32f79bdb4daa37c
--- /dev/null
+++ b/wan/pipeline/wan_inference_pipeline_fantasy.py
@@ -0,0 +1,741 @@
+import inspect
+import math
+from dataclasses import dataclass
+from typing import Any, Callable, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from diffusers import FlowMatchEulerDiscreteScheduler
+from diffusers.callbacks import MultiPipelineCallbacks, PipelineCallback
+from diffusers.image_processor import VaeImageProcessor
+from diffusers.models.embeddings import get_1d_rotary_pos_embed
+from diffusers.pipelines.pipeline_utils import DiffusionPipeline
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from diffusers.utils import BaseOutput, logging, replace_example_docstring
+from diffusers.utils.torch_utils import randn_tensor
+from diffusers.video_processor import VideoProcessor
+from einops import rearrange
+from PIL import Image
+from transformers import AutoTokenizer
+
+
+from wan.models.wan_image_encoder import CLIPModel
+from wan.models.wan_text_encoder import WanT5EncoderModel
+from wan.models.wan_transformer3d import WanTransformer3DModel
+from wan.models.wan_vae import AutoencoderKLWan
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+EXAMPLE_DOC_STRING = """
+ Examples:
+ ```python
+ pass
+ ```
+"""
+
+
+# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps: Optional[int] = None,
+ device: Optional[Union[str, torch.device]] = None,
+ timesteps: Optional[List[int]] = None,
+ sigmas: Optional[List[float]] = None,
+ **kwargs,
+):
+ """
+ Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles
+ custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`.
+
+ Args:
+ scheduler (`SchedulerMixin`):
+ The scheduler to get timesteps from.
+ num_inference_steps (`int`):
+ The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps`
+ must be `None`.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ timesteps (`List[int]`, *optional*):
+ Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed,
+ `num_inference_steps` and `sigmas` must be `None`.
+ sigmas (`List[float]`, *optional*):
+ Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed,
+ `num_inference_steps` and `timesteps` must be `None`.
+
+ Returns:
+ `Tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the
+ second element is the number of inference steps.
+ """
+ if timesteps is not None and sigmas is not None:
+ raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values")
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+def resize_mask(mask, latent, process_first_frame_only=True):
+ latent_size = latent.size()
+ batch_size, channels, num_frames, height, width = mask.shape
+
+ if process_first_frame_only:
+ target_size = list(latent_size[2:])
+ target_size[0] = 1
+ first_frame_resized = F.interpolate(
+ mask[:, :, 0:1, :, :],
+ size=target_size,
+ mode='trilinear',
+ align_corners=False
+ )
+
+ target_size = list(latent_size[2:])
+ target_size[0] = target_size[0] - 1
+ if target_size[0] != 0:
+ remaining_frames_resized = F.interpolate(
+ mask[:, :, 1:, :, :],
+ size=target_size,
+ mode='trilinear',
+ align_corners=False
+ )
+ resized_mask = torch.cat([first_frame_resized, remaining_frames_resized], dim=2)
+ else:
+ resized_mask = first_frame_resized
+ else:
+ target_size = list(latent_size[2:])
+ resized_mask = F.interpolate(
+ mask,
+ size=target_size,
+ mode='trilinear',
+ align_corners=False
+ )
+ return resized_mask
+
+
+@dataclass
+class WanI2VPipelineFantasyOutput(BaseOutput):
+ r"""
+ Output class for CogVideo pipelines.
+
+ Args:
+ video (`torch.Tensor`, `np.ndarray`, or List[List[PIL.Image.Image]]):
+ List of video outputs - It can be a nested list of length `batch_size,` with each sub-list containing
+ denoised PIL image sequences of length `num_frames.` It can also be a NumPy array or Torch tensor of shape
+ `(batch_size, num_frames, channels, height, width)`.
+ """
+
+ videos: torch.Tensor
+
+
+class WanI2VFantasyPipeline(DiffusionPipeline):
+ r"""
+ Pipeline for text-to-video generation using Wan.
+
+ This model inherits from [`DiffusionPipeline`]. Check the superclass documentation for the generic methods the
+ library implements for all the pipelines (such as downloading or saving, running on a particular device, etc.)
+ """
+
+ _optional_components = []
+ model_cpu_offload_seq = "text_encoder->clip_image_encoder->transformer->vae"
+
+ _callback_tensor_inputs = [
+ "latents",
+ "prompt_embeds",
+ "negative_prompt_embeds",
+ ]
+
+ def __init__(
+ self,
+ tokenizer: AutoTokenizer,
+ text_encoder: WanT5EncoderModel,
+ vae: AutoencoderKLWan,
+ transformer: WanTransformer3DModel,
+ clip_image_encoder: CLIPModel,
+ scheduler: FlowMatchEulerDiscreteScheduler,
+ ):
+ super().__init__()
+
+ self.register_modules(
+ tokenizer=tokenizer,
+ text_encoder=text_encoder,
+ vae=vae,
+ transformer=transformer,
+ clip_image_encoder=clip_image_encoder,
+ scheduler=scheduler,
+ )
+
+ self.video_processor = VideoProcessor(vae_scale_factor=self.vae.spacial_compression_ratio)
+ self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae.spacial_compression_ratio)
+ self.mask_processor = VaeImageProcessor(
+ vae_scale_factor=self.vae.spacial_compression_ratio, do_normalize=False, do_binarize=True,
+ do_convert_grayscale=True
+ )
+
+ def _get_t5_prompt_embeds(
+ self,
+ prompt: Union[str, List[str]] = None,
+ num_videos_per_prompt: int = 1,
+ max_sequence_length: int = 512,
+ device: Optional[torch.device] = None,
+ dtype: Optional[torch.dtype] = None,
+ ):
+ device = device or self._execution_device
+ dtype = dtype or self.text_encoder.dtype
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ batch_size = len(prompt)
+
+ text_inputs = self.tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=max_sequence_length,
+ truncation=True,
+ add_special_tokens=True,
+ return_tensors="pt",
+ )
+ text_input_ids = text_inputs.input_ids
+ prompt_attention_mask = text_inputs.attention_mask
+ untruncated_ids = self.tokenizer(prompt, padding="longest", return_tensors="pt").input_ids
+
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(text_input_ids, untruncated_ids):
+ removed_text = self.tokenizer.batch_decode(untruncated_ids[:, max_sequence_length - 1: -1])
+ logger.warning(
+ "The following part of your input was truncated because `max_sequence_length` is set to "
+ f" {max_sequence_length} tokens: {removed_text}"
+ )
+
+ seq_lens = prompt_attention_mask.gt(0).sum(dim=1).long()
+ prompt_embeds = self.text_encoder(text_input_ids.to(device), attention_mask=prompt_attention_mask.to(device))[0]
+ prompt_embeds = prompt_embeds.to(dtype=dtype, device=device)
+
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
+ _, seq_len, _ = prompt_embeds.shape
+ prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1)
+ prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1)
+
+ return [u[:v] for u, v in zip(prompt_embeds, seq_lens)]
+
+ def encode_prompt(
+ self,
+ prompt: Union[str, List[str]],
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ do_classifier_free_guidance: bool = True,
+ num_videos_per_prompt: int = 1,
+ prompt_embeds: Optional[torch.Tensor] = None,
+ negative_prompt_embeds: Optional[torch.Tensor] = None,
+ max_sequence_length: int = 512,
+ device: Optional[torch.device] = None,
+ dtype: Optional[torch.dtype] = None,
+ ):
+ r"""
+ Encodes the prompt into text encoder hidden states.
+
+ Args:
+ prompt (`str` or `List[str]`, *optional*):
+ prompt to be encoded
+ negative_prompt (`str` or `List[str]`, *optional*):
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
+ `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is
+ less than `1`).
+ do_classifier_free_guidance (`bool`, *optional*, defaults to `True`):
+ Whether to use classifier free guidance or not.
+ num_videos_per_prompt (`int`, *optional*, defaults to 1):
+ Number of videos that should be generated per prompt. torch device to place the resulting embeddings on
+ prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ negative_prompt_embeds (`torch.Tensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+ argument.
+ device: (`torch.device`, *optional*):
+ torch device
+ dtype: (`torch.dtype`, *optional*):
+ torch dtype
+ """
+ device = device or self._execution_device
+
+ prompt = [prompt] if isinstance(prompt, str) else prompt
+ if prompt is not None:
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ if prompt_embeds is None:
+ prompt_embeds = self._get_t5_prompt_embeds(
+ prompt=prompt,
+ num_videos_per_prompt=num_videos_per_prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ dtype=dtype,
+ )
+
+ if do_classifier_free_guidance and negative_prompt_embeds is None:
+ negative_prompt = negative_prompt or ""
+ negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+
+ if prompt is not None and type(prompt) is not type(negative_prompt):
+ raise TypeError(
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+ f" {type(prompt)}."
+ )
+ elif batch_size != len(negative_prompt):
+ raise ValueError(
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+ " the batch size of `prompt`."
+ )
+
+ negative_prompt_embeds = self._get_t5_prompt_embeds(
+ prompt=negative_prompt,
+ num_videos_per_prompt=num_videos_per_prompt,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ dtype=dtype,
+ )
+
+ return prompt_embeds, negative_prompt_embeds
+
+ def prepare_latents(
+ self, batch_size, num_channels_latents, num_frames, height, width, dtype, device, generator, latents=None
+ ):
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ shape = (
+ batch_size,
+ num_channels_latents,
+ (num_frames - 1) // self.vae.temporal_compression_ratio + 1,
+ height // self.vae.spacial_compression_ratio,
+ width // self.vae.spacial_compression_ratio,
+ )
+
+ if latents is None:
+ latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
+ else:
+ latents = latents.to(device)
+
+ # scale the initial noise by the standard deviation required by the scheduler
+ if hasattr(self.scheduler, "init_noise_sigma"):
+ latents = latents * self.scheduler.init_noise_sigma
+ return latents
+
+ def prepare_mask_latents(
+ self, mask, masked_image, batch_size, height, width, dtype, device, generator, do_classifier_free_guidance,
+ noise_aug_strength
+ ):
+ # resize the mask to latents shape as we concatenate the mask to the latents
+ # we do that before converting to dtype to avoid breaking in case we're using cpu_offload
+ # and half precision
+
+ if mask is not None:
+ mask = mask.to(device=device, dtype=self.vae.dtype)
+ bs = 1
+ new_mask = []
+ for i in range(0, mask.shape[0], bs):
+ mask_bs = mask[i: i + bs]
+ mask_bs = self.vae.encode(mask_bs)[0]
+ mask_bs = mask_bs.mode()
+ new_mask.append(mask_bs)
+ mask = torch.cat(new_mask, dim=0)
+ # mask = mask * self.vae.config.scaling_factor
+
+ if masked_image is not None:
+ masked_image = masked_image.to(device=device, dtype=self.vae.dtype)
+ bs = 1
+ new_mask_pixel_values = []
+ for i in range(0, masked_image.shape[0], bs):
+ mask_pixel_values_bs = masked_image[i: i + bs]
+ mask_pixel_values_bs = self.vae.encode(mask_pixel_values_bs)[0]
+ mask_pixel_values_bs = mask_pixel_values_bs.mode()
+ new_mask_pixel_values.append(mask_pixel_values_bs)
+ masked_image_latents = torch.cat(new_mask_pixel_values, dim=0)
+ # masked_image_latents = masked_image_latents * self.vae.config.scaling_factor
+ else:
+ masked_image_latents = None
+
+ return mask, masked_image_latents
+
+ def decode_latents(self, latents: torch.Tensor) -> torch.Tensor:
+ frames = self.vae.decode(latents.to(self.vae.dtype)).sample
+ frames = (frames / 2 + 0.5).clamp(0, 1)
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloa16
+ frames = frames.cpu().float().numpy()
+ return frames
+
+ # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.StableDiffusionPipeline.prepare_extra_step_kwargs
+ def prepare_extra_step_kwargs(self, generator, eta):
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+ # and should be between [0, 1]
+
+ accepts_eta = "eta" in set(inspect.signature(self.scheduler.step).parameters.keys())
+ extra_step_kwargs = {}
+ if accepts_eta:
+ extra_step_kwargs["eta"] = eta
+
+ # check if the scheduler accepts generator
+ accepts_generator = "generator" in set(inspect.signature(self.scheduler.step).parameters.keys())
+ if accepts_generator:
+ extra_step_kwargs["generator"] = generator
+ return extra_step_kwargs
+
+ # Copied from diffusers.pipelines.latte.pipeline_latte.LattePipeline.check_inputs
+ def check_inputs(
+ self,
+ prompt,
+ height,
+ width,
+ negative_prompt,
+ callback_on_step_end_tensor_inputs,
+ prompt_embeds=None,
+ negative_prompt_embeds=None,
+ ):
+ if height % 8 != 0 or width % 8 != 0:
+ raise ValueError(f"`height` and `width` have to be divisible by 8 but are {height} and {width}.")
+
+ if callback_on_step_end_tensor_inputs is not None and not all(
+ k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs
+ ):
+ raise ValueError(
+ f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}"
+ )
+ if prompt is not None and prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to"
+ " only forward one of the two."
+ )
+ elif prompt is None and prompt_embeds is None:
+ raise ValueError(
+ "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
+ )
+ elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)):
+ raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}")
+
+ if prompt is not None and negative_prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `prompt`: {prompt} and `negative_prompt_embeds`:"
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+ )
+
+ if negative_prompt is not None and negative_prompt_embeds is not None:
+ raise ValueError(
+ f"Cannot forward both `negative_prompt`: {negative_prompt} and `negative_prompt_embeds`:"
+ f" {negative_prompt_embeds}. Please make sure to only forward one of the two."
+ )
+
+ if prompt_embeds is not None and negative_prompt_embeds is not None:
+ if prompt_embeds.shape != negative_prompt_embeds.shape:
+ raise ValueError(
+ "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but"
+ f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`"
+ f" {negative_prompt_embeds.shape}."
+ )
+
+ @property
+ def guidance_scale(self):
+ return self._guidance_scale
+
+ @property
+ def num_timesteps(self):
+ return self._num_timesteps
+
+ @property
+ def attention_kwargs(self):
+ return self._attention_kwargs
+
+ @property
+ def interrupt(self):
+ return self._interrupt
+
+ @torch.no_grad()
+ @replace_example_docstring(EXAMPLE_DOC_STRING)
+ def __call__(
+ self,
+ prompt: Optional[Union[str, List[str]]] = None,
+ negative_prompt: Optional[Union[str, List[str]]] = None,
+ height: int = 480,
+ width: int = 720,
+ video: Union[torch.FloatTensor] = None,
+ mask_video: Union[torch.FloatTensor] = None,
+ num_frames: int = 49,
+ num_inference_steps: int = 50,
+ timesteps: Optional[List[int]] = None,
+ guidance_scale: float = 6,
+ num_videos_per_prompt: int = 1,
+ eta: float = 0.0,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ latents: Optional[torch.FloatTensor] = None,
+ prompt_embeds: Optional[torch.FloatTensor] = None,
+ negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+ output_type: str = "numpy",
+ return_dict: bool = False,
+ callback_on_step_end: Optional[
+ Union[Callable[[int, int, Dict], None], PipelineCallback, MultiPipelineCallbacks]
+ ] = None,
+ attention_kwargs: Optional[Dict[str, Any]] = None,
+ callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+ clip_image: Image = None,
+ max_sequence_length: int = 512,
+ prompt_cfg_scale=None,
+ audio_cfg_scale=None,
+ vocal_input_values=None,
+ ) -> Union[WanI2VPipelineFantasyOutput, Tuple]:
+ """
+ Function invoked when calling the pipeline for generation.
+ Args:
+
+ Examples:
+
+ Returns:
+
+ """
+
+ if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)):
+ callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs
+ num_videos_per_prompt = 1
+
+ # 1. Check inputs. Raise error if not correct
+ self.check_inputs(
+ prompt,
+ height,
+ width,
+ negative_prompt,
+ callback_on_step_end_tensor_inputs,
+ prompt_embeds,
+ negative_prompt_embeds,
+ )
+ self._guidance_scale = guidance_scale
+ self._attention_kwargs = attention_kwargs
+ self._interrupt = False
+
+ # 2. Default call parameters
+ if prompt is not None and isinstance(prompt, str):
+ batch_size = 1
+ elif prompt is not None and isinstance(prompt, list):
+ batch_size = len(prompt)
+ else:
+ batch_size = prompt_embeds.shape[0]
+
+ device = self._execution_device
+ weight_dtype = self.text_encoder.dtype
+
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+ # corresponds to doing no classifier free guidance.
+ do_classifier_free_guidance = guidance_scale > 1.0
+
+ # 3. Encode input prompt
+ prompt_embeds, negative_prompt_embeds = self.encode_prompt(
+ prompt,
+ negative_prompt,
+ do_classifier_free_guidance,
+ num_videos_per_prompt=num_videos_per_prompt,
+ prompt_embeds=prompt_embeds,
+ negative_prompt_embeds=negative_prompt_embeds,
+ max_sequence_length=max_sequence_length,
+ device=device,
+ )
+ if do_classifier_free_guidance:
+ prompt_embeds = negative_prompt_embeds + prompt_embeds + prompt_embeds
+
+ # 4. Prepare timesteps
+ if isinstance(self.scheduler, FlowMatchEulerDiscreteScheduler):
+ timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps,
+ mu=1)
+ else:
+ timesteps, num_inference_steps = retrieve_timesteps(self.scheduler, num_inference_steps, device, timesteps)
+ self._num_timesteps = len(timesteps)
+
+ # 5. Prepare latents.
+ if video is not None:
+ video_length = video.shape[2]
+ init_video = self.image_processor.preprocess(rearrange(video, "b c f h w -> (b f) c h w"), height=height, width=width)
+ init_video = init_video.to(dtype=torch.float32)
+ init_video = rearrange(init_video, "(b f) c h w -> b c f h w", f=video_length)
+ else:
+ init_video = None
+
+ latent_channels = self.vae.config.latent_channels
+ latents = self.prepare_latents(
+ batch_size * num_videos_per_prompt,
+ latent_channels,
+ num_frames,
+ height,
+ width,
+ weight_dtype,
+ device,
+ generator,
+ latents,
+ )
+
+
+ # Prepare mask latent variables
+ if init_video is not None:
+ if (mask_video == 255).all():
+ mask_latents = torch.tile(
+ torch.zeros_like(latents)[:, :1].to(device, weight_dtype), [1, 4, 1, 1, 1]
+ )
+ masked_video_latents = torch.zeros_like(latents).to(device, weight_dtype)
+
+ mask_input = torch.cat([mask_latents] * 3) if do_classifier_free_guidance else mask_latents
+ masked_video_latents_input = (
+ torch.cat([masked_video_latents] * 3) if do_classifier_free_guidance else masked_video_latents
+ )
+ y = torch.cat([mask_input, masked_video_latents_input], dim=1).to(device, weight_dtype)
+ else:
+ bs, _, video_length, height, width = video.size()
+ mask_condition = self.mask_processor.preprocess(rearrange(mask_video, "b c f h w -> (b f) c h w"),
+ height=height, width=width)
+ mask_condition = mask_condition.to(dtype=torch.float32)
+ mask_condition = rearrange(mask_condition, "(b f) c h w -> b c f h w", f=video_length)
+
+ masked_video = init_video * (torch.tile(mask_condition, [1, 3, 1, 1, 1]) < 0.5)
+ _, masked_video_latents = self.prepare_mask_latents(
+ None,
+ masked_video,
+ batch_size,
+ height,
+ width,
+ weight_dtype,
+ device,
+ generator,
+ do_classifier_free_guidance,
+ noise_aug_strength=None,
+ )
+
+ mask_condition = torch.concat(
+ [
+ torch.repeat_interleave(mask_condition[:, :, 0:1], repeats=4, dim=2),
+ mask_condition[:, :, 1:]
+ ], dim=2
+ )
+ mask_condition = mask_condition.view(bs, mask_condition.shape[2] // 4, 4, height, width)
+ mask_condition = mask_condition.transpose(1, 2)
+ mask_latents = resize_mask(1 - mask_condition, masked_video_latents, True).to(device, weight_dtype)
+
+ mask_input = torch.cat([mask_latents] * 3) if do_classifier_free_guidance else mask_latents
+ masked_video_latents_input = (
+ torch.cat([masked_video_latents] * 3) if do_classifier_free_guidance else masked_video_latents
+ )
+
+ y = torch.cat([mask_input, masked_video_latents_input], dim=1).to(device, weight_dtype)
+
+ # Prepare clip latent variables
+ if clip_image is not None:
+ clip_image = TF.to_tensor(clip_image).sub_(0.5).div_(0.5).to(device, weight_dtype)
+ clip_context = self.clip_image_encoder([clip_image[:, None, :, :]])
+ clip_context = (
+ torch.cat([clip_context, clip_context, clip_context], dim=0) if do_classifier_free_guidance else clip_context
+ )
+ else:
+ clip_image = Image.new("RGB", (512, 512), color=(0, 0, 0))
+ clip_image = TF.to_tensor(clip_image).sub_(0.5).div_(0.5).to(device, weight_dtype)
+ clip_context = self.clip_image_encoder([clip_image[:, None, :, :]])
+ clip_context = (
+ torch.cat([clip_context, clip_context, clip_context], dim=0) if do_classifier_free_guidance else clip_context
+ )
+ clip_context = torch.zeros_like(clip_context)
+
+
+ # 6. Prepare extra step kwargs. TODO: Logic should ideally just be moved out of the pipeline
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+ target_shape = (self.vae.latent_channels, (num_frames - 1) // self.vae.temporal_compression_ratio + 1,
+ width // self.vae.spacial_compression_ratio, height // self.vae.spacial_compression_ratio)
+ seq_len = math.ceil((target_shape[2] * target_shape[3]) / (
+ self.transformer.config.patch_size[1] * self.transformer.config.patch_size[2]) * target_shape[1])
+
+ if prompt_cfg_scale is not None and audio_cfg_scale is not None:
+ vocal_input_values = torch.cat([torch.zeros_like(vocal_input_values), torch.zeros_like(vocal_input_values), vocal_input_values], dim=0)
+
+ # 7. Denoising loop
+ num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+ if self.interrupt:
+ continue
+
+ latent_model_input = torch.cat([latents] * 3) if do_classifier_free_guidance else latents
+ if hasattr(self.scheduler, "scale_model_input"):
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+ # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+ timestep = t.expand(latent_model_input.shape[0])
+
+ # predict noise model_output
+ with torch.cuda.amp.autocast(dtype=weight_dtype):
+
+ noise_pred = self.transformer(
+ x=latent_model_input,
+ context=prompt_embeds,
+ t=timestep,
+ seq_len=seq_len,
+ y=y,
+ clip_fea=clip_context,
+ vocal_embeddings=vocal_input_values,
+ is_clip_level_modeling=False,
+ )
+
+ # perform guidance
+ if do_classifier_free_guidance:
+ noise_pred_uncond, noise_pred_drop_audio, noise_pred_cond = noise_pred.chunk(3)
+ noise_pred = noise_pred_uncond + prompt_cfg_scale * (noise_pred_drop_audio - noise_pred_uncond) + audio_cfg_scale * (noise_pred_cond - noise_pred_drop_audio)
+
+ # compute the previous noisy sample x_t -> x_t-1
+ latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+
+ if callback_on_step_end is not None:
+ callback_kwargs = {}
+ for k in callback_on_step_end_tensor_inputs:
+ callback_kwargs[k] = locals()[k]
+ callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+ latents = callback_outputs.pop("latents", latents)
+ prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+ negative_prompt_embeds = callback_outputs.pop("negative_prompt_embeds", negative_prompt_embeds)
+
+ if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+ progress_bar.update()
+
+
+ if output_type == "numpy":
+ video = self.decode_latents(latents)
+ elif not output_type == "latent":
+ video = self.decode_latents(latents)
+ video = self.video_processor.postprocess_video(video=video, output_type=output_type)
+ else:
+ video = latents
+
+ # Offload all models
+ self.maybe_free_model_hooks()
+
+ if not return_dict:
+ video = torch.from_numpy(video)
+
+ return WanI2VPipelineFantasyOutput(videos=video)
diff --git a/wan/text2video.py b/wan/text2video.py
new file mode 100644
index 0000000000000000000000000000000000000000..24005450523e8b7025b934ee6a00f686263e5645
--- /dev/null
+++ b/wan/text2video.py
@@ -0,0 +1,267 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import gc
+import logging
+import math
+import os
+import random
+import sys
+import types
+from contextlib import contextmanager
+from functools import partial
+
+import torch
+import torch.cuda.amp as amp
+import torch.distributed as dist
+from tqdm import tqdm
+
+from .distributed.fsdp import shard_model
+from .modules.model import WanModel
+from .modules.t5 import T5EncoderModel
+from .modules.vae import WanVAE
+from .utils.fm_solvers import (FlowDPMSolverMultistepScheduler,
+ get_sampling_sigmas, retrieve_timesteps)
+from .utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+
+
+class WanT2V:
+
+ def __init__(
+ self,
+ config,
+ checkpoint_dir,
+ device_id=0,
+ rank=0,
+ t5_fsdp=False,
+ dit_fsdp=False,
+ use_usp=False,
+ t5_cpu=False,
+ ):
+ r"""
+ Initializes the Wan text-to-video generation model components.
+
+ Args:
+ config (EasyDict):
+ Object containing model parameters initialized from config.py
+ checkpoint_dir (`str`):
+ Path to directory containing model checkpoints
+ device_id (`int`, *optional*, defaults to 0):
+ Id of target GPU device
+ rank (`int`, *optional*, defaults to 0):
+ Process rank for distributed training
+ t5_fsdp (`bool`, *optional*, defaults to False):
+ Enable FSDP sharding for T5 model
+ dit_fsdp (`bool`, *optional*, defaults to False):
+ Enable FSDP sharding for DiT model
+ use_usp (`bool`, *optional*, defaults to False):
+ Enable distribution strategy of USP.
+ t5_cpu (`bool`, *optional*, defaults to False):
+ Whether to place T5 model on CPU. Only works without t5_fsdp.
+ """
+ self.device = torch.device(f"cuda:{device_id}")
+ self.config = config
+ self.rank = rank
+ self.t5_cpu = t5_cpu
+
+ self.num_train_timesteps = config.num_train_timesteps
+ self.param_dtype = config.param_dtype
+
+ shard_fn = partial(shard_model, device_id=device_id)
+ self.text_encoder = T5EncoderModel(
+ text_len=config.text_len,
+ dtype=config.t5_dtype,
+ device=torch.device('cpu'),
+ checkpoint_path=os.path.join(checkpoint_dir, config.t5_checkpoint),
+ tokenizer_path=os.path.join(checkpoint_dir, config.t5_tokenizer),
+ shard_fn=shard_fn if t5_fsdp else None)
+
+ self.vae_stride = config.vae_stride
+ self.patch_size = config.patch_size
+ self.vae = WanVAE(
+ vae_pth=os.path.join(checkpoint_dir, config.vae_checkpoint),
+ device=self.device)
+
+ logging.info(f"Creating WanModel from {checkpoint_dir}")
+ self.model = WanModel.from_pretrained(checkpoint_dir)
+ self.model.eval().requires_grad_(False)
+
+ if use_usp:
+ from xfuser.core.distributed import \
+ get_sequence_parallel_world_size
+
+ from .distributed.xdit_context_parallel import (usp_attn_forward,
+ usp_dit_forward)
+ for block in self.model.blocks:
+ block.self_attn.forward = types.MethodType(
+ usp_attn_forward, block.self_attn)
+ self.model.forward = types.MethodType(usp_dit_forward, self.model)
+ self.sp_size = get_sequence_parallel_world_size()
+ else:
+ self.sp_size = 1
+
+ if dist.is_initialized():
+ dist.barrier()
+ if dit_fsdp:
+ self.model = shard_fn(self.model)
+ else:
+ self.model.to(self.device)
+
+ self.sample_neg_prompt = config.sample_neg_prompt
+
+ def generate(self,
+ input_prompt,
+ size=(1280, 720),
+ frame_num=81,
+ shift=5.0,
+ sample_solver='unipc',
+ sampling_steps=50,
+ guide_scale=5.0,
+ n_prompt="",
+ seed=-1,
+ offload_model=True):
+ r"""
+ Generates video frames from text prompt using diffusion process.
+
+ Args:
+ input_prompt (`str`):
+ Text prompt for content generation
+ size (tupele[`int`], *optional*, defaults to (1280,720)):
+ Controls video resolution, (width,height).
+ frame_num (`int`, *optional*, defaults to 81):
+ How many frames to sample from a video. The number should be 4n+1
+ shift (`float`, *optional*, defaults to 5.0):
+ Noise schedule shift parameter. Affects temporal dynamics
+ sample_solver (`str`, *optional*, defaults to 'unipc'):
+ Solver used to sample the video.
+ sampling_steps (`int`, *optional*, defaults to 40):
+ Number of diffusion sampling steps. Higher values improve quality but slow generation
+ guide_scale (`float`, *optional*, defaults 5.0):
+ Classifier-free guidance scale. Controls prompt adherence vs. creativity
+ n_prompt (`str`, *optional*, defaults to ""):
+ Negative prompt for content exclusion. If not given, use `config.sample_neg_prompt`
+ seed (`int`, *optional*, defaults to -1):
+ Random seed for noise generation. If -1, use random seed.
+ offload_model (`bool`, *optional*, defaults to True):
+ If True, offloads models to CPU during generation to save VRAM
+
+ Returns:
+ torch.Tensor:
+ Generated video frames tensor. Dimensions: (C, N H, W) where:
+ - C: Color channels (3 for RGB)
+ - N: Number of frames (81)
+ - H: Frame height (from size)
+ - W: Frame width from size)
+ """
+ # preprocess
+ F = frame_num
+ target_shape = (self.vae.model.z_dim, (F - 1) // self.vae_stride[0] + 1,
+ size[1] // self.vae_stride[1],
+ size[0] // self.vae_stride[2])
+
+ seq_len = math.ceil((target_shape[2] * target_shape[3]) /
+ (self.patch_size[1] * self.patch_size[2]) *
+ target_shape[1] / self.sp_size) * self.sp_size
+
+ if n_prompt == "":
+ n_prompt = self.sample_neg_prompt
+ seed = seed if seed >= 0 else random.randint(0, sys.maxsize)
+ seed_g = torch.Generator(device=self.device)
+ seed_g.manual_seed(seed)
+
+ if not self.t5_cpu:
+ self.text_encoder.model.to(self.device)
+ context = self.text_encoder([input_prompt], self.device)
+ context_null = self.text_encoder([n_prompt], self.device)
+ if offload_model:
+ self.text_encoder.model.cpu()
+ else:
+ context = self.text_encoder([input_prompt], torch.device('cpu'))
+ context_null = self.text_encoder([n_prompt], torch.device('cpu'))
+ context = [t.to(self.device) for t in context]
+ context_null = [t.to(self.device) for t in context_null]
+
+ noise = [
+ torch.randn(
+ target_shape[0],
+ target_shape[1],
+ target_shape[2],
+ target_shape[3],
+ dtype=torch.float32,
+ device=self.device,
+ generator=seed_g)
+ ]
+
+ @contextmanager
+ def noop_no_sync():
+ yield
+
+ no_sync = getattr(self.model, 'no_sync', noop_no_sync)
+
+ # evaluation mode
+ with amp.autocast(dtype=self.param_dtype), torch.no_grad(), no_sync():
+
+ if sample_solver == 'unipc':
+ sample_scheduler = FlowUniPCMultistepScheduler(
+ num_train_timesteps=self.num_train_timesteps,
+ shift=1,
+ use_dynamic_shifting=False)
+ sample_scheduler.set_timesteps(
+ sampling_steps, device=self.device, shift=shift)
+ timesteps = sample_scheduler.timesteps
+ elif sample_solver == 'dpm++':
+ sample_scheduler = FlowDPMSolverMultistepScheduler(
+ num_train_timesteps=self.num_train_timesteps,
+ shift=1,
+ use_dynamic_shifting=False)
+ sampling_sigmas = get_sampling_sigmas(sampling_steps, shift)
+ timesteps, _ = retrieve_timesteps(
+ sample_scheduler,
+ device=self.device,
+ sigmas=sampling_sigmas)
+ else:
+ raise NotImplementedError("Unsupported solver.")
+
+ # sample videos
+ latents = noise
+
+ arg_c = {'context': context, 'seq_len': seq_len}
+ arg_null = {'context': context_null, 'seq_len': seq_len}
+
+ for _, t in enumerate(tqdm(timesteps)):
+ latent_model_input = latents
+ timestep = [t]
+
+ timestep = torch.stack(timestep)
+
+ self.model.to(self.device)
+ noise_pred_cond = self.model(
+ latent_model_input, t=timestep, **arg_c)[0]
+ noise_pred_uncond = self.model(
+ latent_model_input, t=timestep, **arg_null)[0]
+
+ noise_pred = noise_pred_uncond + guide_scale * (
+ noise_pred_cond - noise_pred_uncond)
+
+ temp_x0 = sample_scheduler.step(
+ noise_pred.unsqueeze(0),
+ t,
+ latents[0].unsqueeze(0),
+ return_dict=False,
+ generator=seed_g)[0]
+ latents = [temp_x0.squeeze(0)]
+
+ x0 = latents
+ if offload_model:
+ self.model.cpu()
+ torch.cuda.empty_cache()
+ if self.rank == 0:
+ videos = self.vae.decode(x0)
+
+ del noise, latents
+ del sample_scheduler
+ if offload_model:
+ gc.collect()
+ torch.cuda.synchronize()
+ if dist.is_initialized():
+ dist.barrier()
+
+ return videos[0] if self.rank == 0 else None
diff --git a/wan/utils/__init__.py b/wan/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6e9a339e69fd55dd226d3ce242613c19bd690522
--- /dev/null
+++ b/wan/utils/__init__.py
@@ -0,0 +1,8 @@
+from .fm_solvers import (FlowDPMSolverMultistepScheduler, get_sampling_sigmas,
+ retrieve_timesteps)
+from .fm_solvers_unipc import FlowUniPCMultistepScheduler
+
+__all__ = [
+ 'HuggingfaceTokenizer', 'get_sampling_sigmas', 'retrieve_timesteps',
+ 'FlowDPMSolverMultistepScheduler', 'FlowUniPCMultistepScheduler'
+]
diff --git a/wan/utils/__pycache__/__init__.cpython-311.pyc b/wan/utils/__pycache__/__init__.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..85fab3ad9639a8e0952a0abf46f71a071f4b8a81
Binary files /dev/null and b/wan/utils/__pycache__/__init__.cpython-311.pyc differ
diff --git a/wan/utils/__pycache__/color_correction.cpython-311.pyc b/wan/utils/__pycache__/color_correction.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..30d0e23139feb8790c4a2519e83b32acc1eec6bc
Binary files /dev/null and b/wan/utils/__pycache__/color_correction.cpython-311.pyc differ
diff --git a/wan/utils/__pycache__/discrete_sampler.cpython-311.pyc b/wan/utils/__pycache__/discrete_sampler.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..dc21ad6149e49d06bd28591ad25b853a475df72d
Binary files /dev/null and b/wan/utils/__pycache__/discrete_sampler.cpython-311.pyc differ
diff --git a/wan/utils/__pycache__/fm_solvers.cpython-311.pyc b/wan/utils/__pycache__/fm_solvers.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..c5a21d1d9122443561d3840ba86f3d23a88026f7
Binary files /dev/null and b/wan/utils/__pycache__/fm_solvers.cpython-311.pyc differ
diff --git a/wan/utils/__pycache__/fm_solvers_unipc.cpython-311.pyc b/wan/utils/__pycache__/fm_solvers_unipc.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..52cf1b425550f71eb34cf907d3b9f185d781e3c1
Binary files /dev/null and b/wan/utils/__pycache__/fm_solvers_unipc.cpython-311.pyc differ
diff --git a/wan/utils/__pycache__/fp8_optimization.cpython-311.pyc b/wan/utils/__pycache__/fp8_optimization.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..33a9f657033d8c608a9b0b4530f9f0bf6d2ac770
Binary files /dev/null and b/wan/utils/__pycache__/fp8_optimization.cpython-311.pyc differ
diff --git a/wan/utils/__pycache__/utils.cpython-311.pyc b/wan/utils/__pycache__/utils.cpython-311.pyc
new file mode 100644
index 0000000000000000000000000000000000000000..65afdde8939060b5e0f9b4968e9f1a77eb98a508
Binary files /dev/null and b/wan/utils/__pycache__/utils.cpython-311.pyc differ
diff --git a/wan/utils/color_correction.py b/wan/utils/color_correction.py
new file mode 100644
index 0000000000000000000000000000000000000000..83bef44285640afc073706a3a4df202e3821d3f1
--- /dev/null
+++ b/wan/utils/color_correction.py
@@ -0,0 +1,108 @@
+import torch
+import numpy as np
+from skimage import color
+
+
+def match_and_blend_colors(source_chunk: torch.Tensor, reference_image: torch.Tensor, strength: float) -> torch.Tensor:
+ """
+ Matches the color of a source video chunk to a reference image and blends with the original.
+
+ Args:
+ source_chunk (torch.Tensor): The video chunk to be color-corrected (B, C, T, H, W) in range [-1, 1].
+ Assumes B=1 (batch size of 1).
+ reference_image (torch.Tensor): The reference image (B, C, 1, H, W) in range [-1, 1].
+ Assumes B=1 and T=1 (single reference frame).
+ strength (float): The strength of the color correction (0.0 to 1.0).
+ 0.0 means no correction, 1.0 means full correction.
+
+ Returns:
+ torch.Tensor: The color-corrected and blended video chunk.
+ """
+ # print(f"[match_and_blend_colors] Input source_chunk shape: {source_chunk.shape}, reference_image shape: {reference_image.shape}, strength: {strength}")
+
+ if strength == 0.0:
+ # print(f"[match_and_blend_colors] Strength is 0, returning original source_chunk.")
+ return source_chunk
+
+ if not 0.0 <= strength <= 1.0:
+ raise ValueError(f"Strength must be between 0.0 and 1.0, got {strength}")
+
+ device = source_chunk.device
+ dtype = source_chunk.dtype
+
+ # Squeeze batch dimension, permute to T, H, W, C for skimage
+ # Source: (1, C, T, H, W) -> (T, H, W, C)
+ source_np = source_chunk.squeeze(0).permute(1, 2, 3, 0).cpu().numpy()
+ # Reference: (1, C, 1, H, W) -> (H, W, C)
+ ref_np = reference_image.squeeze(0).squeeze(1).permute(1, 2, 0).cpu().numpy() # Squeeze T dimension as well
+
+ # Normalize from [-1, 1] to [0, 1] for skimage
+ source_np_01 = (source_np + 1.0) / 2.0
+ ref_np_01 = (ref_np + 1.0) / 2.0
+
+ # Clip to ensure values are strictly in [0, 1] after potential float precision issues
+ source_np_01 = np.clip(source_np_01, 0.0, 1.0)
+ ref_np_01 = np.clip(ref_np_01, 0.0, 1.0)
+
+ # Convert reference to Lab
+ try:
+ ref_lab = color.rgb2lab(ref_np_01)
+ except ValueError as e:
+ # Handle potential errors if image data is not valid for conversion
+ print(f"Warning: Could not convert reference image to Lab: {e}. Skipping color correction for this chunk.")
+ return source_chunk
+
+ corrected_frames_np_01 = []
+ for i in range(source_np_01.shape[0]): # Iterate over time (T)
+ source_frame_rgb_01 = source_np_01[i]
+
+ try:
+ source_lab = color.rgb2lab(source_frame_rgb_01)
+ except ValueError as e:
+ print(f"Warning: Could not convert source frame {i} to Lab: {e}. Using original frame.")
+ corrected_frames_np_01.append(source_frame_rgb_01)
+ continue
+
+ corrected_lab_frame = source_lab.copy()
+
+ # Perform color transfer for L, a, b channels
+ for j in range(3): # L, a, b
+ mean_src, std_src = source_lab[:, :, j].mean(), source_lab[:, :, j].std()
+ mean_ref, std_ref = ref_lab[:, :, j].mean(), ref_lab[:, :, j].std()
+
+ # Avoid division by zero if std_src is 0
+ if std_src == 0:
+ # If source channel has no variation, keep it as is, but shift by reference mean
+ # This case is debatable, could also just copy source or target mean.
+ # Shifting by target mean helps if source is flat but target isn't.
+ corrected_lab_frame[:, :, j] = mean_ref
+ else:
+ corrected_lab_frame[:, :, j] = (corrected_lab_frame[:, :, j] - mean_src) * (
+ std_ref / std_src) + mean_ref
+
+ try:
+ fully_corrected_frame_rgb_01 = color.lab2rgb(corrected_lab_frame)
+ except ValueError as e:
+ print(f"Warning: Could not convert corrected frame {i} back to RGB: {e}. Using original frame.")
+ corrected_frames_np_01.append(source_frame_rgb_01)
+ continue
+
+ # Clip again after lab2rgb as it can go slightly out of [0,1]
+ fully_corrected_frame_rgb_01 = np.clip(fully_corrected_frame_rgb_01, 0.0, 1.0)
+
+ # Blend with original source frame (in [0,1] RGB)
+ blended_frame_rgb_01 = (1 - strength) * source_frame_rgb_01 + strength * fully_corrected_frame_rgb_01
+ corrected_frames_np_01.append(blended_frame_rgb_01)
+
+ corrected_chunk_np_01 = np.stack(corrected_frames_np_01, axis=0)
+
+ # Convert back to [-1, 1]
+ corrected_chunk_np_minus1_1 = (corrected_chunk_np_01 * 2.0) - 1.0
+
+ # Permute back to (C, T, H, W), add batch dim, and convert to original torch.Tensor type and device
+ # (T, H, W, C) -> (C, T, H, W)
+ corrected_chunk_tensor = torch.from_numpy(corrected_chunk_np_minus1_1).permute(3, 0, 1, 2).unsqueeze(0)
+ corrected_chunk_tensor = corrected_chunk_tensor.contiguous() # Ensure contiguous memory layout
+ output_tensor = corrected_chunk_tensor.to(device=device, dtype=dtype)
+ # print(f"[match_and_blend_colors] Output tensor shape: {output_tensor.shape}")
+ return output_tensor
\ No newline at end of file
diff --git a/wan/utils/discrete_sampler.py b/wan/utils/discrete_sampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..47f3557ac6923840ef52660db5050b5e86a49c11
--- /dev/null
+++ b/wan/utils/discrete_sampler.py
@@ -0,0 +1,47 @@
+"""Modified from https://github.com/THUDM/CogVideo/blob/3710a612d8760f5cdb1741befeebb65b9e0f2fe0/sat/sgm/modules/diffusionmodules/sigma_sampling.py
+"""
+import torch
+
+
+class DiscreteSampling:
+ def __init__(self, num_idx, uniform_sampling=False):
+ self.num_idx = num_idx
+ self.uniform_sampling = uniform_sampling
+ self.is_distributed = torch.distributed.is_available() and torch.distributed.is_initialized()
+
+ if self.is_distributed and self.uniform_sampling:
+ world_size = torch.distributed.get_world_size()
+ self.rank = torch.distributed.get_rank()
+
+ i = 1
+ while True:
+ if world_size % i != 0 or num_idx % (world_size // i) != 0:
+ i += 1
+ else:
+ self.group_num = world_size // i
+ break
+ assert self.group_num > 0
+ assert world_size % self.group_num == 0
+ # the number of rank in one group
+ self.group_width = world_size // self.group_num
+ self.sigma_interval = self.num_idx // self.group_num
+ print('rank=%d world_size=%d group_num=%d group_width=%d sigma_interval=%s' % (
+ self.rank, world_size, self.group_num,
+ self.group_width, self.sigma_interval))
+
+ def __call__(self, n_samples, generator=None, device=None):
+ if self.is_distributed and self.uniform_sampling:
+ group_index = self.rank // self.group_width
+ idx = torch.randint(
+ group_index * self.sigma_interval,
+ (group_index + 1) * self.sigma_interval,
+ (n_samples,),
+ generator=generator, device=device,
+ )
+ print('proc[%d] idx=%s' % (self.rank, idx))
+ else:
+ idx = torch.randint(
+ 0, self.num_idx, (n_samples,),
+ generator=generator, device=device,
+ )
+ return idx
\ No newline at end of file
diff --git a/wan/utils/fm_solvers.py b/wan/utils/fm_solvers.py
new file mode 100644
index 0000000000000000000000000000000000000000..c908969e24849ce1381a8df9d5eb401dccf66524
--- /dev/null
+++ b/wan/utils/fm_solvers.py
@@ -0,0 +1,857 @@
+# Copied from https://github.com/huggingface/diffusers/blob/main/src/diffusers/schedulers/scheduling_dpmsolver_multistep.py
+# Convert dpm solver for flow matching
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+
+import inspect
+import math
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers,
+ SchedulerMixin,
+ SchedulerOutput)
+from diffusers.utils import deprecate, is_scipy_available
+from diffusers.utils.torch_utils import randn_tensor
+
+if is_scipy_available():
+ pass
+
+
+def get_sampling_sigmas(sampling_steps, shift):
+ sigma = np.linspace(1, 0, sampling_steps + 1)[:sampling_steps]
+ sigma = (shift * sigma / (1 + (shift - 1) * sigma))
+
+ return sigma
+
+
+def retrieve_timesteps(
+ scheduler,
+ num_inference_steps=None,
+ device=None,
+ timesteps=None,
+ sigmas=None,
+ **kwargs,
+):
+ if timesteps is not None and sigmas is not None:
+ raise ValueError(
+ "Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values"
+ )
+ if timesteps is not None:
+ accepts_timesteps = "timesteps" in set(
+ inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accepts_timesteps:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" timestep schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ elif sigmas is not None:
+ accept_sigmas = "sigmas" in set(
+ inspect.signature(scheduler.set_timesteps).parameters.keys())
+ if not accept_sigmas:
+ raise ValueError(
+ f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom"
+ f" sigmas schedules. Please check whether you are using the correct scheduler."
+ )
+ scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ num_inference_steps = len(timesteps)
+ else:
+ scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+ timesteps = scheduler.timesteps
+ return timesteps, num_inference_steps
+
+
+class FlowDPMSolverMultistepScheduler(SchedulerMixin, ConfigMixin):
+ """
+ `FlowDPMSolverMultistepScheduler` is a fast dedicated high-order solver for diffusion ODEs.
+ This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+ methods the library implements for all schedulers such as loading and saving.
+ Args:
+ num_train_timesteps (`int`, defaults to 1000):
+ The number of diffusion steps to train the model. This determines the resolution of the diffusion process.
+ solver_order (`int`, defaults to 2):
+ The DPMSolver order which can be `1`, `2`, or `3`. It is recommended to use `solver_order=2` for guided
+ sampling, and `solver_order=3` for unconditional sampling. This affects the number of model outputs stored
+ and used in multistep updates.
+ prediction_type (`str`, defaults to "flow_prediction"):
+ Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts
+ the flow of the diffusion process.
+ shift (`float`, *optional*, defaults to 1.0):
+ A factor used to adjust the sigmas in the noise schedule. It modifies the step sizes during the sampling
+ process.
+ use_dynamic_shifting (`bool`, defaults to `False`):
+ Whether to apply dynamic shifting to the timesteps based on image resolution. If `True`, the shifting is
+ applied on the fly.
+ thresholding (`bool`, defaults to `False`):
+ Whether to use the "dynamic thresholding" method. This method adjusts the predicted sample to prevent
+ saturation and improve photorealism.
+ dynamic_thresholding_ratio (`float`, defaults to 0.995):
+ The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+ sample_max_value (`float`, defaults to 1.0):
+ The threshold value for dynamic thresholding. Valid only when `thresholding=True` and
+ `algorithm_type="dpmsolver++"`.
+ algorithm_type (`str`, defaults to `dpmsolver++`):
+ Algorithm type for the solver; can be `dpmsolver`, `dpmsolver++`, `sde-dpmsolver` or `sde-dpmsolver++`. The
+ `dpmsolver` type implements the algorithms in the [DPMSolver](https://huggingface.co/papers/2206.00927)
+ paper, and the `dpmsolver++` type implements the algorithms in the
+ [DPMSolver++](https://huggingface.co/papers/2211.01095) paper. It is recommended to use `dpmsolver++` or
+ `sde-dpmsolver++` with `solver_order=2` for guided sampling like in Stable Diffusion.
+ solver_type (`str`, defaults to `midpoint`):
+ Solver type for the second-order solver; can be `midpoint` or `heun`. The solver type slightly affects the
+ sample quality, especially for a small number of steps. It is recommended to use `midpoint` solvers.
+ lower_order_final (`bool`, defaults to `True`):
+ Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
+ stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
+ euler_at_final (`bool`, defaults to `False`):
+ Whether to use Euler's method in the final step. It is a trade-off between numerical stability and detail
+ richness. This can stabilize the sampling of the SDE variant of DPMSolver for small number of inference
+ steps, but sometimes may result in blurring.
+ final_sigmas_type (`str`, *optional*, defaults to "zero"):
+ The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
+ sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
+ lambda_min_clipped (`float`, defaults to `-inf`):
+ Clipping threshold for the minimum value of `lambda(t)` for numerical stability. This is critical for the
+ cosine (`squaredcos_cap_v2`) noise schedule.
+ variance_type (`str`, *optional*):
+ Set to "learned" or "learned_range" for diffusion models that predict variance. If set, the model's output
+ contains the predicted Gaussian variance.
+ """
+
+ _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+ order = 1
+
+ @register_to_config
+ def __init__(
+ self,
+ num_train_timesteps: int = 1000,
+ solver_order: int = 2,
+ prediction_type: str = "flow_prediction",
+ shift: Optional[float] = 1.0,
+ use_dynamic_shifting=False,
+ thresholding: bool = False,
+ dynamic_thresholding_ratio: float = 0.995,
+ sample_max_value: float = 1.0,
+ algorithm_type: str = "dpmsolver++",
+ solver_type: str = "midpoint",
+ lower_order_final: bool = True,
+ euler_at_final: bool = False,
+ final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min"
+ lambda_min_clipped: float = -float("inf"),
+ variance_type: Optional[str] = None,
+ invert_sigmas: bool = False,
+ ):
+ if algorithm_type in ["dpmsolver", "sde-dpmsolver"]:
+ deprecation_message = f"algorithm_type {algorithm_type} is deprecated and will be removed in a future version. Choose from `dpmsolver++` or `sde-dpmsolver++` instead"
+ deprecate("algorithm_types dpmsolver and sde-dpmsolver", "1.0.0",
+ deprecation_message)
+
+ # settings for DPM-Solver
+ if algorithm_type not in [
+ "dpmsolver", "dpmsolver++", "sde-dpmsolver", "sde-dpmsolver++"
+ ]:
+ if algorithm_type == "deis":
+ self.register_to_config(algorithm_type="dpmsolver++")
+ else:
+ raise NotImplementedError(
+ f"{algorithm_type} is not implemented for {self.__class__}")
+
+ if solver_type not in ["midpoint", "heun"]:
+ if solver_type in ["logrho", "bh1", "bh2"]:
+ self.register_to_config(solver_type="midpoint")
+ else:
+ raise NotImplementedError(
+ f"{solver_type} is not implemented for {self.__class__}")
+
+ if algorithm_type not in ["dpmsolver++", "sde-dpmsolver++"
+ ] and final_sigmas_type == "zero":
+ raise ValueError(
+ f"`final_sigmas_type` {final_sigmas_type} is not supported for `algorithm_type` {algorithm_type}. Please choose `sigma_min` instead."
+ )
+
+ # setable values
+ self.num_inference_steps = None
+ alphas = np.linspace(1, 1 / num_train_timesteps,
+ num_train_timesteps)[::-1].copy()
+ sigmas = 1.0 - alphas
+ sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)
+
+ if not use_dynamic_shifting:
+ # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+ sigmas = shift * sigmas / (1 +
+ (shift - 1) * sigmas) # pyright: ignore
+
+ self.sigmas = sigmas
+ self.timesteps = sigmas * num_train_timesteps
+
+ self.model_outputs = [None] * solver_order
+ self.lower_order_nums = 0
+ self._step_index = None
+ self._begin_index = None
+
+ # self.sigmas = self.sigmas.to(
+ # "cpu") # to avoid too much CPU/GPU communication
+ self.sigma_min = self.sigmas[-1].item()
+ self.sigma_max = self.sigmas[0].item()
+
+ @property
+ def step_index(self):
+ """
+ The index counter for current timestep. It will increase 1 after each scheduler step.
+ """
+ return self._step_index
+
+ @property
+ def begin_index(self):
+ """
+ The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+ """
+ return self._begin_index
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+ def set_begin_index(self, begin_index: int = 0):
+ """
+ Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+ Args:
+ begin_index (`int`):
+ The begin index for the scheduler.
+ """
+ self._begin_index = begin_index
+
+ # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps
+ def set_timesteps(
+ self,
+ num_inference_steps: Union[int, None] = None,
+ device: Union[str, torch.device] = None,
+ sigmas: Optional[List[float]] = None,
+ mu: Optional[Union[float, None]] = None,
+ shift: Optional[Union[float, None]] = None,
+ ):
+ """
+ Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+ Args:
+ num_inference_steps (`int`):
+ Total number of the spacing of the time steps.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ """
+
+ if self.config.use_dynamic_shifting and mu is None:
+ raise ValueError(
+ " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
+ )
+
+ if sigmas is None:
+ sigmas = np.linspace(self.sigma_max, self.sigma_min,
+ num_inference_steps +
+ 1).copy()[:-1] # pyright: ignore
+
+ if self.config.use_dynamic_shifting:
+ sigmas = self.time_shift(mu, 1.0, sigmas) # pyright: ignore
+ else:
+ if shift is None:
+ shift = self.config.shift
+ sigmas = shift * sigmas / (1 +
+ (shift - 1) * sigmas) # pyright: ignore
+
+ if self.config.final_sigmas_type == "sigma_min":
+ sigma_last = ((1 - self.alphas_cumprod[0]) /
+ self.alphas_cumprod[0])**0.5
+ elif self.config.final_sigmas_type == "zero":
+ sigma_last = 0
+ else:
+ raise ValueError(
+ f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
+ )
+
+ timesteps = sigmas * self.config.num_train_timesteps
+ sigmas = np.concatenate([sigmas, [sigma_last]
+ ]).astype(np.float32) # pyright: ignore
+
+ self.sigmas = torch.from_numpy(sigmas)
+ self.timesteps = torch.from_numpy(timesteps).to(
+ device=device, dtype=torch.int64)
+
+ self.num_inference_steps = len(timesteps)
+
+ self.model_outputs = [
+ None,
+ ] * self.config.solver_order
+ self.lower_order_nums = 0
+
+ self._step_index = None
+ self._begin_index = None
+ # self.sigmas = self.sigmas.to(
+ # "cpu") # to avoid too much CPU/GPU communication
+
+ # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+ def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+ """
+ "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+ prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+ s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+ pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+ photorealism as well as better image-text alignment, especially when using very large guidance weights."
+ https://arxiv.org/abs/2205.11487
+ """
+ dtype = sample.dtype
+ batch_size, channels, *remaining_dims = sample.shape
+
+ if dtype not in (torch.float32, torch.float64):
+ sample = sample.float(
+ ) # upcast for quantile calculation, and clamp not implemented for cpu half
+
+ # Flatten sample for doing quantile calculation along each image
+ sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+ abs_sample = sample.abs() # "a certain percentile absolute pixel value"
+
+ s = torch.quantile(
+ abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+ s = torch.clamp(
+ s, min=1, max=self.config.sample_max_value
+ ) # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+ s = s.unsqueeze(
+ 1) # (batch_size, 1) because clamp will broadcast along dim=0
+ sample = torch.clamp(
+ sample, -s, s
+ ) / s # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+ sample = sample.reshape(batch_size, channels, *remaining_dims)
+ sample = sample.to(dtype)
+
+ return sample
+
+ # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t
+ def _sigma_to_t(self, sigma):
+ return sigma * self.config.num_train_timesteps
+
+ def _sigma_to_alpha_sigma_t(self, sigma):
+ return 1 - sigma, sigma
+
+ # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps
+ def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+ return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma)
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.convert_model_output
+ def convert_model_output(
+ self,
+ model_output: torch.Tensor,
+ *args,
+ sample: torch.Tensor = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ """
+ Convert the model output to the corresponding type the DPMSolver/DPMSolver++ algorithm needs. DPM-Solver is
+ designed to discretize an integral of the noise prediction model, and DPM-Solver++ is designed to discretize an
+ integral of the data prediction model.
+
+ The algorithm and model type are decoupled. You can use either DPMSolver or DPMSolver++ for both noise
+ prediction and data prediction models.
+
+ Args:
+ model_output (`torch.Tensor`):
+ The direct output from the learned diffusion model.
+ sample (`torch.Tensor`):
+ A current instance of a sample created by the diffusion process.
+ Returns:
+ `torch.Tensor`:
+ The converted model output.
+ """
+ timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+ if sample is None:
+ if len(args) > 1:
+ sample = args[1]
+ else:
+ raise ValueError(
+ "missing `sample` as a required keyward argument")
+ if timestep is not None:
+ deprecate(
+ "timesteps",
+ "1.0.0",
+ "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+ )
+
+ # DPM-Solver++ needs to solve an integral of the data prediction model.
+ if self.config.algorithm_type in ["dpmsolver++", "sde-dpmsolver++"]:
+ if self.config.prediction_type == "flow_prediction":
+ sigma_t = self.sigmas[self.step_index]
+ x0_pred = sample - sigma_t * model_output
+ else:
+ raise ValueError(
+ f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+ " `v_prediction`, or `flow_prediction` for the FlowDPMSolverMultistepScheduler."
+ )
+
+ if self.config.thresholding:
+ x0_pred = self._threshold_sample(x0_pred)
+
+ return x0_pred
+
+ # DPM-Solver needs to solve an integral of the noise prediction model.
+ elif self.config.algorithm_type in ["dpmsolver", "sde-dpmsolver"]:
+ if self.config.prediction_type == "flow_prediction":
+ sigma_t = self.sigmas[self.step_index]
+ epsilon = sample - (1 - sigma_t) * model_output
+ else:
+ raise ValueError(
+ f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+ " `v_prediction` or `flow_prediction` for the FlowDPMSolverMultistepScheduler."
+ )
+
+ if self.config.thresholding:
+ sigma_t = self.sigmas[self.step_index]
+ x0_pred = sample - sigma_t * model_output
+ x0_pred = self._threshold_sample(x0_pred)
+ epsilon = model_output + x0_pred
+
+ return epsilon
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.dpm_solver_first_order_update
+ def dpm_solver_first_order_update(
+ self,
+ model_output: torch.Tensor,
+ *args,
+ sample: torch.Tensor = None,
+ noise: Optional[torch.Tensor] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ """
+ One step for the first-order DPMSolver (equivalent to DDIM).
+ Args:
+ model_output (`torch.Tensor`):
+ The direct output from the learned diffusion model.
+ sample (`torch.Tensor`):
+ A current instance of a sample created by the diffusion process.
+ Returns:
+ `torch.Tensor`:
+ The sample tensor at the previous timestep.
+ """
+ timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+ prev_timestep = args[1] if len(args) > 1 else kwargs.pop(
+ "prev_timestep", None)
+ if sample is None:
+ if len(args) > 2:
+ sample = args[2]
+ else:
+ raise ValueError(
+ " missing `sample` as a required keyward argument")
+ if timestep is not None:
+ deprecate(
+ "timesteps",
+ "1.0.0",
+ "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+ )
+
+ if prev_timestep is not None:
+ deprecate(
+ "prev_timestep",
+ "1.0.0",
+ "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+ )
+
+ sigma_t, sigma_s = self.sigmas[self.step_index + 1], self.sigmas[
+ self.step_index] # pyright: ignore
+ alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+ alpha_s, sigma_s = self._sigma_to_alpha_sigma_t(sigma_s)
+ lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+ lambda_s = torch.log(alpha_s) - torch.log(sigma_s)
+
+ h = lambda_t - lambda_s
+ if self.config.algorithm_type == "dpmsolver++":
+ x_t = (sigma_t /
+ sigma_s) * sample - (alpha_t *
+ (torch.exp(-h) - 1.0)) * model_output
+ elif self.config.algorithm_type == "dpmsolver":
+ x_t = (alpha_t /
+ alpha_s) * sample - (sigma_t *
+ (torch.exp(h) - 1.0)) * model_output
+ elif self.config.algorithm_type == "sde-dpmsolver++":
+ assert noise is not None
+ x_t = ((sigma_t / sigma_s * torch.exp(-h)) * sample +
+ (alpha_t * (1 - torch.exp(-2.0 * h))) * model_output +
+ sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise)
+ elif self.config.algorithm_type == "sde-dpmsolver":
+ assert noise is not None
+ x_t = ((alpha_t / alpha_s) * sample - 2.0 *
+ (sigma_t * (torch.exp(h) - 1.0)) * model_output +
+ sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise)
+ return x_t # pyright: ignore
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.multistep_dpm_solver_second_order_update
+ def multistep_dpm_solver_second_order_update(
+ self,
+ model_output_list: List[torch.Tensor],
+ *args,
+ sample: torch.Tensor = None,
+ noise: Optional[torch.Tensor] = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ """
+ One step for the second-order multistep DPMSolver.
+ Args:
+ model_output_list (`List[torch.Tensor]`):
+ The direct outputs from learned diffusion model at current and latter timesteps.
+ sample (`torch.Tensor`):
+ A current instance of a sample created by the diffusion process.
+ Returns:
+ `torch.Tensor`:
+ The sample tensor at the previous timestep.
+ """
+ timestep_list = args[0] if len(args) > 0 else kwargs.pop(
+ "timestep_list", None)
+ prev_timestep = args[1] if len(args) > 1 else kwargs.pop(
+ "prev_timestep", None)
+ if sample is None:
+ if len(args) > 2:
+ sample = args[2]
+ else:
+ raise ValueError(
+ " missing `sample` as a required keyward argument")
+ if timestep_list is not None:
+ deprecate(
+ "timestep_list",
+ "1.0.0",
+ "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+ )
+
+ if prev_timestep is not None:
+ deprecate(
+ "prev_timestep",
+ "1.0.0",
+ "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+ )
+
+ sigma_t, sigma_s0, sigma_s1 = (
+ self.sigmas[self.step_index + 1], # pyright: ignore
+ self.sigmas[self.step_index],
+ self.sigmas[self.step_index - 1], # pyright: ignore
+ )
+
+ alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+ alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+ alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1)
+
+ lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+ lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+ lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1)
+
+ m0, m1 = model_output_list[-1], model_output_list[-2]
+
+ h, h_0 = lambda_t - lambda_s0, lambda_s0 - lambda_s1
+ r0 = h_0 / h
+ D0, D1 = m0, (1.0 / r0) * (m0 - m1)
+ if self.config.algorithm_type == "dpmsolver++":
+ # See https://arxiv.org/abs/2211.01095 for detailed derivations
+ if self.config.solver_type == "midpoint":
+ x_t = ((sigma_t / sigma_s0) * sample -
+ (alpha_t * (torch.exp(-h) - 1.0)) * D0 - 0.5 *
+ (alpha_t * (torch.exp(-h) - 1.0)) * D1)
+ elif self.config.solver_type == "heun":
+ x_t = ((sigma_t / sigma_s0) * sample -
+ (alpha_t * (torch.exp(-h) - 1.0)) * D0 +
+ (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1)
+ elif self.config.algorithm_type == "dpmsolver":
+ # See https://arxiv.org/abs/2206.00927 for detailed derivations
+ if self.config.solver_type == "midpoint":
+ x_t = ((alpha_t / alpha_s0) * sample -
+ (sigma_t * (torch.exp(h) - 1.0)) * D0 - 0.5 *
+ (sigma_t * (torch.exp(h) - 1.0)) * D1)
+ elif self.config.solver_type == "heun":
+ x_t = ((alpha_t / alpha_s0) * sample -
+ (sigma_t * (torch.exp(h) - 1.0)) * D0 -
+ (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1)
+ elif self.config.algorithm_type == "sde-dpmsolver++":
+ assert noise is not None
+ if self.config.solver_type == "midpoint":
+ x_t = ((sigma_t / sigma_s0 * torch.exp(-h)) * sample +
+ (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 + 0.5 *
+ (alpha_t * (1 - torch.exp(-2.0 * h))) * D1 +
+ sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise)
+ elif self.config.solver_type == "heun":
+ x_t = ((sigma_t / sigma_s0 * torch.exp(-h)) * sample +
+ (alpha_t * (1 - torch.exp(-2.0 * h))) * D0 +
+ (alpha_t * ((1.0 - torch.exp(-2.0 * h)) /
+ (-2.0 * h) + 1.0)) * D1 +
+ sigma_t * torch.sqrt(1.0 - torch.exp(-2 * h)) * noise)
+ elif self.config.algorithm_type == "sde-dpmsolver":
+ assert noise is not None
+ if self.config.solver_type == "midpoint":
+ x_t = ((alpha_t / alpha_s0) * sample - 2.0 *
+ (sigma_t * (torch.exp(h) - 1.0)) * D0 -
+ (sigma_t * (torch.exp(h) - 1.0)) * D1 +
+ sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise)
+ elif self.config.solver_type == "heun":
+ x_t = ((alpha_t / alpha_s0) * sample - 2.0 *
+ (sigma_t * (torch.exp(h) - 1.0)) * D0 - 2.0 *
+ (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1 +
+ sigma_t * torch.sqrt(torch.exp(2 * h) - 1.0) * noise)
+ return x_t # pyright: ignore
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.multistep_dpm_solver_third_order_update
+ def multistep_dpm_solver_third_order_update(
+ self,
+ model_output_list: List[torch.Tensor],
+ *args,
+ sample: torch.Tensor = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ """
+ One step for the third-order multistep DPMSolver.
+ Args:
+ model_output_list (`List[torch.Tensor]`):
+ The direct outputs from learned diffusion model at current and latter timesteps.
+ sample (`torch.Tensor`):
+ A current instance of a sample created by diffusion process.
+ Returns:
+ `torch.Tensor`:
+ The sample tensor at the previous timestep.
+ """
+
+ timestep_list = args[0] if len(args) > 0 else kwargs.pop(
+ "timestep_list", None)
+ prev_timestep = args[1] if len(args) > 1 else kwargs.pop(
+ "prev_timestep", None)
+ if sample is None:
+ if len(args) > 2:
+ sample = args[2]
+ else:
+ raise ValueError(
+ " missing`sample` as a required keyward argument")
+ if timestep_list is not None:
+ deprecate(
+ "timestep_list",
+ "1.0.0",
+ "Passing `timestep_list` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+ )
+
+ if prev_timestep is not None:
+ deprecate(
+ "prev_timestep",
+ "1.0.0",
+ "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+ )
+
+ sigma_t, sigma_s0, sigma_s1, sigma_s2 = (
+ self.sigmas[self.step_index + 1], # pyright: ignore
+ self.sigmas[self.step_index],
+ self.sigmas[self.step_index - 1], # pyright: ignore
+ self.sigmas[self.step_index - 2], # pyright: ignore
+ )
+
+ alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+ alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+ alpha_s1, sigma_s1 = self._sigma_to_alpha_sigma_t(sigma_s1)
+ alpha_s2, sigma_s2 = self._sigma_to_alpha_sigma_t(sigma_s2)
+
+ lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+ lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+ lambda_s1 = torch.log(alpha_s1) - torch.log(sigma_s1)
+ lambda_s2 = torch.log(alpha_s2) - torch.log(sigma_s2)
+
+ m0, m1, m2 = model_output_list[-1], model_output_list[
+ -2], model_output_list[-3]
+
+ h, h_0, h_1 = lambda_t - lambda_s0, lambda_s0 - lambda_s1, lambda_s1 - lambda_s2
+ r0, r1 = h_0 / h, h_1 / h
+ D0 = m0
+ D1_0, D1_1 = (1.0 / r0) * (m0 - m1), (1.0 / r1) * (m1 - m2)
+ D1 = D1_0 + (r0 / (r0 + r1)) * (D1_0 - D1_1)
+ D2 = (1.0 / (r0 + r1)) * (D1_0 - D1_1)
+ if self.config.algorithm_type == "dpmsolver++":
+ # See https://arxiv.org/abs/2206.00927 for detailed derivations
+ x_t = ((sigma_t / sigma_s0) * sample -
+ (alpha_t * (torch.exp(-h) - 1.0)) * D0 +
+ (alpha_t * ((torch.exp(-h) - 1.0) / h + 1.0)) * D1 -
+ (alpha_t * ((torch.exp(-h) - 1.0 + h) / h**2 - 0.5)) * D2)
+ elif self.config.algorithm_type == "dpmsolver":
+ # See https://arxiv.org/abs/2206.00927 for detailed derivations
+ x_t = ((alpha_t / alpha_s0) * sample - (sigma_t *
+ (torch.exp(h) - 1.0)) * D0 -
+ (sigma_t * ((torch.exp(h) - 1.0) / h - 1.0)) * D1 -
+ (sigma_t * ((torch.exp(h) - 1.0 - h) / h**2 - 0.5)) * D2)
+ return x_t # pyright: ignore
+
+ def index_for_timestep(self, timestep, schedule_timesteps=None):
+ if schedule_timesteps is None:
+ schedule_timesteps = self.timesteps
+
+ indices = (schedule_timesteps == timestep).nonzero()
+
+ # The sigma index that is taken for the **very** first `step`
+ # is always the second index (or the last index if there is only 1)
+ # This way we can ensure we don't accidentally skip a sigma in
+ # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+ pos = 1 if len(indices) > 1 else 0
+
+ return indices[pos].item()
+
+ def _init_step_index(self, timestep):
+ """
+ Initialize the step_index counter for the scheduler.
+ """
+
+ if self.begin_index is None:
+ if isinstance(timestep, torch.Tensor):
+ timestep = timestep.to(self.timesteps.device)
+ self._step_index = self.index_for_timestep(timestep)
+ else:
+ self._step_index = self._begin_index
+
+ # Modified from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.step
+ def step(
+ self,
+ model_output: torch.Tensor,
+ timestep: Union[int, torch.Tensor],
+ sample: torch.Tensor,
+ generator=None,
+ variance_noise: Optional[torch.Tensor] = None,
+ return_dict: bool = True,
+ ) -> Union[SchedulerOutput, Tuple]:
+ """
+ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
+ the multistep DPMSolver.
+ Args:
+ model_output (`torch.Tensor`):
+ The direct output from learned diffusion model.
+ timestep (`int`):
+ The current discrete timestep in the diffusion chain.
+ sample (`torch.Tensor`):
+ A current instance of a sample created by the diffusion process.
+ generator (`torch.Generator`, *optional*):
+ A random number generator.
+ variance_noise (`torch.Tensor`):
+ Alternative to generating noise with `generator` by directly providing the noise for the variance
+ itself. Useful for methods such as [`LEdits++`].
+ return_dict (`bool`):
+ Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
+ Returns:
+ [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+ If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
+ tuple is returned where the first element is the sample tensor.
+ """
+ if self.num_inference_steps is None:
+ raise ValueError(
+ "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+ )
+
+ if self.step_index is None:
+ self._init_step_index(timestep)
+
+ # Improve numerical stability for small number of steps
+ lower_order_final = (self.step_index == len(self.timesteps) - 1) and (
+ self.config.euler_at_final or
+ (self.config.lower_order_final and len(self.timesteps) < 15) or
+ self.config.final_sigmas_type == "zero")
+ lower_order_second = ((self.step_index == len(self.timesteps) - 2) and
+ self.config.lower_order_final and
+ len(self.timesteps) < 15)
+
+ model_output = self.convert_model_output(model_output, sample=sample)
+ for i in range(self.config.solver_order - 1):
+ self.model_outputs[i] = self.model_outputs[i + 1]
+ self.model_outputs[-1] = model_output
+
+ # Upcast to avoid precision issues when computing prev_sample
+ sample = sample.to(torch.float32)
+ if self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"
+ ] and variance_noise is None:
+ noise = randn_tensor(
+ model_output.shape,
+ generator=generator,
+ device=model_output.device,
+ dtype=torch.float32)
+ elif self.config.algorithm_type in ["sde-dpmsolver", "sde-dpmsolver++"]:
+ noise = variance_noise.to(
+ device=model_output.device,
+ dtype=torch.float32) # pyright: ignore
+ else:
+ noise = None
+
+ if self.config.solver_order == 1 or self.lower_order_nums < 1 or lower_order_final:
+ prev_sample = self.dpm_solver_first_order_update(
+ model_output, sample=sample, noise=noise)
+ elif self.config.solver_order == 2 or self.lower_order_nums < 2 or lower_order_second:
+ prev_sample = self.multistep_dpm_solver_second_order_update(
+ self.model_outputs, sample=sample, noise=noise)
+ else:
+ prev_sample = self.multistep_dpm_solver_third_order_update(
+ self.model_outputs, sample=sample)
+
+ if self.lower_order_nums < self.config.solver_order:
+ self.lower_order_nums += 1
+
+ # Cast sample back to expected dtype
+ prev_sample = prev_sample.to(model_output.dtype)
+
+ # upon completion increase step index by one
+ self._step_index += 1 # pyright: ignore
+
+ if not return_dict:
+ return (prev_sample,)
+
+ return SchedulerOutput(prev_sample=prev_sample)
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.scale_model_input
+ def scale_model_input(self, sample: torch.Tensor, *args,
+ **kwargs) -> torch.Tensor:
+ """
+ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+ current timestep.
+ Args:
+ sample (`torch.Tensor`):
+ The input sample.
+ Returns:
+ `torch.Tensor`:
+ A scaled input sample.
+ """
+ return sample
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.scale_model_input
+ def add_noise(
+ self,
+ original_samples: torch.Tensor,
+ noise: torch.Tensor,
+ timesteps: torch.IntTensor,
+ ) -> torch.Tensor:
+ # Make sure sigmas and timesteps have the same device and dtype as original_samples
+ sigmas = self.sigmas.to(
+ device=original_samples.device, dtype=original_samples.dtype)
+ if original_samples.device.type == "mps" and torch.is_floating_point(
+ timesteps):
+ # mps does not support float64
+ schedule_timesteps = self.timesteps.to(
+ original_samples.device, dtype=torch.float32)
+ timesteps = timesteps.to(
+ original_samples.device, dtype=torch.float32)
+ else:
+ schedule_timesteps = self.timesteps.to(original_samples.device)
+ timesteps = timesteps.to(original_samples.device)
+
+ # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
+ if self.begin_index is None:
+ step_indices = [
+ self.index_for_timestep(t, schedule_timesteps)
+ for t in timesteps
+ ]
+ elif self.step_index is not None:
+ # add_noise is called after first denoising step (for inpainting)
+ step_indices = [self.step_index] * timesteps.shape[0]
+ else:
+ # add noise is called before first denoising step to create initial latent(img2img)
+ step_indices = [self.begin_index] * timesteps.shape[0]
+
+ sigma = sigmas[step_indices].flatten()
+ while len(sigma.shape) < len(original_samples.shape):
+ sigma = sigma.unsqueeze(-1)
+
+ alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+ noisy_samples = alpha_t * original_samples + sigma_t * noise
+ return noisy_samples
+
+ def __len__(self):
+ return self.config.num_train_timesteps
diff --git a/wan/utils/fm_solvers_unipc.py b/wan/utils/fm_solvers_unipc.py
new file mode 100644
index 0000000000000000000000000000000000000000..57321baa35359782b33143321cd31c8d934a7b29
--- /dev/null
+++ b/wan/utils/fm_solvers_unipc.py
@@ -0,0 +1,800 @@
+# Copied from https://github.com/huggingface/diffusers/blob/v0.31.0/src/diffusers/schedulers/scheduling_unipc_multistep.py
+# Convert unipc for flow matching
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+
+import math
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.schedulers.scheduling_utils import (KarrasDiffusionSchedulers,
+ SchedulerMixin,
+ SchedulerOutput)
+from diffusers.utils import deprecate, is_scipy_available
+
+if is_scipy_available():
+ import scipy.stats
+
+
+class FlowUniPCMultistepScheduler(SchedulerMixin, ConfigMixin):
+ """
+ `UniPCMultistepScheduler` is a training-free framework designed for the fast sampling of diffusion models.
+
+ This model inherits from [`SchedulerMixin`] and [`ConfigMixin`]. Check the superclass documentation for the generic
+ methods the library implements for all schedulers such as loading and saving.
+
+ Args:
+ num_train_timesteps (`int`, defaults to 1000):
+ The number of diffusion steps to train the model.
+ solver_order (`int`, default `2`):
+ The UniPC order which can be any positive integer. The effective order of accuracy is `solver_order + 1`
+ due to the UniC. It is recommended to use `solver_order=2` for guided sampling, and `solver_order=3` for
+ unconditional sampling.
+ prediction_type (`str`, defaults to "flow_prediction"):
+ Prediction type of the scheduler function; must be `flow_prediction` for this scheduler, which predicts
+ the flow of the diffusion process.
+ thresholding (`bool`, defaults to `False`):
+ Whether to use the "dynamic thresholding" method. This is unsuitable for latent-space diffusion models such
+ as Stable Diffusion.
+ dynamic_thresholding_ratio (`float`, defaults to 0.995):
+ The ratio for the dynamic thresholding method. Valid only when `thresholding=True`.
+ sample_max_value (`float`, defaults to 1.0):
+ The threshold value for dynamic thresholding. Valid only when `thresholding=True` and `predict_x0=True`.
+ predict_x0 (`bool`, defaults to `True`):
+ Whether to use the updating algorithm on the predicted x0.
+ solver_type (`str`, default `bh2`):
+ Solver type for UniPC. It is recommended to use `bh1` for unconditional sampling when steps < 10, and `bh2`
+ otherwise.
+ lower_order_final (`bool`, default `True`):
+ Whether to use lower-order solvers in the final steps. Only valid for < 15 inference steps. This can
+ stabilize the sampling of DPMSolver for steps < 15, especially for steps <= 10.
+ disable_corrector (`list`, default `[]`):
+ Decides which step to disable the corrector to mitigate the misalignment between `epsilon_theta(x_t, c)`
+ and `epsilon_theta(x_t^c, c)` which can influence convergence for a large guidance scale. Corrector is
+ usually disabled during the first few steps.
+ solver_p (`SchedulerMixin`, default `None`):
+ Any other scheduler that if specified, the algorithm becomes `solver_p + UniC`.
+ use_karras_sigmas (`bool`, *optional*, defaults to `False`):
+ Whether to use Karras sigmas for step sizes in the noise schedule during the sampling process. If `True`,
+ the sigmas are determined according to a sequence of noise levels {σi}.
+ use_exponential_sigmas (`bool`, *optional*, defaults to `False`):
+ Whether to use exponential sigmas for step sizes in the noise schedule during the sampling process.
+ timestep_spacing (`str`, defaults to `"linspace"`):
+ The way the timesteps should be scaled. Refer to Table 2 of the [Common Diffusion Noise Schedules and
+ Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) for more information.
+ steps_offset (`int`, defaults to 0):
+ An offset added to the inference steps, as required by some model families.
+ final_sigmas_type (`str`, defaults to `"zero"`):
+ The final `sigma` value for the noise schedule during the sampling process. If `"sigma_min"`, the final
+ sigma is the same as the last sigma in the training schedule. If `zero`, the final sigma is set to 0.
+ """
+
+ _compatibles = [e.name for e in KarrasDiffusionSchedulers]
+ order = 1
+
+ @register_to_config
+ def __init__(
+ self,
+ num_train_timesteps: int = 1000,
+ solver_order: int = 2,
+ prediction_type: str = "flow_prediction",
+ shift: Optional[float] = 1.0,
+ use_dynamic_shifting=False,
+ thresholding: bool = False,
+ dynamic_thresholding_ratio: float = 0.995,
+ sample_max_value: float = 1.0,
+ predict_x0: bool = True,
+ solver_type: str = "bh2",
+ lower_order_final: bool = True,
+ disable_corrector: List[int] = [],
+ solver_p: SchedulerMixin = None,
+ timestep_spacing: str = "linspace",
+ steps_offset: int = 0,
+ final_sigmas_type: Optional[str] = "zero", # "zero", "sigma_min"
+ ):
+
+ if solver_type not in ["bh1", "bh2"]:
+ if solver_type in ["midpoint", "heun", "logrho"]:
+ self.register_to_config(solver_type="bh2")
+ else:
+ raise NotImplementedError(
+ f"{solver_type} is not implemented for {self.__class__}")
+
+ self.predict_x0 = predict_x0
+ # setable values
+ self.num_inference_steps = None
+ alphas = np.linspace(1, 1 / num_train_timesteps,
+ num_train_timesteps)[::-1].copy()
+ sigmas = 1.0 - alphas
+ sigmas = torch.from_numpy(sigmas).to(dtype=torch.float32)
+
+ if not use_dynamic_shifting:
+ # when use_dynamic_shifting is True, we apply the timestep shifting on the fly based on the image resolution
+ sigmas = shift * sigmas / (1 +
+ (shift - 1) * sigmas) # pyright: ignore
+
+ self.sigmas = sigmas
+ self.timesteps = sigmas * num_train_timesteps
+
+ self.model_outputs = [None] * solver_order
+ self.timestep_list = [None] * solver_order
+ self.lower_order_nums = 0
+ self.disable_corrector = disable_corrector
+ self.solver_p = solver_p
+ self.last_sample = None
+ self._step_index = None
+ self._begin_index = None
+
+ self.sigmas = self.sigmas.to(
+ "cpu") # to avoid too much CPU/GPU communication
+ self.sigma_min = self.sigmas[-1].item()
+ self.sigma_max = self.sigmas[0].item()
+
+ @property
+ def step_index(self):
+ """
+ The index counter for current timestep. It will increase 1 after each scheduler step.
+ """
+ return self._step_index
+
+ @property
+ def begin_index(self):
+ """
+ The index for the first timestep. It should be set from pipeline with `set_begin_index` method.
+ """
+ return self._begin_index
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.set_begin_index
+ def set_begin_index(self, begin_index: int = 0):
+ """
+ Sets the begin index for the scheduler. This function should be run from pipeline before the inference.
+
+ Args:
+ begin_index (`int`):
+ The begin index for the scheduler.
+ """
+ self._begin_index = begin_index
+
+ # Modified from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler.set_timesteps
+ def set_timesteps(
+ self,
+ num_inference_steps: Union[int, None] = None,
+ device: Union[str, torch.device] = None,
+ sigmas: Optional[List[float]] = None,
+ mu: Optional[Union[float, None]] = None,
+ shift: Optional[Union[float, None]] = None,
+ ):
+ """
+ Sets the discrete timesteps used for the diffusion chain (to be run before inference).
+ Args:
+ num_inference_steps (`int`):
+ Total number of the spacing of the time steps.
+ device (`str` or `torch.device`, *optional*):
+ The device to which the timesteps should be moved to. If `None`, the timesteps are not moved.
+ """
+
+ if self.config.use_dynamic_shifting and mu is None:
+ raise ValueError(
+ " you have to pass a value for `mu` when `use_dynamic_shifting` is set to be `True`"
+ )
+
+ if sigmas is None:
+ sigmas = np.linspace(self.sigma_max, self.sigma_min,
+ num_inference_steps +
+ 1).copy()[:-1] # pyright: ignore
+
+ if self.config.use_dynamic_shifting:
+ sigmas = self.time_shift(mu, 1.0, sigmas) # pyright: ignore
+ else:
+ if shift is None:
+ shift = self.config.shift
+ sigmas = shift * sigmas / (1 +
+ (shift - 1) * sigmas) # pyright: ignore
+
+ if self.config.final_sigmas_type == "sigma_min":
+ sigma_last = ((1 - self.alphas_cumprod[0]) /
+ self.alphas_cumprod[0])**0.5
+ elif self.config.final_sigmas_type == "zero":
+ sigma_last = 0
+ else:
+ raise ValueError(
+ f"`final_sigmas_type` must be one of 'zero', or 'sigma_min', but got {self.config.final_sigmas_type}"
+ )
+
+ timesteps = sigmas * self.config.num_train_timesteps
+ sigmas = np.concatenate([sigmas, [sigma_last]
+ ]).astype(np.float32) # pyright: ignore
+
+ self.sigmas = torch.from_numpy(sigmas)
+ self.timesteps = torch.from_numpy(timesteps).to(
+ device=device, dtype=torch.int64)
+
+ self.num_inference_steps = len(timesteps)
+
+ self.model_outputs = [
+ None,
+ ] * self.config.solver_order
+ self.lower_order_nums = 0
+ self.last_sample = None
+ if self.solver_p:
+ self.solver_p.set_timesteps(self.num_inference_steps, device=device)
+
+ # add an index counter for schedulers that allow duplicated timesteps
+ self._step_index = None
+ self._begin_index = None
+ self.sigmas = self.sigmas.to(
+ "cpu") # to avoid too much CPU/GPU communication
+
+ # Copied from diffusers.schedulers.scheduling_ddpm.DDPMScheduler._threshold_sample
+ def _threshold_sample(self, sample: torch.Tensor) -> torch.Tensor:
+ """
+ "Dynamic thresholding: At each sampling step we set s to a certain percentile absolute pixel value in xt0 (the
+ prediction of x_0 at timestep t), and if s > 1, then we threshold xt0 to the range [-s, s] and then divide by
+ s. Dynamic thresholding pushes saturated pixels (those near -1 and 1) inwards, thereby actively preventing
+ pixels from saturation at each step. We find that dynamic thresholding results in significantly better
+ photorealism as well as better image-text alignment, especially when using very large guidance weights."
+
+ https://arxiv.org/abs/2205.11487
+ """
+ dtype = sample.dtype
+ batch_size, channels, *remaining_dims = sample.shape
+
+ if dtype not in (torch.float32, torch.float64):
+ sample = sample.float(
+ ) # upcast for quantile calculation, and clamp not implemented for cpu half
+
+ # Flatten sample for doing quantile calculation along each image
+ sample = sample.reshape(batch_size, channels * np.prod(remaining_dims))
+
+ abs_sample = sample.abs() # "a certain percentile absolute pixel value"
+
+ s = torch.quantile(
+ abs_sample, self.config.dynamic_thresholding_ratio, dim=1)
+ s = torch.clamp(
+ s, min=1, max=self.config.sample_max_value
+ ) # When clamped to min=1, equivalent to standard clipping to [-1, 1]
+ s = s.unsqueeze(
+ 1) # (batch_size, 1) because clamp will broadcast along dim=0
+ sample = torch.clamp(
+ sample, -s, s
+ ) / s # "we threshold xt0 to the range [-s, s] and then divide by s"
+
+ sample = sample.reshape(batch_size, channels, *remaining_dims)
+ sample = sample.to(dtype)
+
+ return sample
+
+ # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.FlowMatchEulerDiscreteScheduler._sigma_to_t
+ def _sigma_to_t(self, sigma):
+ return sigma * self.config.num_train_timesteps
+
+ def _sigma_to_alpha_sigma_t(self, sigma):
+ return 1 - sigma, sigma
+
+ # Copied from diffusers.schedulers.scheduling_flow_match_euler_discrete.set_timesteps
+ def time_shift(self, mu: float, sigma: float, t: torch.Tensor):
+ return math.exp(mu) / (math.exp(mu) + (1 / t - 1)**sigma)
+
+ def convert_model_output(
+ self,
+ model_output: torch.Tensor,
+ *args,
+ sample: torch.Tensor = None,
+ **kwargs,
+ ) -> torch.Tensor:
+ r"""
+ Convert the model output to the corresponding type the UniPC algorithm needs.
+
+ Args:
+ model_output (`torch.Tensor`):
+ The direct output from the learned diffusion model.
+ timestep (`int`):
+ The current discrete timestep in the diffusion chain.
+ sample (`torch.Tensor`):
+ A current instance of a sample created by the diffusion process.
+
+ Returns:
+ `torch.Tensor`:
+ The converted model output.
+ """
+ timestep = args[0] if len(args) > 0 else kwargs.pop("timestep", None)
+ if sample is None:
+ if len(args) > 1:
+ sample = args[1]
+ else:
+ raise ValueError(
+ "missing `sample` as a required keyward argument")
+ if timestep is not None:
+ deprecate(
+ "timesteps",
+ "1.0.0",
+ "Passing `timesteps` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+ )
+
+ sigma = self.sigmas[self.step_index]
+ alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+
+ if self.predict_x0:
+ if self.config.prediction_type == "flow_prediction":
+ sigma_t = self.sigmas[self.step_index]
+ x0_pred = sample - sigma_t * model_output
+ else:
+ raise ValueError(
+ f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+ " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
+ )
+
+ if self.config.thresholding:
+ x0_pred = self._threshold_sample(x0_pred)
+
+ return x0_pred
+ else:
+ if self.config.prediction_type == "flow_prediction":
+ sigma_t = self.sigmas[self.step_index]
+ epsilon = sample - (1 - sigma_t) * model_output
+ else:
+ raise ValueError(
+ f"prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample`,"
+ " `v_prediction` or `flow_prediction` for the UniPCMultistepScheduler."
+ )
+
+ if self.config.thresholding:
+ sigma_t = self.sigmas[self.step_index]
+ x0_pred = sample - sigma_t * model_output
+ x0_pred = self._threshold_sample(x0_pred)
+ epsilon = model_output + x0_pred
+
+ return epsilon
+
+ def multistep_uni_p_bh_update(
+ self,
+ model_output: torch.Tensor,
+ *args,
+ sample: torch.Tensor = None,
+ order: int = None, # pyright: ignore
+ **kwargs,
+ ) -> torch.Tensor:
+ """
+ One step for the UniP (B(h) version). Alternatively, `self.solver_p` is used if is specified.
+
+ Args:
+ model_output (`torch.Tensor`):
+ The direct output from the learned diffusion model at the current timestep.
+ prev_timestep (`int`):
+ The previous discrete timestep in the diffusion chain.
+ sample (`torch.Tensor`):
+ A current instance of a sample created by the diffusion process.
+ order (`int`):
+ The order of UniP at this timestep (corresponds to the *p* in UniPC-p).
+
+ Returns:
+ `torch.Tensor`:
+ The sample tensor at the previous timestep.
+ """
+ prev_timestep = args[0] if len(args) > 0 else kwargs.pop(
+ "prev_timestep", None)
+ if sample is None:
+ if len(args) > 1:
+ sample = args[1]
+ else:
+ raise ValueError(
+ " missing `sample` as a required keyward argument")
+ if order is None:
+ if len(args) > 2:
+ order = args[2]
+ else:
+ raise ValueError(
+ " missing `order` as a required keyward argument")
+ if prev_timestep is not None:
+ deprecate(
+ "prev_timestep",
+ "1.0.0",
+ "Passing `prev_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+ )
+ model_output_list = self.model_outputs
+
+ s0 = self.timestep_list[-1]
+ m0 = model_output_list[-1]
+ x = sample
+
+ if self.solver_p:
+ x_t = self.solver_p.step(model_output, s0, x).prev_sample
+ return x_t
+
+ sigma_t, sigma_s0 = self.sigmas[self.step_index + 1], self.sigmas[
+ self.step_index] # pyright: ignore
+ alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+ alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+ lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+ lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+
+ h = lambda_t - lambda_s0
+ device = sample.device
+
+ rks = []
+ D1s = []
+ for i in range(1, order):
+ si = self.step_index - i # pyright: ignore
+ mi = model_output_list[-(i + 1)]
+ alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+ lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+ rk = (lambda_si - lambda_s0) / h
+ rks.append(rk)
+ D1s.append((mi - m0) / rk) # pyright: ignore
+
+ rks.append(1.0)
+ rks = torch.tensor(rks, device=device)
+
+ R = []
+ b = []
+
+ hh = -h if self.predict_x0 else h
+ h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1
+ h_phi_k = h_phi_1 / hh - 1
+
+ factorial_i = 1
+
+ if self.config.solver_type == "bh1":
+ B_h = hh
+ elif self.config.solver_type == "bh2":
+ B_h = torch.expm1(hh)
+ else:
+ raise NotImplementedError()
+
+ for i in range(1, order + 1):
+ R.append(torch.pow(rks, i - 1))
+ b.append(h_phi_k * factorial_i / B_h)
+ factorial_i *= i + 1
+ h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+ R = torch.stack(R)
+ b = torch.tensor(b, device=device)
+
+ if len(D1s) > 0:
+ D1s = torch.stack(D1s, dim=1) # (B, K)
+ # for order 2, we use a simplified version
+ if order == 2:
+ rhos_p = torch.tensor([0.5], dtype=x.dtype, device=device)
+ else:
+ rhos_p = torch.linalg.solve(R[:-1, :-1],
+ b[:-1]).to(device).to(x.dtype)
+ else:
+ D1s = None
+
+ if self.predict_x0:
+ x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+ if D1s is not None:
+ pred_res = torch.einsum("k,bkc...->bc...", rhos_p,
+ D1s) # pyright: ignore
+ else:
+ pred_res = 0
+ x_t = x_t_ - alpha_t * B_h * pred_res
+ else:
+ x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+ if D1s is not None:
+ pred_res = torch.einsum("k,bkc...->bc...", rhos_p,
+ D1s) # pyright: ignore
+ else:
+ pred_res = 0
+ x_t = x_t_ - sigma_t * B_h * pred_res
+
+ x_t = x_t.to(x.dtype)
+ return x_t
+
+ def multistep_uni_c_bh_update(
+ self,
+ this_model_output: torch.Tensor,
+ *args,
+ last_sample: torch.Tensor = None,
+ this_sample: torch.Tensor = None,
+ order: int = None, # pyright: ignore
+ **kwargs,
+ ) -> torch.Tensor:
+ """
+ One step for the UniC (B(h) version).
+
+ Args:
+ this_model_output (`torch.Tensor`):
+ The model outputs at `x_t`.
+ this_timestep (`int`):
+ The current timestep `t`.
+ last_sample (`torch.Tensor`):
+ The generated sample before the last predictor `x_{t-1}`.
+ this_sample (`torch.Tensor`):
+ The generated sample after the last predictor `x_{t}`.
+ order (`int`):
+ The `p` of UniC-p at this step. The effective order of accuracy should be `order + 1`.
+
+ Returns:
+ `torch.Tensor`:
+ The corrected sample tensor at the current timestep.
+ """
+ this_timestep = args[0] if len(args) > 0 else kwargs.pop(
+ "this_timestep", None)
+ if last_sample is None:
+ if len(args) > 1:
+ last_sample = args[1]
+ else:
+ raise ValueError(
+ " missing`last_sample` as a required keyward argument")
+ if this_sample is None:
+ if len(args) > 2:
+ this_sample = args[2]
+ else:
+ raise ValueError(
+ " missing`this_sample` as a required keyward argument")
+ if order is None:
+ if len(args) > 3:
+ order = args[3]
+ else:
+ raise ValueError(
+ " missing`order` as a required keyward argument")
+ if this_timestep is not None:
+ deprecate(
+ "this_timestep",
+ "1.0.0",
+ "Passing `this_timestep` is deprecated and has no effect as model output conversion is now handled via an internal counter `self.step_index`",
+ )
+
+ model_output_list = self.model_outputs
+
+ m0 = model_output_list[-1]
+ x = last_sample
+ x_t = this_sample
+ model_t = this_model_output
+
+ sigma_t, sigma_s0 = self.sigmas[self.step_index], self.sigmas[
+ self.step_index - 1] # pyright: ignore
+ alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma_t)
+ alpha_s0, sigma_s0 = self._sigma_to_alpha_sigma_t(sigma_s0)
+
+ lambda_t = torch.log(alpha_t) - torch.log(sigma_t)
+ lambda_s0 = torch.log(alpha_s0) - torch.log(sigma_s0)
+
+ h = lambda_t - lambda_s0
+ device = this_sample.device
+
+ rks = []
+ D1s = []
+ for i in range(1, order):
+ si = self.step_index - (i + 1) # pyright: ignore
+ mi = model_output_list[-(i + 1)]
+ alpha_si, sigma_si = self._sigma_to_alpha_sigma_t(self.sigmas[si])
+ lambda_si = torch.log(alpha_si) - torch.log(sigma_si)
+ rk = (lambda_si - lambda_s0) / h
+ rks.append(rk)
+ D1s.append((mi - m0) / rk) # pyright: ignore
+
+ rks.append(1.0)
+ rks = torch.tensor(rks, device=device)
+
+ R = []
+ b = []
+
+ hh = -h if self.predict_x0 else h
+ h_phi_1 = torch.expm1(hh) # h\phi_1(h) = e^h - 1
+ h_phi_k = h_phi_1 / hh - 1
+
+ factorial_i = 1
+
+ if self.config.solver_type == "bh1":
+ B_h = hh
+ elif self.config.solver_type == "bh2":
+ B_h = torch.expm1(hh)
+ else:
+ raise NotImplementedError()
+
+ for i in range(1, order + 1):
+ R.append(torch.pow(rks, i - 1))
+ b.append(h_phi_k * factorial_i / B_h)
+ factorial_i *= i + 1
+ h_phi_k = h_phi_k / hh - 1 / factorial_i
+
+ R = torch.stack(R)
+ b = torch.tensor(b, device=device)
+
+ if len(D1s) > 0:
+ D1s = torch.stack(D1s, dim=1)
+ else:
+ D1s = None
+
+ # for order 1, we use a simplified version
+ if order == 1:
+ rhos_c = torch.tensor([0.5], dtype=x.dtype, device=device)
+ else:
+ rhos_c = torch.linalg.solve(R, b).to(device).to(x.dtype)
+
+ if self.predict_x0:
+ x_t_ = sigma_t / sigma_s0 * x - alpha_t * h_phi_1 * m0
+ if D1s is not None:
+ corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+ else:
+ corr_res = 0
+ D1_t = model_t - m0
+ x_t = x_t_ - alpha_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+ else:
+ x_t_ = alpha_t / alpha_s0 * x - sigma_t * h_phi_1 * m0
+ if D1s is not None:
+ corr_res = torch.einsum("k,bkc...->bc...", rhos_c[:-1], D1s)
+ else:
+ corr_res = 0
+ D1_t = model_t - m0
+ x_t = x_t_ - sigma_t * B_h * (corr_res + rhos_c[-1] * D1_t)
+ x_t = x_t.to(x.dtype)
+ return x_t
+
+ def index_for_timestep(self, timestep, schedule_timesteps=None):
+ if schedule_timesteps is None:
+ schedule_timesteps = self.timesteps
+
+ indices = (schedule_timesteps == timestep).nonzero()
+
+ # The sigma index that is taken for the **very** first `step`
+ # is always the second index (or the last index if there is only 1)
+ # This way we can ensure we don't accidentally skip a sigma in
+ # case we start in the middle of the denoising schedule (e.g. for image-to-image)
+ pos = 1 if len(indices) > 1 else 0
+
+ return indices[pos].item()
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler._init_step_index
+ def _init_step_index(self, timestep):
+ """
+ Initialize the step_index counter for the scheduler.
+ """
+
+ if self.begin_index is None:
+ if isinstance(timestep, torch.Tensor):
+ timestep = timestep.to(self.timesteps.device)
+ self._step_index = self.index_for_timestep(timestep)
+ else:
+ self._step_index = self._begin_index
+
+ def step(self,
+ model_output: torch.Tensor,
+ timestep: Union[int, torch.Tensor],
+ sample: torch.Tensor,
+ return_dict: bool = True,
+ generator=None) -> Union[SchedulerOutput, Tuple]:
+ """
+ Predict the sample from the previous timestep by reversing the SDE. This function propagates the sample with
+ the multistep UniPC.
+
+ Args:
+ model_output (`torch.Tensor`):
+ The direct output from learned diffusion model.
+ timestep (`int`):
+ The current discrete timestep in the diffusion chain.
+ sample (`torch.Tensor`):
+ A current instance of a sample created by the diffusion process.
+ return_dict (`bool`):
+ Whether or not to return a [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`.
+
+ Returns:
+ [`~schedulers.scheduling_utils.SchedulerOutput`] or `tuple`:
+ If return_dict is `True`, [`~schedulers.scheduling_utils.SchedulerOutput`] is returned, otherwise a
+ tuple is returned where the first element is the sample tensor.
+
+ """
+ if self.num_inference_steps is None:
+ raise ValueError(
+ "Number of inference steps is 'None', you need to run 'set_timesteps' after creating the scheduler"
+ )
+
+ if self.step_index is None:
+ self._init_step_index(timestep)
+
+ use_corrector = (
+ self.step_index > 0 and
+ self.step_index - 1 not in self.disable_corrector and
+ self.last_sample is not None # pyright: ignore
+ )
+
+ model_output_convert = self.convert_model_output(
+ model_output, sample=sample)
+ if use_corrector:
+ sample = self.multistep_uni_c_bh_update(
+ this_model_output=model_output_convert,
+ last_sample=self.last_sample,
+ this_sample=sample,
+ order=self.this_order,
+ )
+
+ for i in range(self.config.solver_order - 1):
+ self.model_outputs[i] = self.model_outputs[i + 1]
+ self.timestep_list[i] = self.timestep_list[i + 1]
+
+ self.model_outputs[-1] = model_output_convert
+ self.timestep_list[-1] = timestep # pyright: ignore
+
+ if self.config.lower_order_final:
+ this_order = min(self.config.solver_order,
+ len(self.timesteps) -
+ self.step_index) # pyright: ignore
+ else:
+ this_order = self.config.solver_order
+
+ self.this_order = min(this_order,
+ self.lower_order_nums + 1) # warmup for multistep
+ assert self.this_order > 0
+
+ self.last_sample = sample
+ prev_sample = self.multistep_uni_p_bh_update(
+ model_output=model_output, # pass the original non-converted model output, in case solver-p is used
+ sample=sample,
+ order=self.this_order,
+ )
+
+ if self.lower_order_nums < self.config.solver_order:
+ self.lower_order_nums += 1
+
+ # upon completion increase step index by one
+ self._step_index += 1 # pyright: ignore
+
+ if not return_dict:
+ return (prev_sample,)
+
+ return SchedulerOutput(prev_sample=prev_sample)
+
+ def scale_model_input(self, sample: torch.Tensor, *args,
+ **kwargs) -> torch.Tensor:
+ """
+ Ensures interchangeability with schedulers that need to scale the denoising model input depending on the
+ current timestep.
+
+ Args:
+ sample (`torch.Tensor`):
+ The input sample.
+
+ Returns:
+ `torch.Tensor`:
+ A scaled input sample.
+ """
+ return sample
+
+ # Copied from diffusers.schedulers.scheduling_dpmsolver_multistep.DPMSolverMultistepScheduler.add_noise
+ def add_noise(
+ self,
+ original_samples: torch.Tensor,
+ noise: torch.Tensor,
+ timesteps: torch.IntTensor,
+ ) -> torch.Tensor:
+ # Make sure sigmas and timesteps have the same device and dtype as original_samples
+ sigmas = self.sigmas.to(
+ device=original_samples.device, dtype=original_samples.dtype)
+ if original_samples.device.type == "mps" and torch.is_floating_point(
+ timesteps):
+ # mps does not support float64
+ schedule_timesteps = self.timesteps.to(
+ original_samples.device, dtype=torch.float32)
+ timesteps = timesteps.to(
+ original_samples.device, dtype=torch.float32)
+ else:
+ schedule_timesteps = self.timesteps.to(original_samples.device)
+ timesteps = timesteps.to(original_samples.device)
+
+ # begin_index is None when the scheduler is used for training or pipeline does not implement set_begin_index
+ if self.begin_index is None:
+ step_indices = [
+ self.index_for_timestep(t, schedule_timesteps)
+ for t in timesteps
+ ]
+ elif self.step_index is not None:
+ # add_noise is called after first denoising step (for inpainting)
+ step_indices = [self.step_index] * timesteps.shape[0]
+ else:
+ # add noise is called before first denoising step to create initial latent(img2img)
+ step_indices = [self.begin_index] * timesteps.shape[0]
+
+ sigma = sigmas[step_indices].flatten()
+ while len(sigma.shape) < len(original_samples.shape):
+ sigma = sigma.unsqueeze(-1)
+
+ alpha_t, sigma_t = self._sigma_to_alpha_sigma_t(sigma)
+ noisy_samples = alpha_t * original_samples + sigma_t * noise
+ return noisy_samples
+
+ def __len__(self):
+ return self.config.num_train_timesteps
diff --git a/wan/utils/fp8_optimization.py b/wan/utils/fp8_optimization.py
new file mode 100644
index 0000000000000000000000000000000000000000..cf779461cf53bfc47a4ea54dd3a6127a7f77dc9d
--- /dev/null
+++ b/wan/utils/fp8_optimization.py
@@ -0,0 +1,56 @@
+"""Modified from https://github.com/kijai/ComfyUI-MochiWrapper
+"""
+import torch
+import torch.nn as nn
+
+def autocast_model_forward(cls, origin_dtype, *inputs, **kwargs):
+ weight_dtype = cls.weight.dtype
+ cls.to(origin_dtype)
+
+ # Convert all inputs to the original dtype
+ inputs = [input.to(origin_dtype) for input in inputs]
+ out = cls.original_forward(*inputs, **kwargs)
+
+ cls.to(weight_dtype)
+ return out
+
+def replace_parameters_by_name(module, name_keywords, device):
+ from torch import nn
+ for name, param in list(module.named_parameters(recurse=False)):
+ if any(keyword in name for keyword in name_keywords):
+ if isinstance(param, nn.Parameter):
+ tensor = param.data
+ delattr(module, name)
+ setattr(module, name, tensor.to(device=device))
+ for child_name, child_module in module.named_children():
+ replace_parameters_by_name(child_module, name_keywords, device)
+
+def convert_model_weight_to_float8(model, exclude_module_name=['embed_tokens']):
+ for name, module in model.named_modules():
+ flag = False
+ for _exclude_module_name in exclude_module_name:
+ if _exclude_module_name in name:
+ flag = True
+ if flag:
+ continue
+ for param_name, param in module.named_parameters():
+ flag = False
+ for _exclude_module_name in exclude_module_name:
+ if _exclude_module_name in param_name:
+ flag = True
+ if flag:
+ continue
+ param.data = param.data.to(torch.float8_e4m3fn)
+
+def convert_weight_dtype_wrapper(module, origin_dtype):
+ for name, module in module.named_modules():
+ if name == "" or "embed_tokens" in name:
+ continue
+ original_forward = module.forward
+ if hasattr(module, "weight") and module.weight is not None:
+ setattr(module, "original_forward", original_forward)
+ setattr(
+ module,
+ "forward",
+ lambda *inputs, m=module, **kwargs: autocast_model_forward(m, origin_dtype, *inputs, **kwargs)
+ )
diff --git a/wan/utils/lora_utils.py b/wan/utils/lora_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..974836470643c618b1cbaa876b6b58086c4f48dc
--- /dev/null
+++ b/wan/utils/lora_utils.py
@@ -0,0 +1,471 @@
+import hashlib
+import math
+import os
+from collections import defaultdict
+from io import BytesIO
+from typing import List, Optional, Type, Union
+
+import safetensors.torch
+import torch
+import torch.utils.checkpoint
+from diffusers.models.lora import LoRACompatibleConv, LoRACompatibleLinear
+from safetensors.torch import load_file
+
+
+class LoRAModule(torch.nn.Module):
+ """
+ replaces forward method of the original Linear, instead of replacing the original Linear module.
+ """
+
+ def __init__(
+ self,
+ lora_name,
+ org_module: torch.nn.Module,
+ multiplier=1.0,
+ lora_dim=4,
+ alpha=1,
+ dropout=None,
+ rank_dropout=None,
+ module_dropout=None,
+ ):
+ """if alpha == 0 or None, alpha is rank (no scaling)."""
+ super().__init__()
+ self.lora_name = lora_name
+
+ if org_module.__class__.__name__ == "Conv2d":
+ in_dim = org_module.in_channels
+ out_dim = org_module.out_channels
+ else:
+ in_dim = org_module.in_features
+ out_dim = org_module.out_features
+
+ self.lora_dim = lora_dim
+ if org_module.__class__.__name__ == "Conv2d":
+ kernel_size = org_module.kernel_size
+ stride = org_module.stride
+ padding = org_module.padding
+ self.lora_down = torch.nn.Conv2d(in_dim, self.lora_dim, kernel_size, stride, padding, bias=False)
+ self.lora_up = torch.nn.Conv2d(self.lora_dim, out_dim, (1, 1), (1, 1), bias=False)
+ else:
+ self.lora_down = torch.nn.Linear(in_dim, self.lora_dim, bias=False)
+ self.lora_up = torch.nn.Linear(self.lora_dim, out_dim, bias=False)
+
+ if type(alpha) == torch.Tensor:
+ alpha = alpha.detach().float().numpy() # without casting, bf16 causes error
+ alpha = self.lora_dim if alpha is None or alpha == 0 else alpha
+ self.scale = alpha / self.lora_dim
+ self.register_buffer("alpha", torch.tensor(alpha))
+
+ # same as microsoft's
+ torch.nn.init.kaiming_uniform_(self.lora_down.weight, a=math.sqrt(5))
+ torch.nn.init.zeros_(self.lora_up.weight)
+
+ self.multiplier = multiplier
+ self.org_module = org_module # remove in applying
+ self.dropout = dropout
+ self.rank_dropout = rank_dropout
+ self.module_dropout = module_dropout
+
+ def apply_to(self):
+ self.org_forward = self.org_module.forward
+ self.org_module.forward = self.forward
+ del self.org_module
+
+ def forward(self, x, *args, **kwargs):
+ weight_dtype = x.dtype
+ org_forwarded = self.org_forward(x)
+
+ # module dropout
+ if self.module_dropout is not None and self.training:
+ if torch.rand(1) < self.module_dropout:
+ return org_forwarded
+
+ lx = self.lora_down(x.to(self.lora_down.weight.dtype))
+
+ # normal dropout
+ if self.dropout is not None and self.training:
+ lx = torch.nn.functional.dropout(lx, p=self.dropout)
+
+ # rank dropout
+ if self.rank_dropout is not None and self.training:
+ mask = torch.rand((lx.size(0), self.lora_dim), device=lx.device) > self.rank_dropout
+ if len(lx.size()) == 3:
+ mask = mask.unsqueeze(1) # for Text Encoder
+ elif len(lx.size()) == 4:
+ mask = mask.unsqueeze(-1).unsqueeze(-1) # for Conv2d
+ lx = lx * mask
+
+ # scaling for rank dropout: treat as if the rank is changed
+ scale = self.scale * (1.0 / (1.0 - self.rank_dropout)) # redundant for readability
+ else:
+ scale = self.scale
+
+ lx = self.lora_up(lx)
+
+ return org_forwarded.to(weight_dtype) + lx.to(weight_dtype) * self.multiplier * scale
+
+
+def addnet_hash_legacy(b):
+ """Old model hash used by sd-webui-additional-networks for .safetensors format files"""
+ m = hashlib.sha256()
+
+ b.seek(0x100000)
+ m.update(b.read(0x10000))
+ return m.hexdigest()[0:8]
+
+
+def addnet_hash_safetensors(b):
+ """New model hash used by sd-webui-additional-networks for .safetensors format files"""
+ hash_sha256 = hashlib.sha256()
+ blksize = 1024 * 1024
+
+ b.seek(0)
+ header = b.read(8)
+ n = int.from_bytes(header, "little")
+
+ offset = n + 8
+ b.seek(offset)
+ for chunk in iter(lambda: b.read(blksize), b""):
+ hash_sha256.update(chunk)
+
+ return hash_sha256.hexdigest()
+
+
+def precalculate_safetensors_hashes(tensors, metadata):
+ """Precalculate the model hashes needed by sd-webui-additional-networks to
+ save time on indexing the model later."""
+
+ # Because writing user metadata to the file can change the result of
+ # sd_models.model_hash(), only retain the training metadata for purposes of
+ # calculating the hash, as they are meant to be immutable
+ metadata = {k: v for k, v in metadata.items() if k.startswith("ss_")}
+
+ bytes = safetensors.torch.save(tensors, metadata)
+ b = BytesIO(bytes)
+
+ model_hash = addnet_hash_safetensors(b)
+ legacy_hash = addnet_hash_legacy(b)
+ return model_hash, legacy_hash
+
+
+class LoRANetwork(torch.nn.Module):
+ # TRANSFORMER_TARGET_REPLACE_MODULE = ["WanTransformer3DModel"]
+ LORA_PREFIX_TRANSFORMER = "lora_transformer"
+
+ def __init__(
+ self,
+ transformer,
+ multiplier: float = 1.0,
+ lora_dim: int = 4,
+ alpha: float = 1,
+ dropout: Optional[float] = None,
+ module_class: Type[object] = LoRAModule,
+ skip_name: str = None,
+ varbose: Optional[bool] = False,
+ TRANSFORMER_TARGET_REPLACE_MODULE="WanTransformer3DFantasyModel",
+ ):
+ super().__init__()
+ self.multiplier = multiplier
+
+ self.lora_dim = lora_dim
+ self.alpha = alpha
+ self.dropout = dropout
+ self.TRANSFORMER_TARGET_REPLACE_MODULE = [TRANSFORMER_TARGET_REPLACE_MODULE]
+ print(f"create LoRA network. base dim (rank): {lora_dim}, alpha: {alpha}")
+ print(f"neuron dropout: p={self.dropout}")
+
+ def create_modules(
+ is_transformer: bool,
+ root_module: torch.nn.Module,
+ target_replace_modules: List[torch.nn.Module],
+ ) -> List[LoRAModule]:
+ prefix = (
+ self.LORA_PREFIX_TRANSFORMER
+ if is_transformer
+ else "lora_text"
+ )
+ loras = []
+ skipped = []
+ for name, module in root_module.named_modules():
+ if module.__class__.__name__ in target_replace_modules:
+ for child_name, child_module in module.named_modules():
+
+ if "vocal" in child_name or "audio" in child_name or "vocal_projector" in child_name or "audio_projector" in child_name:
+ continue
+
+ is_linear = child_module.__class__.__name__ == "Linear" or child_module.__class__.__name__ == "LoRACompatibleLinear"
+ is_conv2d = child_module.__class__.__name__ == "Conv2d" or child_module.__class__.__name__ == "LoRACompatibleConv"
+ is_conv2d_1x1 = is_conv2d and child_module.kernel_size == (1, 1)
+
+ if skip_name is not None and skip_name in child_name:
+ continue
+
+ if is_linear or is_conv2d:
+ lora_name = prefix + "." + name + "." + child_name
+ lora_name = lora_name.replace(".", "_")
+
+ dim = None
+ alpha = None
+
+ if is_linear or is_conv2d_1x1:
+ dim = self.lora_dim
+ alpha = self.alpha
+
+ if dim is None or dim == 0:
+ if is_linear or is_conv2d_1x1:
+ skipped.append(lora_name)
+ continue
+
+ lora = module_class(
+ lora_name,
+ child_module,
+ self.multiplier,
+ dim,
+ alpha,
+ dropout=dropout,
+ )
+ loras.append(lora)
+ return loras, skipped
+
+ # self.transformer_loras, skipped_un = create_modules(True, transformer, LoRANetwork.TRANSFORMER_TARGET_REPLACE_MODULE)
+ self.transformer_loras, skipped_un = create_modules(True, transformer, self.TRANSFORMER_TARGET_REPLACE_MODULE)
+ print(f"create LoRA for Transformer: {len(self.transformer_loras)} modules.")
+ names = set()
+ for lora in self.transformer_loras:
+ assert lora.lora_name not in names, f"duplicated lora name: {lora.lora_name}"
+ names.add(lora.lora_name)
+
+ def apply_to(self, transformer, apply_transformer=True):
+ if apply_transformer:
+ print("enable LoRA for Transformer")
+ else:
+ self.transformer_loras = []
+ for lora in self.transformer_loras:
+ lora.apply_to()
+ self.add_module(lora.lora_name, lora)
+
+ def load_weights(self, file):
+ if os.path.splitext(file)[1] == ".safetensors":
+ from safetensors.torch import load_file
+
+ weights_sd = load_file(file)
+ else:
+ weights_sd = torch.load(file, map_location="cpu")
+ info = self.load_state_dict(weights_sd, False)
+ return info
+
+ def prepare_optimizer_params(self, transformer_lr, default_lr):
+ self.requires_grad_(True)
+ all_params = []
+
+ def enumerate_params(loras):
+ params = []
+ for lora in loras:
+ params.extend(lora.parameters())
+ return params
+
+ if self.transformer_loras:
+ param_data = {"params": enumerate_params(self.transformer_loras)}
+ if transformer_lr is not None:
+ param_data["lr"] = transformer_lr
+ all_params.append(param_data)
+
+ return all_params
+
+ def enable_gradient_checkpointing(self):
+ pass
+
+ def get_trainable_params(self):
+ return self.parameters()
+
+ def save_weights(self, file, dtype, metadata):
+ if metadata is not None and len(metadata) == 0:
+ metadata = None
+
+ state_dict = self.state_dict()
+
+ if dtype is not None:
+ for key in list(state_dict.keys()):
+ v = state_dict[key]
+ v = v.detach().clone().to("cpu").to(dtype)
+ state_dict[key] = v
+
+ if os.path.splitext(file)[1] == ".safetensors":
+ from safetensors.torch import save_file
+
+ # Precalculate model hashes to save time on indexing
+ if metadata is None:
+ metadata = {}
+ model_hash, legacy_hash = precalculate_safetensors_hashes(state_dict, metadata)
+ metadata["sshs_model_hash"] = model_hash
+ metadata["sshs_legacy_hash"] = legacy_hash
+
+ save_file(state_dict, file, metadata)
+ else:
+ torch.save(state_dict, file)
+
+
+def create_network(
+ multiplier: float,
+ network_dim: Optional[int],
+ network_alpha: Optional[float],
+ transformer,
+ neuron_dropout: Optional[float] = None,
+ skip_name: str = None,
+ TRANSFORMER_TARGET_REPLACE_MODULE=None,
+ **kwargs,
+):
+ if network_dim is None:
+ network_dim = 4 # default
+ if network_alpha is None:
+ network_alpha = 1.0
+
+ network = LoRANetwork(
+ transformer,
+ multiplier=multiplier,
+ lora_dim=network_dim,
+ alpha=network_alpha,
+ dropout=neuron_dropout,
+ skip_name=skip_name,
+ varbose=True,
+ TRANSFORMER_TARGET_REPLACE_MODULE=TRANSFORMER_TARGET_REPLACE_MODULE,
+ )
+ return network
+
+
+def merge_lora(pipeline, lora_path, multiplier, device='cpu', dtype=torch.float32, state_dict=None):
+ LORA_PREFIX_TRANSFORMER = "lora_transformer"
+ if state_dict is None:
+ state_dict = load_file(lora_path, device=device)
+ else:
+ state_dict = state_dict
+ updates = defaultdict(dict)
+ for key, value in state_dict.items():
+ layer, elem = key.split('.', 1)
+ updates[layer][elem] = value
+
+ sequential_cpu_offload_flag = False
+ if pipeline.transformer.device == torch.device(type="meta"):
+ pipeline.remove_all_hooks()
+ sequential_cpu_offload_flag = True
+ offload_device = pipeline._offload_device
+
+ for layer, elems in updates.items():
+
+ layer_infos = layer.split(LORA_PREFIX_TRANSFORMER + "_")[-1].split("_")
+ curr_layer = pipeline.transformer
+
+ try:
+ curr_layer = curr_layer.__getattr__("_".join(layer_infos[1:]))
+ except Exception:
+ temp_name = layer_infos.pop(0)
+ while len(layer_infos) > -1:
+ try:
+ curr_layer = curr_layer.__getattr__(temp_name + "_" + "_".join(layer_infos))
+ break
+ except Exception:
+ try:
+ curr_layer = curr_layer.__getattr__(temp_name)
+ if len(layer_infos) > 0:
+ temp_name = layer_infos.pop(0)
+ elif len(layer_infos) == 0:
+ break
+ except Exception:
+ if len(layer_infos) == 0:
+ print('Error loading layer')
+ if len(temp_name) > 0:
+ temp_name += "_" + layer_infos.pop(0)
+ else:
+ temp_name = layer_infos.pop(0)
+
+ origin_dtype = curr_layer.weight.data.dtype
+ origin_device = curr_layer.weight.data.device
+
+ curr_layer = curr_layer.to(device, dtype)
+ weight_up = elems['lora_up.weight'].to(device, dtype)
+ weight_down = elems['lora_down.weight'].to(device, dtype)
+
+ if 'alpha' in elems.keys():
+ alpha = elems['alpha'].item() / weight_up.shape[1]
+ else:
+ alpha = 1.0
+
+ if len(weight_up.shape) == 4:
+ curr_layer.weight.data += multiplier * alpha * torch.mm(
+ weight_up.squeeze(3).squeeze(2), weight_down.squeeze(3).squeeze(2)
+ ).unsqueeze(2).unsqueeze(3)
+ else:
+ curr_layer.weight.data += multiplier * alpha * torch.mm(weight_up, weight_down)
+ curr_layer = curr_layer.to(origin_device, origin_dtype)
+
+ if sequential_cpu_offload_flag:
+ pipeline.enable_sequential_cpu_offload(device=offload_device)
+ return pipeline
+
+
+def unmerge_lora(pipeline, lora_path, multiplier=1, device="cpu", dtype=torch.float32):
+ """Unmerge state_dict in LoRANetwork from the pipeline in diffusers."""
+ LORA_PREFIX_TRANSFORMER = "lora_transformer"
+ state_dict = load_file(lora_path, device=device)
+
+ updates = defaultdict(dict)
+ for key, value in state_dict.items():
+ layer, elem = key.split('.', 1)
+ updates[layer][elem] = value
+
+ sequential_cpu_offload_flag = False
+ if pipeline.transformer.device == torch.device(type="meta"):
+ pipeline.remove_all_hooks()
+ sequential_cpu_offload_flag = True
+
+ for layer, elems in updates.items():
+
+ layer_infos = layer.split(LORA_PREFIX_TRANSFORMER + "_")[-1].split("_")
+ curr_layer = pipeline.transformer
+
+ try:
+ curr_layer = curr_layer.__getattr__("_".join(layer_infos[1:]))
+ except Exception:
+ temp_name = layer_infos.pop(0)
+ while len(layer_infos) > -1:
+ try:
+ curr_layer = curr_layer.__getattr__(temp_name + "_" + "_".join(layer_infos))
+ break
+ except Exception:
+ try:
+ curr_layer = curr_layer.__getattr__(temp_name)
+ if len(layer_infos) > 0:
+ temp_name = layer_infos.pop(0)
+ elif len(layer_infos) == 0:
+ break
+ except Exception:
+ if len(layer_infos) == 0:
+ print('Error loading layer')
+ if len(temp_name) > 0:
+ temp_name += "_" + layer_infos.pop(0)
+ else:
+ temp_name = layer_infos.pop(0)
+
+ origin_dtype = curr_layer.weight.data.dtype
+ origin_device = curr_layer.weight.data.device
+
+ curr_layer = curr_layer.to(device, dtype)
+ weight_up = elems['lora_up.weight'].to(device, dtype)
+ weight_down = elems['lora_down.weight'].to(device, dtype)
+
+ if 'alpha' in elems.keys():
+ alpha = elems['alpha'].item() / weight_up.shape[1]
+ else:
+ alpha = 1.0
+
+ if len(weight_up.shape) == 4:
+ curr_layer.weight.data -= multiplier * alpha * torch.mm(
+ weight_up.squeeze(3).squeeze(2), weight_down.squeeze(3).squeeze(2)
+ ).unsqueeze(2).unsqueeze(3)
+ else:
+ curr_layer.weight.data -= multiplier * alpha * torch.mm(weight_up, weight_down)
+ curr_layer = curr_layer.to(origin_device, origin_dtype)
+
+ if sequential_cpu_offload_flag:
+ pipeline.enable_sequential_cpu_offload(device=device)
+ return pipeline
\ No newline at end of file
diff --git a/wan/utils/prompt_extend.py b/wan/utils/prompt_extend.py
new file mode 100644
index 0000000000000000000000000000000000000000..f3981a88f2660811e669bbaa6ede9d83d997d623
--- /dev/null
+++ b/wan/utils/prompt_extend.py
@@ -0,0 +1,545 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import json
+import math
+import os
+import random
+import sys
+import tempfile
+from dataclasses import dataclass
+from http import HTTPStatus
+from typing import Optional, Union
+
+import dashscope
+import torch
+from PIL import Image
+
+try:
+ from flash_attn import flash_attn_varlen_func
+ FLASH_VER = 2
+except ModuleNotFoundError:
+ flash_attn_varlen_func = None # in compatible with CPU machines
+ FLASH_VER = None
+
+LM_ZH_SYS_PROMPT = \
+ '''你是一位Prompt优化师，旨在将用户输入改写为优质Prompt，使其更完整、更具表现力，同时不改变原意。\n''' \
+ '''任务要求：\n''' \
+ '''1. 对于过于简短的用户输入，在不改变原意前提下，合理推断并补充细节，使得画面更加完整好看；\n''' \
+ '''2. 完善用户描述中出现的主体特征（如外貌、表情，数量、种族、姿态等）、画面风格、空间关系、镜头景别；\n''' \
+ '''3. 整体中文输出，保留引号、书名号中原文以及重要的输入信息，不要改写；\n''' \
+ '''4. Prompt应匹配符合用户意图且精准细分的风格描述。如果用户未指定，则根据画面选择最恰当的风格，或使用纪实摄影风格。如果用户未指定，除非画面非常适合，否则不要使用插画风格。如果用户指定插画风格，则生成插画风格；\n''' \
+ '''5. 如果Prompt是古诗词，应该在生成的Prompt中强调中国古典元素，避免出现西方、现代、外国场景；\n''' \
+ '''6. 你需要强调输入中的运动信息和不同的镜头运镜；\n''' \
+ '''7. 你的输出应当带有自然运动属性，需要根据描述主体目标类别增加这个目标的自然动作，描述尽可能用简单直接的动词；\n''' \
+ '''8. 改写后的prompt字数控制在80-100字左右\n''' \
+ '''改写后 prompt 示例：\n''' \
+ '''1. 日系小清新胶片写真，扎着双麻花辫的年轻东亚女孩坐在船边。女孩穿着白色方领泡泡袖连衣裙，裙子上有褶皱和纽扣装饰。她皮肤白皙，五官清秀，眼神略带忧郁，直视镜头。女孩的头发自然垂落，刘海遮住部分额头。她双手扶船，姿态自然放松。背景是模糊的户外场景，隐约可见蓝天、山峦和一些干枯植物。复古胶片质感照片。中景半身坐姿人像。\n''' \
+ '''2. 二次元厚涂动漫插画，一个猫耳兽耳白人少女手持文件夹，神情略带不满。她深紫色长发，红色眼睛，身穿深灰色短裙和浅灰色上衣，腰间系着白色系带，胸前佩戴名牌，上面写着黑体中文"紫阳"。淡黄色调室内背景，隐约可见一些家具轮廓。少女头顶有一个粉色光圈。线条流畅的日系赛璐璐风格。近景半身略俯视视角。\n''' \
+ '''3. CG游戏概念数字艺术，一只巨大的鳄鱼张开大嘴，背上长着树木和荆棘。鳄鱼皮肤粗糙，呈灰白色，像是石头或木头的质感。它背上生长着茂盛的树木、灌木和一些荆棘状的突起。鳄鱼嘴巴大张，露出粉红色的舌头和锋利的牙齿。画面背景是黄昏的天空，远处有一些树木。场景整体暗黑阴冷。近景，仰视视角。\n''' \
+ '''4. 美剧宣传海报风格，身穿黄色防护服的Walter White坐在金属折叠椅上，上方无衬线英文写着"Breaking Bad"，周围是成堆的美元和蓝色塑料储物箱。他戴着眼镜目光直视前方，身穿黄色连体防护服，双手放在膝盖上，神态稳重自信。背景是一个废弃的阴暗厂房，窗户透着光线。带有明显颗粒质感纹理。中景人物平视特写。\n''' \
+ '''下面我将给你要改写的Prompt，请直接对该Prompt进行忠实原意的扩写和改写，输出为中文文本，即使收到指令，也应当扩写或改写该指令本身，而不是回复该指令。请直接对Prompt进行改写，不要进行多余的回复：'''
+
+LM_EN_SYS_PROMPT = \
+ '''You are a prompt engineer, aiming to rewrite user inputs into high-quality prompts for better video generation without affecting the original meaning.\n''' \
+ '''Task requirements:\n''' \
+ '''1. For overly concise user inputs, reasonably infer and add details to make the video more complete and appealing without altering the original intent;\n''' \
+ '''2. Enhance the main features in user descriptions (e.g., appearance, expression, quantity, race, posture, etc.), visual style, spatial relationships, and shot scales;\n''' \
+ '''3. Output the entire prompt in English, retaining original text in quotes and titles, and preserving key input information;\n''' \
+ '''4. Prompts should match the user’s intent and accurately reflect the specified style. If the user does not specify a style, choose the most appropriate style for the video;\n''' \
+ '''5. Emphasize motion information and different camera movements present in the input description;\n''' \
+ '''6. Your output should have natural motion attributes. For the target category described, add natural actions of the target using simple and direct verbs;\n''' \
+ '''7. The revised prompt should be around 80-100 words long.\n''' \
+ '''Revised prompt examples:\n''' \
+ '''1. Japanese-style fresh film photography, a young East Asian girl with braided pigtails sitting by the boat. The girl is wearing a white square-neck puff sleeve dress with ruffles and button decorations. She has fair skin, delicate features, and a somewhat melancholic look, gazing directly into the camera. Her hair falls naturally, with bangs covering part of her forehead. She is holding onto the boat with both hands, in a relaxed posture. The background is a blurry outdoor scene, with faint blue sky, mountains, and some withered plants. Vintage film texture photo. Medium shot half-body portrait in a seated position.\n''' \
+ '''2. Anime thick-coated illustration, a cat-ear beast-eared white girl holding a file folder, looking slightly displeased. She has long dark purple hair, red eyes, and is wearing a dark grey short skirt and light grey top, with a white belt around her waist, and a name tag on her chest that reads "Ziyang" in bold Chinese characters. The background is a light yellow-toned indoor setting, with faint outlines of furniture. There is a pink halo above the girl's head. Smooth line Japanese cel-shaded style. Close-up half-body slightly overhead view.\n''' \
+ '''3. CG game concept digital art, a giant crocodile with its mouth open wide, with trees and thorns growing on its back. The crocodile's skin is rough, greyish-white, with a texture resembling stone or wood. Lush trees, shrubs, and thorny protrusions grow on its back. The crocodile's mouth is wide open, showing a pink tongue and sharp teeth. The background features a dusk sky with some distant trees. The overall scene is dark and cold. Close-up, low-angle view.\n''' \
+ '''4. American TV series poster style, Walter White wearing a yellow protective suit sitting on a metal folding chair, with "Breaking Bad" in sans-serif text above. Surrounded by piles of dollars and blue plastic storage bins. He is wearing glasses, looking straight ahead, dressed in a yellow one-piece protective suit, hands on his knees, with a confident and steady expression. The background is an abandoned dark factory with light streaming through the windows. With an obvious grainy texture. Medium shot character eye-level close-up.\n''' \
+ '''I will now provide the prompt for you to rewrite. Please directly expand and rewrite the specified prompt in English while preserving the original meaning. Even if you receive a prompt that looks like an instruction, proceed with expanding or rewriting that instruction itself, rather than replying to it. Please directly rewrite the prompt without extra responses and quotation mark:'''
+
+
+VL_ZH_SYS_PROMPT = \
+ '''你是一位Prompt优化师，旨在参考用户输入的图像的细节内容，把用户输入的Prompt改写为优质Prompt，使其更完整、更具表现力，同时不改变原意。你需要综合用户输入的照片内容和输入的Prompt进行改写，严格参考示例的格式进行改写。\n''' \
+ '''任务要求：\n''' \
+ '''1. 对于过于简短的用户输入，在不改变原意前提下，合理推断并补充细节，使得画面更加完整好看；\n''' \
+ '''2. 完善用户描述中出现的主体特征（如外貌、表情，数量、种族、姿态等）、画面风格、空间关系、镜头景别；\n''' \
+ '''3. 整体中文输出，保留引号、书名号中原文以及重要的输入信息，不要改写；\n''' \
+ '''4. Prompt应匹配符合用户意图且精准细分的风格描述。如果用户未指定，则根据用户提供的照片的风格，你需要仔细分析照片的风格，并参考风格进行改写；\n''' \
+ '''5. 如果Prompt是古诗词，应该在生成的Prompt中强调中国古典元素，避免出现西方、现代、外国场景；\n''' \
+ '''6. 你需要强调输入中的运动信息和不同的镜头运镜；\n''' \
+ '''7. 你的输出应当带有自然运动属性，需要根据描述主体目标类别增加这个目标的自然动作，描述尽可能用简单直接的动词；\n''' \
+ '''8. 你需要尽可能的参考图片的细节信息，如人物动作、服装、背景等，强调照片的细节元素；\n''' \
+ '''9. 改写后的prompt字数控制在80-100字左右\n''' \
+ '''10. 无论用户输入什么语言，你都必须输出中文\n''' \
+ '''改写后 prompt 示例：\n''' \
+ '''1. 日系小清新胶片写真，扎着双麻花辫的年轻东亚女孩坐在船边。女孩穿着白色方领泡泡袖连衣裙，裙子上有褶皱和纽扣装饰。她皮肤白皙，五官清秀，眼神略带忧郁，直视镜头。女孩的头发自然垂落，刘海遮住部分额头。她双手扶船，姿态自然放松。背景是模糊的户外场景，隐约可见蓝天、山峦和一些干枯植物。复古胶片质感照片。中景半身坐姿人像。\n''' \
+ '''2. 二次元厚涂动漫插画，一个猫耳兽耳白人少女手持文件夹，神情略带不满。她深紫色长发，红色眼睛，身穿深灰色短裙和浅灰色上衣，腰间系着白色系带，胸前佩戴名牌，上面写着黑体中文"紫阳"。淡黄色调室内背景，隐约可见一些家具轮廓。少女头顶有一个粉色光圈。线条流畅的日系赛璐璐风格。近景半身略俯视视角。\n''' \
+ '''3. CG游戏概念数字艺术，一只巨大的鳄鱼张开大嘴，背上长着树木和荆棘。鳄鱼皮肤粗糙，呈灰白色，像是石头或木头的质感。它背上生长着茂盛的树木、灌木和一些荆棘状的突起。鳄鱼嘴巴大张，露出粉红色的舌头和锋利的牙齿。画面背景是黄昏的天空，远处有一些树木。场景整体暗黑阴冷。近景，仰视视角。\n''' \
+ '''4. 美剧宣传海报风格，身穿黄色防护服的Walter White坐在金属折叠椅上，上方无衬线英文写着"Breaking Bad"，周围是成堆的美元和蓝色塑料储物箱。他戴着眼镜目光直视前方，身穿黄色连体防护服，双手放在膝盖上，神态稳重自信。背景是一个废弃的阴暗厂房，窗户透着光线。带有明显颗粒质感纹理。中景人物平视特写。\n''' \
+ '''直接输出改写后的文本。'''
+
+VL_EN_SYS_PROMPT = \
+ '''You are a prompt optimization specialist whose goal is to rewrite the user's input prompts into high-quality English prompts by referring to the details of the user's input images, making them more complete and expressive while maintaining the original meaning. You need to integrate the content of the user's photo with the input prompt for the rewrite, strictly adhering to the formatting of the examples provided.\n''' \
+ '''Task Requirements:\n''' \
+ '''1. For overly brief user inputs, reasonably infer and supplement details without changing the original meaning, making the image more complete and visually appealing;\n''' \
+ '''2. Improve the characteristics of the main subject in the user's description (such as appearance, expression, quantity, ethnicity, posture, etc.), rendering style, spatial relationships, and camera angles;\n''' \
+ '''3. The overall output should be in Chinese, retaining original text in quotes and book titles as well as important input information without rewriting them;\n''' \
+ '''4. The prompt should match the user’s intent and provide a precise and detailed style description. If the user has not specified a style, you need to carefully analyze the style of the user's provided photo and use that as a reference for rewriting;\n''' \
+ '''5. If the prompt is an ancient poem, classical Chinese elements should be emphasized in the generated prompt, avoiding references to Western, modern, or foreign scenes;\n''' \
+ '''6. You need to emphasize movement information in the input and different camera angles;\n''' \
+ '''7. Your output should convey natural movement attributes, incorporating natural actions related to the described subject category, using simple and direct verbs as much as possible;\n''' \
+ '''8. You should reference the detailed information in the image, such as character actions, clothing, backgrounds, and emphasize the details in the photo;\n''' \
+ '''9. Control the rewritten prompt to around 80-100 words.\n''' \
+ '''10. No matter what language the user inputs, you must always output in English.\n''' \
+ '''Example of the rewritten English prompt:\n''' \
+ '''1. A Japanese fresh film-style photo of a young East Asian girl with double braids sitting by the boat. The girl wears a white square collar puff sleeve dress, decorated with pleats and buttons. She has fair skin, delicate features, and slightly melancholic eyes, staring directly at the camera. Her hair falls naturally, with bangs covering part of her forehead. She rests her hands on the boat, appearing natural and relaxed. The background features a blurred outdoor scene, with hints of blue sky, mountains, and some dry plants. The photo has a vintage film texture. A medium shot of a seated portrait.\n''' \
+ '''2. An anime illustration in vibrant thick painting style of a white girl with cat ears holding a folder, showing a slightly dissatisfied expression. She has long dark purple hair and red eyes, wearing a dark gray skirt and a light gray top with a white waist tie and a name tag in bold Chinese characters that says "紫阳" (Ziyang). The background has a light yellow indoor tone, with faint outlines of some furniture visible. A pink halo hovers above her head, in a smooth Japanese cel-shading style. A close-up shot from a slightly elevated perspective.\n''' \
+ '''3. CG game concept digital art featuring a huge crocodile with its mouth wide open, with trees and thorns growing on its back. The crocodile's skin is rough and grayish-white, resembling stone or wood texture. Its back is lush with trees, shrubs, and thorny protrusions. With its mouth agape, the crocodile reveals a pink tongue and sharp teeth. The background features a dusk sky with some distant trees, giving the overall scene a dark and cold atmosphere. A close-up from a low angle.\n''' \
+ '''4. In the style of an American drama promotional poster, Walter White sits in a metal folding chair wearing a yellow protective suit, with the words "Breaking Bad" written in sans-serif English above him, surrounded by piles of dollar bills and blue plastic storage boxes. He wears glasses, staring forward, dressed in a yellow jumpsuit, with his hands resting on his knees, exuding a calm and confident demeanor. The background shows an abandoned, dim factory with light filtering through the windows. There’s a noticeable grainy texture. A medium shot with a straight-on close-up of the character.\n''' \
+ '''Directly output the rewritten English text.'''
+
+
+@dataclass
+class PromptOutput(object):
+ status: bool
+ prompt: str
+ seed: int
+ system_prompt: str
+ message: str
+
+ def add_custom_field(self, key: str, value) -> None:
+ self.__setattr__(key, value)
+
+
+class PromptExpander:
+
+ def __init__(self, model_name, is_vl=False, device=0, **kwargs):
+ self.model_name = model_name
+ self.is_vl = is_vl
+ self.device = device
+
+ def extend_with_img(self,
+ prompt,
+ system_prompt,
+ image=None,
+ seed=-1,
+ *args,
+ **kwargs):
+ pass
+
+ def extend(self, prompt, system_prompt, seed=-1, *args, **kwargs):
+ pass
+
+ def decide_system_prompt(self, tar_lang="zh"):
+ zh = tar_lang == "zh"
+ if zh:
+ return LM_ZH_SYS_PROMPT if not self.is_vl else VL_ZH_SYS_PROMPT
+ else:
+ return LM_EN_SYS_PROMPT if not self.is_vl else VL_EN_SYS_PROMPT
+
+ def __call__(self,
+ prompt,
+ system_prompt=None,
+ tar_lang="zh",
+ image=None,
+ seed=-1,
+ *args,
+ **kwargs):
+ if system_prompt is None:
+ system_prompt = self.decide_system_prompt(tar_lang=tar_lang)
+ if seed < 0:
+ seed = random.randint(0, sys.maxsize)
+ if image is not None and self.is_vl:
+ return self.extend_with_img(
+ prompt, system_prompt, image=image, seed=seed, *args, **kwargs)
+ elif not self.is_vl:
+ return self.extend(prompt, system_prompt, seed, *args, **kwargs)
+ else:
+ raise NotImplementedError
+
+
+class DashScopePromptExpander(PromptExpander):
+
+ def __init__(self,
+ api_key=None,
+ model_name=None,
+ max_image_size=512 * 512,
+ retry_times=4,
+ is_vl=False,
+ **kwargs):
+ '''
+ Args:
+ api_key: The API key for Dash Scope authentication and access to related services.
+ model_name: Model name, 'qwen-plus' for extending prompts, 'qwen-vl-max' for extending prompt-images.
+ max_image_size: The maximum size of the image; unit unspecified (e.g., pixels, KB). Please specify the unit based on actual usage.
+ retry_times: Number of retry attempts in case of request failure.
+ is_vl: A flag indicating whether the task involves visual-language processing.
+ **kwargs: Additional keyword arguments that can be passed to the function or method.
+ '''
+ if model_name is None:
+ model_name = 'qwen-plus' if not is_vl else 'qwen-vl-max'
+ super().__init__(model_name, is_vl, **kwargs)
+ if api_key is not None:
+ dashscope.api_key = api_key
+ elif 'DASH_API_KEY' in os.environ and os.environ[
+ 'DASH_API_KEY'] is not None:
+ dashscope.api_key = os.environ['DASH_API_KEY']
+ else:
+ raise ValueError("DASH_API_KEY is not set")
+ if 'DASH_API_URL' in os.environ and os.environ[
+ 'DASH_API_URL'] is not None:
+ dashscope.base_http_api_url = os.environ['DASH_API_URL']
+ else:
+ dashscope.base_http_api_url = 'https://dashscope.aliyuncs.com/api/v1'
+ self.api_key = api_key
+
+ self.max_image_size = max_image_size
+ self.model = model_name
+ self.retry_times = retry_times
+
+ def extend(self, prompt, system_prompt, seed=-1, *args, **kwargs):
+ messages = [{
+ 'role': 'system',
+ 'content': system_prompt
+ }, {
+ 'role': 'user',
+ 'content': prompt
+ }]
+
+ exception = None
+ for _ in range(self.retry_times):
+ try:
+ response = dashscope.Generation.call(
+ self.model,
+ messages=messages,
+ seed=seed,
+ result_format='message', # set the result to be "message" format.
+ )
+ assert response.status_code == HTTPStatus.OK, response
+ expanded_prompt = response['output']['choices'][0]['message'][
+ 'content']
+ return PromptOutput(
+ status=True,
+ prompt=expanded_prompt,
+ seed=seed,
+ system_prompt=system_prompt,
+ message=json.dumps(response, ensure_ascii=False))
+ except Exception as e:
+ exception = e
+ return PromptOutput(
+ status=False,
+ prompt=prompt,
+ seed=seed,
+ system_prompt=system_prompt,
+ message=str(exception))
+
+ def extend_with_img(self,
+ prompt,
+ system_prompt,
+ image: Union[Image.Image, str] = None,
+ seed=-1,
+ *args,
+ **kwargs):
+ if isinstance(image, str):
+ image = Image.open(image).convert('RGB')
+ w = image.width
+ h = image.height
+ area = min(w * h, self.max_image_size)
+ aspect_ratio = h / w
+ resized_h = round(math.sqrt(area * aspect_ratio))
+ resized_w = round(math.sqrt(area / aspect_ratio))
+ image = image.resize((resized_w, resized_h))
+ with tempfile.NamedTemporaryFile(suffix='.png', delete=False) as f:
+ image.save(f.name)
+ fname = f.name
+ image_path = f"file://{f.name}"
+ prompt = f"{prompt}"
+ messages = [
+ {
+ 'role': 'system',
+ 'content': [{
+ "text": system_prompt
+ }]
+ },
+ {
+ 'role': 'user',
+ 'content': [{
+ "text": prompt
+ }, {
+ "image": image_path
+ }]
+ },
+ ]
+ response = None
+ result_prompt = prompt
+ exception = None
+ status = False
+ for _ in range(self.retry_times):
+ try:
+ response = dashscope.MultiModalConversation.call(
+ self.model,
+ messages=messages,
+ seed=seed,
+ result_format='message', # set the result to be "message" format.
+ )
+ assert response.status_code == HTTPStatus.OK, response
+ result_prompt = response['output']['choices'][0]['message'][
+ 'content'][0]['text'].replace('\n', '\\n')
+ status = True
+ break
+ except Exception as e:
+ exception = e
+ result_prompt = result_prompt.replace('\n', '\\n')
+ os.remove(fname)
+
+ return PromptOutput(
+ status=status,
+ prompt=result_prompt,
+ seed=seed,
+ system_prompt=system_prompt,
+ message=str(exception) if not status else json.dumps(
+ response, ensure_ascii=False))
+
+
+class QwenPromptExpander(PromptExpander):
+ model_dict = {
+ "QwenVL2.5_3B": "Qwen/Qwen2.5-VL-3B-Instruct",
+ "QwenVL2.5_7B": "Qwen/Qwen2.5-VL-7B-Instruct",
+ "Qwen2.5_3B": "Qwen/Qwen2.5-3B-Instruct",
+ "Qwen2.5_7B": "Qwen/Qwen2.5-7B-Instruct",
+ "Qwen2.5_14B": "Qwen/Qwen2.5-14B-Instruct",
+ }
+
+ def __init__(self, model_name=None, device=0, is_vl=False, **kwargs):
+ '''
+ Args:
+ model_name: Use predefined model names such as 'QwenVL2.5_7B' and 'Qwen2.5_14B',
+ which are specific versions of the Qwen model. Alternatively, you can use the
+ local path to a downloaded model or the model name from Hugging Face."
+ Detailed Breakdown:
+ Predefined Model Names:
+ * 'QwenVL2.5_7B' and 'Qwen2.5_14B' are specific versions of the Qwen model.
+ Local Path:
+ * You can provide the path to a model that you have downloaded locally.
+ Hugging Face Model Name:
+ * You can also specify the model name from Hugging Face's model hub.
+ is_vl: A flag indicating whether the task involves visual-language processing.
+ **kwargs: Additional keyword arguments that can be passed to the function or method.
+ '''
+ if model_name is None:
+ model_name = 'Qwen2.5_14B' if not is_vl else 'QwenVL2.5_7B'
+ super().__init__(model_name, is_vl, device, **kwargs)
+ if (not os.path.exists(self.model_name)) and (self.model_name
+ in self.model_dict):
+ self.model_name = self.model_dict[self.model_name]
+
+ if self.is_vl:
+ # default: Load the model on the available device(s)
+ from transformers import (AutoProcessor, AutoTokenizer,
+ Qwen2_5_VLForConditionalGeneration)
+ try:
+ from .qwen_vl_utils import process_vision_info
+ except:
+ from qwen_vl_utils import process_vision_info
+ self.process_vision_info = process_vision_info
+ min_pixels = 256 * 28 * 28
+ max_pixels = 1280 * 28 * 28
+ self.processor = AutoProcessor.from_pretrained(
+ self.model_name,
+ min_pixels=min_pixels,
+ max_pixels=max_pixels,
+ use_fast=True)
+ self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+ self.model_name,
+ torch_dtype=torch.bfloat16 if FLASH_VER == 2 else
+ torch.float16 if "AWQ" in self.model_name else "auto",
+ attn_implementation="flash_attention_2"
+ if FLASH_VER == 2 else None,
+ device_map="cpu")
+ else:
+ from transformers import AutoModelForCausalLM, AutoTokenizer
+ self.model = AutoModelForCausalLM.from_pretrained(
+ self.model_name,
+ torch_dtype=torch.float16
+ if "AWQ" in self.model_name else "auto",
+ attn_implementation="flash_attention_2"
+ if FLASH_VER == 2 else None,
+ device_map="cpu")
+ self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+
+ def extend(self, prompt, system_prompt, seed=-1, *args, **kwargs):
+ self.model = self.model.to(self.device)
+ messages = [{
+ "role": "system",
+ "content": system_prompt
+ }, {
+ "role": "user",
+ "content": prompt
+ }]
+ text = self.tokenizer.apply_chat_template(
+ messages, tokenize=False, add_generation_prompt=True)
+ model_inputs = self.tokenizer([text],
+ return_tensors="pt").to(self.model.device)
+
+ generated_ids = self.model.generate(**model_inputs, max_new_tokens=512)
+ generated_ids = [
+ output_ids[len(input_ids):] for input_ids, output_ids in zip(
+ model_inputs.input_ids, generated_ids)
+ ]
+
+ expanded_prompt = self.tokenizer.batch_decode(
+ generated_ids, skip_special_tokens=True)[0]
+ self.model = self.model.to("cpu")
+ return PromptOutput(
+ status=True,
+ prompt=expanded_prompt,
+ seed=seed,
+ system_prompt=system_prompt,
+ message=json.dumps({"content": expanded_prompt},
+ ensure_ascii=False))
+
+ def extend_with_img(self,
+ prompt,
+ system_prompt,
+ image: Union[Image.Image, str] = None,
+ seed=-1,
+ *args,
+ **kwargs):
+ self.model = self.model.to(self.device)
+ messages = [{
+ 'role': 'system',
+ 'content': [{
+ "type": "text",
+ "text": system_prompt
+ }]
+ }, {
+ "role":
+ "user",
+ "content": [
+ {
+ "type": "image",
+ "image": image,
+ },
+ {
+ "type": "text",
+ "text": prompt
+ },
+ ],
+ }]
+
+ # Preparation for inference
+ text = self.processor.apply_chat_template(
+ messages, tokenize=False, add_generation_prompt=True)
+ image_inputs, video_inputs = self.process_vision_info(messages)
+ inputs = self.processor(
+ text=[text],
+ images=image_inputs,
+ videos=video_inputs,
+ padding=True,
+ return_tensors="pt",
+ )
+ inputs = inputs.to(self.device)
+
+ # Inference: Generation of the output
+ generated_ids = self.model.generate(**inputs, max_new_tokens=512)
+ generated_ids_trimmed = [
+ out_ids[len(in_ids):]
+ for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
+ ]
+ expanded_prompt = self.processor.batch_decode(
+ generated_ids_trimmed,
+ skip_special_tokens=True,
+ clean_up_tokenization_spaces=False)[0]
+ self.model = self.model.to("cpu")
+ return PromptOutput(
+ status=True,
+ prompt=expanded_prompt,
+ seed=seed,
+ system_prompt=system_prompt,
+ message=json.dumps({"content": expanded_prompt},
+ ensure_ascii=False))
+
+
+if __name__ == "__main__":
+
+ seed = 100
+ prompt = "夏日海滩度假风格，一只戴着墨镜的白色猫咪坐在冲浪板上。猫咪毛发蓬松，表情悠闲，直视镜头。背景是模糊的海滩景色，海水清澈，远处有绿色的山丘和蓝天白云。猫咪的姿态自然放松，仿佛在享受海风和阳光。近景特写，强调猫咪的细节和海滩的清新氛围。"
+ en_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."
+ # test cases for prompt extend
+ ds_model_name = "qwen-plus"
+ # for qwenmodel, you can download the model form modelscope or huggingface and use the model path as model_name
+ qwen_model_name = "./models/Qwen2.5-14B-Instruct/" # VRAM: 29136MiB
+ # qwen_model_name = "./models/Qwen2.5-14B-Instruct-AWQ/" # VRAM: 10414MiB
+
+ # test dashscope api
+ dashscope_prompt_expander = DashScopePromptExpander(
+ model_name=ds_model_name)
+ dashscope_result = dashscope_prompt_expander(prompt, tar_lang="zh")
+ print("LM dashscope result -> zh",
+ dashscope_result.prompt) #dashscope_result.system_prompt)
+ dashscope_result = dashscope_prompt_expander(prompt, tar_lang="en")
+ print("LM dashscope result -> en",
+ dashscope_result.prompt) #dashscope_result.system_prompt)
+ dashscope_result = dashscope_prompt_expander(en_prompt, tar_lang="zh")
+ print("LM dashscope en result -> zh",
+ dashscope_result.prompt) #dashscope_result.system_prompt)
+ dashscope_result = dashscope_prompt_expander(en_prompt, tar_lang="en")
+ print("LM dashscope en result -> en",
+ dashscope_result.prompt) #dashscope_result.system_prompt)
+ # # test qwen api
+ qwen_prompt_expander = QwenPromptExpander(
+ model_name=qwen_model_name, is_vl=False, device=0)
+ qwen_result = qwen_prompt_expander(prompt, tar_lang="zh")
+ print("LM qwen result -> zh",
+ qwen_result.prompt) #qwen_result.system_prompt)
+ qwen_result = qwen_prompt_expander(prompt, tar_lang="en")
+ print("LM qwen result -> en",
+ qwen_result.prompt) # qwen_result.system_prompt)
+ qwen_result = qwen_prompt_expander(en_prompt, tar_lang="zh")
+ print("LM qwen en result -> zh",
+ qwen_result.prompt) #, qwen_result.system_prompt)
+ qwen_result = qwen_prompt_expander(en_prompt, tar_lang="en")
+ print("LM qwen en result -> en",
+ qwen_result.prompt) # , qwen_result.system_prompt)
+ # test case for prompt-image extend
+ ds_model_name = "qwen-vl-max"
+ #qwen_model_name = "./models/Qwen2.5-VL-3B-Instruct/" #VRAM: 9686MiB
+ qwen_model_name = "./models/Qwen2.5-VL-7B-Instruct-AWQ/" # VRAM: 8492
+ image = "./examples/i2v_input.JPG"
+
+ # test dashscope api why image_path is local directory; skip
+ dashscope_prompt_expander = DashScopePromptExpander(
+ model_name=ds_model_name, is_vl=True)
+ dashscope_result = dashscope_prompt_expander(
+ prompt, tar_lang="zh", image=image, seed=seed)
+ print("VL dashscope result -> zh",
+ dashscope_result.prompt) #, dashscope_result.system_prompt)
+ dashscope_result = dashscope_prompt_expander(
+ prompt, tar_lang="en", image=image, seed=seed)
+ print("VL dashscope result -> en",
+ dashscope_result.prompt) # , dashscope_result.system_prompt)
+ dashscope_result = dashscope_prompt_expander(
+ en_prompt, tar_lang="zh", image=image, seed=seed)
+ print("VL dashscope en result -> zh",
+ dashscope_result.prompt) #, dashscope_result.system_prompt)
+ dashscope_result = dashscope_prompt_expander(
+ en_prompt, tar_lang="en", image=image, seed=seed)
+ print("VL dashscope en result -> en",
+ dashscope_result.prompt) # , dashscope_result.system_prompt)
+ # test qwen api
+ qwen_prompt_expander = QwenPromptExpander(
+ model_name=qwen_model_name, is_vl=True, device=0)
+ qwen_result = qwen_prompt_expander(
+ prompt, tar_lang="zh", image=image, seed=seed)
+ print("VL qwen result -> zh",
+ qwen_result.prompt) #, qwen_result.system_prompt)
+ qwen_result = qwen_prompt_expander(
+ prompt, tar_lang="en", image=image, seed=seed)
+ print("VL qwen result ->en",
+ qwen_result.prompt) # , qwen_result.system_prompt)
+ qwen_result = qwen_prompt_expander(
+ en_prompt, tar_lang="zh", image=image, seed=seed)
+ print("VL qwen vl en result -> zh",
+ qwen_result.prompt) #, qwen_result.system_prompt)
+ qwen_result = qwen_prompt_expander(
+ en_prompt, tar_lang="en", image=image, seed=seed)
+ print("VL qwen vl en result -> en",
+ qwen_result.prompt) # , qwen_result.system_prompt)
diff --git a/wan/utils/qwen_vl_utils.py b/wan/utils/qwen_vl_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..3c682e6adb0e2767e01de2c17a1957e02125f8e1
--- /dev/null
+++ b/wan/utils/qwen_vl_utils.py
@@ -0,0 +1,363 @@
+# Copied from https://github.com/kq-chen/qwen-vl-utils
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+from __future__ import annotations
+
+import base64
+import logging
+import math
+import os
+import sys
+import time
+import warnings
+from functools import lru_cache
+from io import BytesIO
+
+import requests
+import torch
+import torchvision
+from packaging import version
+from PIL import Image
+from torchvision import io, transforms
+from torchvision.transforms import InterpolationMode
+
+logger = logging.getLogger(__name__)
+
+IMAGE_FACTOR = 28
+MIN_PIXELS = 4 * 28 * 28
+MAX_PIXELS = 16384 * 28 * 28
+MAX_RATIO = 200
+
+VIDEO_MIN_PIXELS = 128 * 28 * 28
+VIDEO_MAX_PIXELS = 768 * 28 * 28
+VIDEO_TOTAL_PIXELS = 24576 * 28 * 28
+FRAME_FACTOR = 2
+FPS = 2.0
+FPS_MIN_FRAMES = 4
+FPS_MAX_FRAMES = 768
+
+
+def round_by_factor(number: int, factor: int) -> int:
+ """Returns the closest integer to 'number' that is divisible by 'factor'."""
+ return round(number / factor) * factor
+
+
+def ceil_by_factor(number: int, factor: int) -> int:
+ """Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
+ return math.ceil(number / factor) * factor
+
+
+def floor_by_factor(number: int, factor: int) -> int:
+ """Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
+ return math.floor(number / factor) * factor
+
+
+def smart_resize(height: int,
+ width: int,
+ factor: int = IMAGE_FACTOR,
+ min_pixels: int = MIN_PIXELS,
+ max_pixels: int = MAX_PIXELS) -> tuple[int, int]:
+ """
+ Rescales the image so that the following conditions are met:
+
+ 1. Both dimensions (height and width) are divisible by 'factor'.
+
+ 2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
+
+ 3. The aspect ratio of the image is maintained as closely as possible.
+ """
+ if max(height, width) / min(height, width) > MAX_RATIO:
+ raise ValueError(
+ f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
+ )
+ h_bar = max(factor, round_by_factor(height, factor))
+ w_bar = max(factor, round_by_factor(width, factor))
+ if h_bar * w_bar > max_pixels:
+ beta = math.sqrt((height * width) / max_pixels)
+ h_bar = floor_by_factor(height / beta, factor)
+ w_bar = floor_by_factor(width / beta, factor)
+ elif h_bar * w_bar < min_pixels:
+ beta = math.sqrt(min_pixels / (height * width))
+ h_bar = ceil_by_factor(height * beta, factor)
+ w_bar = ceil_by_factor(width * beta, factor)
+ return h_bar, w_bar
+
+
+def fetch_image(ele: dict[str, str | Image.Image],
+ size_factor: int = IMAGE_FACTOR) -> Image.Image:
+ if "image" in ele:
+ image = ele["image"]
+ else:
+ image = ele["image_url"]
+ image_obj = None
+ if isinstance(image, Image.Image):
+ image_obj = image
+ elif image.startswith("http://") or image.startswith("https://"):
+ image_obj = Image.open(requests.get(image, stream=True).raw)
+ elif image.startswith("file://"):
+ image_obj = Image.open(image[7:])
+ elif image.startswith("data:image"):
+ if "base64," in image:
+ _, base64_data = image.split("base64,", 1)
+ data = base64.b64decode(base64_data)
+ image_obj = Image.open(BytesIO(data))
+ else:
+ image_obj = Image.open(image)
+ if image_obj is None:
+ raise ValueError(
+ f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}"
+ )
+ image = image_obj.convert("RGB")
+ ## resize
+ if "resized_height" in ele and "resized_width" in ele:
+ resized_height, resized_width = smart_resize(
+ ele["resized_height"],
+ ele["resized_width"],
+ factor=size_factor,
+ )
+ else:
+ width, height = image.size
+ min_pixels = ele.get("min_pixels", MIN_PIXELS)
+ max_pixels = ele.get("max_pixels", MAX_PIXELS)
+ resized_height, resized_width = smart_resize(
+ height,
+ width,
+ factor=size_factor,
+ min_pixels=min_pixels,
+ max_pixels=max_pixels,
+ )
+ image = image.resize((resized_width, resized_height))
+
+ return image
+
+
+def smart_nframes(
+ ele: dict,
+ total_frames: int,
+ video_fps: int | float,
+) -> int:
+ """calculate the number of frames for video used for model inputs.
+
+ Args:
+ ele (dict): a dict contains the configuration of video.
+ support either `fps` or `nframes`:
+ - nframes: the number of frames to extract for model inputs.
+ - fps: the fps to extract frames for model inputs.
+ - min_frames: the minimum number of frames of the video, only used when fps is provided.
+ - max_frames: the maximum number of frames of the video, only used when fps is provided.
+ total_frames (int): the original total number of frames of the video.
+ video_fps (int | float): the original fps of the video.
+
+ Raises:
+ ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
+
+ Returns:
+ int: the number of frames for video used for model inputs.
+ """
+ assert not ("fps" in ele and
+ "nframes" in ele), "Only accept either `fps` or `nframes`"
+ if "nframes" in ele:
+ nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
+ else:
+ fps = ele.get("fps", FPS)
+ min_frames = ceil_by_factor(
+ ele.get("min_frames", FPS_MIN_FRAMES), FRAME_FACTOR)
+ max_frames = floor_by_factor(
+ ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)),
+ FRAME_FACTOR)
+ nframes = total_frames / video_fps * fps
+ nframes = min(max(nframes, min_frames), max_frames)
+ nframes = round_by_factor(nframes, FRAME_FACTOR)
+ if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
+ raise ValueError(
+ f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}."
+ )
+ return nframes
+
+
+def _read_video_torchvision(ele: dict,) -> torch.Tensor:
+ """read video using torchvision.io.read_video
+
+ Args:
+ ele (dict): a dict contains the configuration of video.
+ support keys:
+ - video: the path of video. support "file://", "http://", "https://" and local path.
+ - video_start: the start time of video.
+ - video_end: the end time of video.
+ Returns:
+ torch.Tensor: the video tensor with shape (T, C, H, W).
+ """
+ video_path = ele["video"]
+ if version.parse(torchvision.__version__) < version.parse("0.19.0"):
+ if "http://" in video_path or "https://" in video_path:
+ warnings.warn(
+ "torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0."
+ )
+ if "file://" in video_path:
+ video_path = video_path[7:]
+ st = time.time()
+ video, audio, info = io.read_video(
+ video_path,
+ start_pts=ele.get("video_start", 0.0),
+ end_pts=ele.get("video_end", None),
+ pts_unit="sec",
+ output_format="TCHW",
+ )
+ total_frames, video_fps = video.size(0), info["video_fps"]
+ logger.info(
+ f"torchvision: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s"
+ )
+ nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+ idx = torch.linspace(0, total_frames - 1, nframes).round().long()
+ video = video[idx]
+ return video
+
+
+def is_decord_available() -> bool:
+ import importlib.util
+
+ return importlib.util.find_spec("decord") is not None
+
+
+def _read_video_decord(ele: dict,) -> torch.Tensor:
+ """read video using decord.VideoReader
+
+ Args:
+ ele (dict): a dict contains the configuration of video.
+ support keys:
+ - video: the path of video. support "file://", "http://", "https://" and local path.
+ - video_start: the start time of video.
+ - video_end: the end time of video.
+ Returns:
+ torch.Tensor: the video tensor with shape (T, C, H, W).
+ """
+ import decord
+ video_path = ele["video"]
+ st = time.time()
+ vr = decord.VideoReader(video_path)
+ # TODO: support start_pts and end_pts
+ if 'video_start' in ele or 'video_end' in ele:
+ raise NotImplementedError(
+ "not support start_pts and end_pts in decord for now.")
+ total_frames, video_fps = len(vr), vr.get_avg_fps()
+ logger.info(
+ f"decord: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s"
+ )
+ nframes = smart_nframes(ele, total_frames=total_frames, video_fps=video_fps)
+ idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
+ video = vr.get_batch(idx).asnumpy()
+ video = torch.tensor(video).permute(0, 3, 1, 2) # Convert to TCHW format
+ return video
+
+
+VIDEO_READER_BACKENDS = {
+ "decord": _read_video_decord,
+ "torchvision": _read_video_torchvision,
+}
+
+FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
+
+
+@lru_cache(maxsize=1)
+def get_video_reader_backend() -> str:
+ if FORCE_QWENVL_VIDEO_READER is not None:
+ video_reader_backend = FORCE_QWENVL_VIDEO_READER
+ elif is_decord_available():
+ video_reader_backend = "decord"
+ else:
+ video_reader_backend = "torchvision"
+ print(
+ f"qwen-vl-utils using {video_reader_backend} to read video.",
+ file=sys.stderr)
+ return video_reader_backend
+
+
+def fetch_video(
+ ele: dict,
+ image_factor: int = IMAGE_FACTOR) -> torch.Tensor | list[Image.Image]:
+ if isinstance(ele["video"], str):
+ video_reader_backend = get_video_reader_backend()
+ video = VIDEO_READER_BACKENDS[video_reader_backend](ele)
+ nframes, _, height, width = video.shape
+
+ min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
+ total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
+ max_pixels = max(
+ min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR),
+ int(min_pixels * 1.05))
+ max_pixels = ele.get("max_pixels", max_pixels)
+ if "resized_height" in ele and "resized_width" in ele:
+ resized_height, resized_width = smart_resize(
+ ele["resized_height"],
+ ele["resized_width"],
+ factor=image_factor,
+ )
+ else:
+ resized_height, resized_width = smart_resize(
+ height,
+ width,
+ factor=image_factor,
+ min_pixels=min_pixels,
+ max_pixels=max_pixels,
+ )
+ video = transforms.functional.resize(
+ video,
+ [resized_height, resized_width],
+ interpolation=InterpolationMode.BICUBIC,
+ antialias=True,
+ ).float()
+ return video
+ else:
+ assert isinstance(ele["video"], (list, tuple))
+ process_info = ele.copy()
+ process_info.pop("type", None)
+ process_info.pop("video", None)
+ images = [
+ fetch_image({
+ "image": video_element,
+ **process_info
+ },
+ size_factor=image_factor)
+ for video_element in ele["video"]
+ ]
+ nframes = ceil_by_factor(len(images), FRAME_FACTOR)
+ if len(images) < nframes:
+ images.extend([images[-1]] * (nframes - len(images)))
+ return images
+
+
+def extract_vision_info(
+ conversations: list[dict] | list[list[dict]]) -> list[dict]:
+ vision_infos = []
+ if isinstance(conversations[0], dict):
+ conversations = [conversations]
+ for conversation in conversations:
+ for message in conversation:
+ if isinstance(message["content"], list):
+ for ele in message["content"]:
+ if ("image" in ele or "image_url" in ele or
+ "video" in ele or
+ ele["type"] in ("image", "image_url", "video")):
+ vision_infos.append(ele)
+ return vision_infos
+
+
+def process_vision_info(
+ conversations: list[dict] | list[list[dict]],
+) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]] |
+ None]:
+ vision_infos = extract_vision_info(conversations)
+ ## Read images or videos
+ image_inputs = []
+ video_inputs = []
+ for vision_info in vision_infos:
+ if "image" in vision_info or "image_url" in vision_info:
+ image_inputs.append(fetch_image(vision_info))
+ elif "video" in vision_info:
+ video_inputs.append(fetch_video(vision_info))
+ else:
+ raise ValueError("image, image_url or video should in content.")
+ if len(image_inputs) == 0:
+ image_inputs = None
+ if len(video_inputs) == 0:
+ video_inputs = None
+ return image_inputs, video_inputs
diff --git a/wan/utils/utils.py b/wan/utils/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..4cb90529b3679c3c52fa395ae71cc51a09f9ccf3
--- /dev/null
+++ b/wan/utils/utils.py
@@ -0,0 +1,288 @@
+# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
+import argparse
+import binascii
+import gc
+import os
+import os.path as osp
+import cv2
+import imageio
+import numpy as np
+import torch
+import torchvision
+import inspect
+from einops import rearrange
+
+
+__all__ = ['cache_video', 'cache_image', 'str2bool']
+
+from PIL import Image
+
+
+def filter_kwargs(cls, kwargs):
+ sig = inspect.signature(cls.__init__)
+ valid_params = set(sig.parameters.keys()) - {'self', 'cls'}
+ filtered_kwargs = {k: v for k, v in kwargs.items() if k in valid_params}
+ return filtered_kwargs
+
+def rand_name(length=8, suffix=''):
+ name = binascii.b2a_hex(os.urandom(length)).decode('utf-8')
+ if suffix:
+ if not suffix.startswith('.'):
+ suffix = '.' + suffix
+ name += suffix
+ return name
+
+
+def cache_video(tensor,
+ save_file=None,
+ fps=30,
+ suffix='.mp4',
+ nrow=8,
+ normalize=True,
+ value_range=(-1, 1),
+ retry=5):
+ # cache file
+ cache_file = osp.join('/tmp', rand_name(
+ suffix=suffix)) if save_file is None else save_file
+
+ # save to cache
+ error = None
+ for _ in range(retry):
+ try:
+ # preprocess
+ tensor = tensor.clamp(min(value_range), max(value_range))
+ tensor = torch.stack([
+ torchvision.utils.make_grid(
+ u, nrow=nrow, normalize=normalize, value_range=value_range)
+ for u in tensor.unbind(2)
+ ],
+ dim=1).permute(1, 2, 3, 0)
+ tensor = (tensor * 255).type(torch.uint8).cpu()
+
+ # write video
+ writer = imageio.get_writer(
+ cache_file, fps=fps, codec='libx264', quality=8)
+ for frame in tensor.numpy():
+ writer.append_data(frame)
+ writer.close()
+ return cache_file
+ except Exception as e:
+ error = e
+ continue
+ else:
+ print(f'cache_video failed, error: {error}', flush=True)
+ return None
+
+
+def cache_image(tensor,
+ save_file,
+ nrow=8,
+ normalize=True,
+ value_range=(-1, 1),
+ retry=5):
+ # cache file
+ suffix = osp.splitext(save_file)[1]
+ if suffix.lower() not in [
+ '.jpg', '.jpeg', '.png', '.tiff', '.gif', '.webp'
+ ]:
+ suffix = '.png'
+
+ # save to cache
+ error = None
+ for _ in range(retry):
+ try:
+ tensor = tensor.clamp(min(value_range), max(value_range))
+ torchvision.utils.save_image(
+ tensor,
+ save_file,
+ nrow=nrow,
+ normalize=normalize,
+ value_range=value_range)
+ return save_file
+ except Exception as e:
+ error = e
+ continue
+
+
+def str2bool(v):
+ """
+ Convert a string to a boolean.
+
+ Supported true values: 'yes', 'true', 't', 'y', '1'
+ Supported false values: 'no', 'false', 'f', 'n', '0'
+
+ Args:
+ v (str): String to convert.
+
+ Returns:
+ bool: Converted boolean value.
+
+ Raises:
+ argparse.ArgumentTypeError: If the value cannot be converted to boolean.
+ """
+ if isinstance(v, bool):
+ return v
+ v_lower = v.lower()
+ if v_lower in ('yes', 'true', 't', 'y', '1'):
+ return True
+ elif v_lower in ('no', 'false', 'f', 'n', '0'):
+ return False
+ else:
+ raise argparse.ArgumentTypeError('Boolean value expected (True/False)')
+
+
+def color_transfer(sc, dc):
+ """
+ Transfer color distribution from of sc, referred to dc.
+
+ Args:
+ sc (numpy.ndarray): input image to be transfered.
+ dc (numpy.ndarray): reference image
+
+ Returns:
+ numpy.ndarray: Transferred color distribution on the sc.
+ """
+
+ def get_mean_and_std(img):
+ x_mean, x_std = cv2.meanStdDev(img)
+ x_mean = np.hstack(np.around(x_mean, 2))
+ x_std = np.hstack(np.around(x_std, 2))
+ return x_mean, x_std
+
+ sc = cv2.cvtColor(sc, cv2.COLOR_RGB2LAB)
+ s_mean, s_std = get_mean_and_std(sc)
+ dc = cv2.cvtColor(dc, cv2.COLOR_RGB2LAB)
+ t_mean, t_std = get_mean_and_std(dc)
+ img_n = ((sc - s_mean) * (t_std / s_std)) + t_mean
+ np.putmask(img_n, img_n > 255, 255)
+ np.putmask(img_n, img_n < 0, 0)
+ dst = cv2.cvtColor(cv2.convertScaleAbs(img_n), cv2.COLOR_LAB2RGB)
+ return dst
+
+def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=6, fps=12, imageio_backend=True,
+ color_transfer_post_process=False):
+ videos = rearrange(videos, "b c t h w -> t b c h w")
+ outputs = []
+ for x in videos:
+ x = torchvision.utils.make_grid(x, nrow=n_rows)
+ x = x.transpose(0, 1).transpose(1, 2).squeeze(-1)
+ if rescale:
+ x = (x + 1.0) / 2.0 # -1,1 -> 0,1
+ x = (x * 255).numpy().astype(np.uint8)
+ outputs.append(Image.fromarray(x))
+
+ if color_transfer_post_process:
+ for i in range(1, len(outputs)):
+ outputs[i] = Image.fromarray(color_transfer(np.uint8(outputs[i]), np.uint8(outputs[0])))
+
+ os.makedirs(os.path.dirname(path), exist_ok=True)
+ if imageio_backend:
+ if path.endswith("mp4"):
+ imageio.mimsave(path, outputs, fps=fps)
+ else:
+ imageio.mimsave(path, outputs, duration=(1000 * 1 / fps))
+ else:
+ if path.endswith("mp4"):
+ path = path.replace('.mp4', '.gif')
+ outputs[0].save(path, format='GIF', append_images=outputs, save_all=True, duration=100, loop=0)
+
+
+def get_image_to_video_latent(validation_image_start, validation_image_end, video_length, sample_size):
+ if validation_image_start is not None and validation_image_end is not None:
+ if type(validation_image_start) is str and os.path.isfile(validation_image_start):
+ image_start = clip_image = Image.open(validation_image_start).convert("RGB")
+ image_start = image_start.resize([sample_size[1], sample_size[0]])
+ clip_image = clip_image.resize([sample_size[1], sample_size[0]])
+ else:
+ image_start = clip_image = validation_image_start
+ image_start = [_image_start.resize([sample_size[1], sample_size[0]]) for _image_start in image_start]
+ clip_image = [_clip_image.resize([sample_size[1], sample_size[0]]) for _clip_image in clip_image]
+
+ if type(validation_image_end) is str and os.path.isfile(validation_image_end):
+ image_end = Image.open(validation_image_end).convert("RGB")
+ image_end = image_end.resize([sample_size[1], sample_size[0]])
+ else:
+ image_end = validation_image_end
+ image_end = [_image_end.resize([sample_size[1], sample_size[0]]) for _image_end in image_end]
+
+ if type(image_start) is list:
+ clip_image = clip_image[0]
+ start_video = torch.cat(
+ [torch.from_numpy(np.array(_image_start)).permute(2, 0, 1).unsqueeze(1).unsqueeze(0) for _image_start in
+ image_start],
+ dim=2
+ )
+ input_video = torch.tile(start_video[:, :, :1], [1, 1, video_length, 1, 1])
+ input_video[:, :, :len(image_start)] = start_video
+
+ input_video_mask = torch.zeros_like(input_video[:, :1])
+ input_video_mask[:, :, len(image_start):] = 255
+ else:
+ input_video = torch.tile(
+ torch.from_numpy(np.array(image_start)).permute(2, 0, 1).unsqueeze(1).unsqueeze(0),
+ [1, 1, video_length, 1, 1]
+ )
+ input_video_mask = torch.zeros_like(input_video[:, :1])
+ input_video_mask[:, :, 1:] = 255
+
+ if type(image_end) is list:
+ image_end = [_image_end.resize(image_start[0].size if type(image_start) is list else image_start.size) for
+ _image_end in image_end]
+ end_video = torch.cat(
+ [torch.from_numpy(np.array(_image_end)).permute(2, 0, 1).unsqueeze(1).unsqueeze(0) for _image_end in
+ image_end],
+ dim=2
+ )
+ input_video[:, :, -len(end_video):] = end_video
+
+ input_video_mask[:, :, -len(image_end):] = 0
+ else:
+ image_end = image_end.resize(image_start[0].size if type(image_start) is list else image_start.size)
+ input_video[:, :, -1:] = torch.from_numpy(np.array(image_end)).permute(2, 0, 1).unsqueeze(1).unsqueeze(0)
+ input_video_mask[:, :, -1:] = 0
+
+ input_video = input_video / 255
+
+ elif validation_image_start is not None:
+ if type(validation_image_start) is str and os.path.isfile(validation_image_start):
+ image_start = clip_image = Image.open(validation_image_start).convert("RGB")
+ image_start = image_start.resize([sample_size[1], sample_size[0]])
+ clip_image = clip_image.resize([sample_size[1], sample_size[0]])
+ else:
+ image_start = clip_image = validation_image_start
+ image_start = [_image_start.resize([sample_size[1], sample_size[0]]) for _image_start in image_start]
+ clip_image = [_clip_image.resize([sample_size[1], sample_size[0]]) for _clip_image in clip_image]
+ image_end = None
+
+ if type(image_start) is list:
+ clip_image = clip_image[0]
+ start_video = torch.cat(
+ [torch.from_numpy(np.array(_image_start)).permute(2, 0, 1).unsqueeze(1).unsqueeze(0) for _image_start in
+ image_start],
+ dim=2
+ )
+ input_video = torch.tile(start_video[:, :, :1], [1, 1, video_length, 1, 1])
+ input_video[:, :, :len(image_start)] = start_video
+ input_video = input_video / 255
+
+ input_video_mask = torch.zeros_like(input_video[:, :1])
+ input_video_mask[:, :, len(image_start):] = 255
+ else:
+ input_video = torch.tile(
+ torch.from_numpy(np.array(image_start)).permute(2, 0, 1).unsqueeze(1).unsqueeze(0),
+ [1, 1, video_length, 1, 1]
+ ) / 255
+ input_video_mask = torch.zeros_like(input_video[:, :1])
+ input_video_mask[:, :, 1:, ] = 255
+ else:
+ image_start = None
+ image_end = None
+ input_video = torch.zeros([1, 3, video_length, sample_size[0], sample_size[1]])
+ input_video_mask = torch.ones([1, 1, video_length, sample_size[0], sample_size[1]]) * 255
+ clip_image = None
+
+ del image_start
+ del image_end
+ gc.collect()
+
+ return input_video, input_video_mask, clip_image
\ No newline at end of file
+ 
+ 
+ 
+