bugfix

2026-03-22 16:50:47 +00:00 · 2025-11-26 14:59:30 +08:00
parent 0336551544
commit 0eead33ed7
22 changed files with 71 additions and 42 deletions
--- a/examples/wanvideo/model_inference/Wan2.2-VACE-Fun-A14B.py
+++ b/examples/wanvideo/model_inference/Wan2.2-VACE-Fun-A14B.py
@@ -1,5 +1,6 @@
 # Without VRAM Management, 80G VRAM is not enough to run this example.
 # We recommend to use `examples/wanvideo/model_inference_low_vram/Wan2.2-VACE-Fun-A14B.py`.
+# CPU Offload is enabled in this example.
 import torch
 from PIL import Image
 from diffsynth.utils.data import save_video, VideoData
@@ -7,16 +8,27 @@ from diffsynth.pipelines.wan_video import WanVideoPipeline, ModelConfig
 from modelscope import dataset_snapshot_download


+vram_config = {
+    "offload_dtype": torch.bfloat16,
+    "offload_device": "cpu",
+    "onload_dtype": torch.bfloat16,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.bfloat16,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
 pipe = WanVideoPipeline.from_pretrained(
    torch_dtype=torch.bfloat16,
    device="cuda",
    model_configs=[
-        ModelConfig(model_id="PAI/Wan2.2-VACE-Fun-A14B", origin_file_pattern="high_noise_model/diffusion_pytorch_model*.safetensors"),
-        ModelConfig(model_id="PAI/Wan2.2-VACE-Fun-A14B", origin_file_pattern="low_noise_model/diffusion_pytorch_model*.safetensors"),
-        ModelConfig(model_id="PAI/Wan2.2-VACE-Fun-A14B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth"),
-        ModelConfig(model_id="PAI/Wan2.2-VACE-Fun-A14B", origin_file_pattern="Wan2.1_VAE.pth"),
+        ModelConfig(model_id="PAI/Wan2.2-VACE-Fun-A14B", origin_file_pattern="high_noise_model/diffusion_pytorch_model*.safetensors", **vram_config),
+        ModelConfig(model_id="PAI/Wan2.2-VACE-Fun-A14B", origin_file_pattern="low_noise_model/diffusion_pytorch_model*.safetensors", **vram_config),
+        ModelConfig(model_id="PAI/Wan2.2-VACE-Fun-A14B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth", **vram_config),
+        ModelConfig(model_id="PAI/Wan2.2-VACE-Fun-A14B", origin_file_pattern="Wan2.1_VAE.pth", **vram_config),
    ],
    tokenizer_config=ModelConfig(model_id="Wan-AI/Wan2.1-T2V-1.3B", origin_file_pattern="google/umt5-xxl/"),
+    vram_limit=torch.cuda.mem_get_info("cuda")[1] / (1024 ** 3) - 2,
 )


--- a/examples/wanvideo/model_inference_low_vram/Wan2.2-Animate-14B.py
+++ b/examples/wanvideo/model_inference_low_vram/Wan2.2-Animate-14B.py
@@ -53,6 +53,7 @@ save_video(video, "video_1_Wan2.2-Animate-14B.mp4", fps=15, quality=5)
 # Replace
 snapshot_download("Wan-AI/Wan2.2-Animate-14B", allow_file_pattern="relighting_lora.ckpt", local_dir="models/Wan-AI/Wan2.2-Animate-14B")
 lora_state_dict = load_state_dict("models/Wan-AI/Wan2.2-Animate-14B/relighting_lora.ckpt", torch_dtype=torch.bfloat16, device="cuda")["state_dict"]
+lora_state_dict = {i: lora_state_dict[i].to(torch.bfloat16) for i in lora_state_dict}
 pipe.load_lora(pipe.dit, state_dict=lora_state_dict)
 input_image = Image.open("data/examples/wan/animate/replace_input_image.png")
 animate_pose_video = VideoData("data/examples/wan/animate/replace_pose_video.mp4").raw_data()[:81-4]
--- a/examples/wanvideo/model_inference_low_vram/Wan2.2-VACE-Fun-A14B.py
+++ b/examples/wanvideo/model_inference_low_vram/Wan2.2-VACE-Fun-A14B.py
@@ -1,5 +1,3 @@
-# Without VRAM Management, 80G VRAM is not enough to run this example.
-# We recommend to use `examples/wanvideo/model_inference_low_vram/Wan2.2-VACE-Fun-A14B.py`.
 import torch
 from PIL import Image
 from diffsynth.utils.data import save_video, VideoData
--- a/examples/wanvideo/model_training/full/Wan2.1-FLF2V-14B-720P.sh
+++ b/examples/wanvideo/model_training/full/Wan2.1-FLF2V-14B-720P.sh
@@ -10,4 +10,5 @@ accelerate launch --config_file examples/wanvideo/model_training/full/accelerate
  --remove_prefix_in_ckpt "pipe.dit." \
  --output_path "./models/train/Wan2.1-FLF2V-14B-720P_full" \
  --trainable_models "dit" \
-  --extra_inputs "input_image,end_image"
+  --extra_inputs "input_image,end_image" \
+  --initialize_model_on_cpu
--- a/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-480P.sh
+++ b/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-480P.sh
@@ -10,4 +10,5 @@ accelerate launch --config_file examples/wanvideo/model_training/full/accelerate
  --remove_prefix_in_ckpt "pipe.dit." \
  --output_path "./models/train/Wan2.1-I2V-14B-480P_full" \
  --trainable_models "dit" \
-  --extra_inputs "input_image"
+  --extra_inputs "input_image" \
+  --initialize_model_on_cpu
--- a/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-720P.sh
+++ b/examples/wanvideo/model_training/full/Wan2.1-I2V-14B-720P.sh
@@ -12,4 +12,5 @@ accelerate launch --config_file examples/wanvideo/model_training/full/accelerate
  --output_path "./models/train/Wan2.1-I2V-14B-720P_full" \
  --trainable_models "dit" \
  --extra_inputs "input_image" \
-  --use_gradient_checkpointing_offload
+  --use_gradient_checkpointing_offload \
+  --initialize_model_on_cpu
--- a/examples/wanvideo/model_training/full/Wan2.2-S2V-14B.sh
+++ b/examples/wanvideo/model_training/full/Wan2.2-S2V-14B.sh
@@ -1,6 +1,6 @@
 accelerate launch --config_file examples/wanvideo/model_training/full/accelerate_config_14B.yaml examples/wanvideo/model_training/train.py \
-  --dataset_base_path data/example_video_dataset/wans2v \
-  --dataset_metadata_path data/example_video_dataset/wans2v/metadata.csv \
+  --dataset_base_path data/example_video_dataset \
+  --dataset_metadata_path data/example_video_dataset/metadata_s2v.csv \
  --data_file_keys "video,input_audio,s2v_pose_video" \
  --height 448 \
  --width 832 \
--- a/examples/wanvideo/model_training/full/Wan2.2-VACE-Fun-A14B.sh
+++ b/examples/wanvideo/model_training/full/Wan2.2-VACE-Fun-A14B.sh
@@ -15,7 +15,8 @@ accelerate launch --config_file examples/wanvideo/model_training/full/accelerate
  --extra_inputs "vace_video,vace_reference_image" \
  --use_gradient_checkpointing_offload \
  --max_timestep_boundary 0.358 \
-  --min_timestep_boundary 0
+  --min_timestep_boundary 0 \
+  --initialize_model_on_cpu
 # boundary corresponds to timesteps [900, 1000]


@@ -36,5 +37,6 @@ accelerate launch --config_file examples/wanvideo/model_training/full/accelerate
  --extra_inputs "vace_video,vace_reference_image" \
  --use_gradient_checkpointing_offload \
  --max_timestep_boundary 1 \
-  --min_timestep_boundary 0.358
+  --min_timestep_boundary 0.358 \
+  --initialize_model_on_cpu
 # boundary corresponds to timesteps [0, 900]
--- a/examples/wanvideo/model_training/lora/Wan2.1-I2V-14B-720P.sh
+++ b/examples/wanvideo/model_training/lora/Wan2.1-I2V-14B-720P.sh
@@ -1,9 +1,10 @@
-accelerate launch examples/wanvideo/model_training/train.py \
+# 1*80G GPU cannot train Wan2.2-Animate-14B LoRA
+# We tested on 8*80G GPUs
+accelerate launch --config_file examples/wanvideo/model_training/full/accelerate_config_14B.yaml examples/wanvideo/model_training/train.py \
  --dataset_base_path data/example_video_dataset \
  --dataset_metadata_path data/example_video_dataset/metadata.csv \
  --height 720 \
  --width 1280 \
-  --num_frames 49 \
  --dataset_repeat 100 \
  --model_id_with_origin_paths "Wan-AI/Wan2.1-I2V-14B-720P:diffusion_pytorch_model*.safetensors,Wan-AI/Wan2.1-I2V-14B-720P:models_t5_umt5-xxl-enc-bf16.pth,Wan-AI/Wan2.1-I2V-14B-720P:Wan2.1_VAE.pth,Wan-AI/Wan2.1-I2V-14B-720P:models_clip_open-clip-xlm-roberta-large-vit-huge-14.pth" \
  --learning_rate 1e-4 \
@@ -14,4 +15,5 @@ accelerate launch examples/wanvideo/model_training/train.py \
  --lora_target_modules "q,k,v,o,ffn.0,ffn.2" \
  --lora_rank 32 \
  --extra_inputs "input_image" \
-  --use_gradient_checkpointing_offload
+  --use_gradient_checkpointing_offload \
+  --initialize_model_on_cpu
--- a/examples/wanvideo/model_training/lora/Wan2.2-S2V-14B.sh
+++ b/examples/wanvideo/model_training/lora/Wan2.2-S2V-14B.sh
@@ -1,6 +1,6 @@
 accelerate launch --config_file examples/wanvideo/model_training/full/accelerate_config_14B.yaml examples/wanvideo/model_training/train.py \
-  --dataset_base_path data/example_video_dataset/wans2v \
-  --dataset_metadata_path data/example_video_dataset/wans2v/metadata.csv \
+  --dataset_base_path data/example_video_dataset \
+  --dataset_metadata_path data/example_video_dataset/metadata_s2v.csv \
  --data_file_keys "video,input_audio,s2v_pose_video" \
  --height 448 \
  --width 832 \
--- a/examples/wanvideo/model_training/train.py
+++ b/examples/wanvideo/model_training/train.py
@@ -116,6 +116,7 @@ def wan_parser():
    parser.add_argument("--audio_processor_path", type=str, default=None, help="Path to the audio processor. If provided, the processor will be used for Wan2.2-S2V model.")
    parser.add_argument("--max_timestep_boundary", type=float, default=1.0, help="Max timestep boundary (for mixed models, e.g., Wan-AI/Wan2.2-I2V-A14B).")
    parser.add_argument("--min_timestep_boundary", type=float, default=0.0, help="Min timestep boundary (for mixed models, e.g., Wan-AI/Wan2.2-I2V-A14B).")
+    parser.add_argument("--initialize_model_on_cpu", default=False, action="store_true", help="Whether to initialize models on CPU.")
    return parser


@@ -165,7 +166,7 @@ if __name__ == "__main__":
        fp8_models=args.fp8_models,
        offload_models=args.offload_models,
        task=args.task,
-        device=accelerator.device,
+        device="cpu" if args.initialize_model_on_cpu else accelerator.device,
        max_timestep_boundary=args.max_timestep_boundary,
        min_timestep_boundary=args.min_timestep_boundary,
    )
--- a/examples/wanvideo/model_training/validate_full/Wan2.2-VACE-Fun-A14B.py
+++ b/examples/wanvideo/model_training/validate_full/Wan2.2-VACE-Fun-A14B.py
@@ -5,19 +5,29 @@ from diffsynth.core import load_state_dict
 from diffsynth.pipelines.wan_video import WanVideoPipeline, ModelConfig


+vram_config = {
+    "offload_dtype": torch.bfloat16,
+    "offload_device": "cpu",
+    "onload_dtype": torch.bfloat16,
+    "onload_device": "cpu",
+    "preparing_dtype": torch.bfloat16,
+    "preparing_device": "cuda",
+    "computation_dtype": torch.bfloat16,
+    "computation_device": "cuda",
+}
 pipe = WanVideoPipeline.from_pretrained(
    torch_dtype=torch.bfloat16,
    device="cuda",
    model_configs=[
        ModelConfig(model_id="PAI/Wan2.2-VACE-Fun-A14B", origin_file_pattern="high_noise_model/diffusion_pytorch_model*.safetensors"),
        ModelConfig(model_id="PAI/Wan2.2-VACE-Fun-A14B", origin_file_pattern="low_noise_model/diffusion_pytorch_model*.safetensors"),
-        ModelConfig(model_id="PAI/Wan2.2-VACE-Fun-A14B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth"),
-        ModelConfig(model_id="PAI/Wan2.2-VACE-Fun-A14B", origin_file_pattern="Wan2.1_VAE.pth"),
+        ModelConfig(model_id="PAI/Wan2.2-VACE-Fun-A14B", origin_file_pattern="models_t5_umt5-xxl-enc-bf16.pth", **vram_config),
+        ModelConfig(model_id="PAI/Wan2.2-VACE-Fun-A14B", origin_file_pattern="Wan2.1_VAE.pth", **vram_config),
    ],
 )
-state_dict = load_state_dict("models/train/Wan2.2-VACE-Fun-A14B_high_noise_full/epoch-1.safetensors")
+state_dict = load_state_dict("models/train/Wan2.2-VACE-Fun-A14B_high_noise_full/epoch-1.safetensors", torch_dtype=torch.bfloat16, device="cpu")
 pipe.vace.load_state_dict(state_dict)
-state_dict = load_state_dict("models/train/Wan2.2-VACE-Fun-A14B_low_noise_full/epoch-1.safetensors")
+state_dict = load_state_dict("models/train/Wan2.2-VACE-Fun-A14B_low_noise_full/epoch-1.safetensors", torch_dtype=torch.bfloat16, device="cpu")
 pipe.vace2.load_state_dict(state_dict)

 video = VideoData("data/example_video_dataset/video1_softedge.mp4", height=480, width=832)