Merge branch 'main' into speed

2026-09-04 21:47:41 +08:00 · 2026-08-26 17:22:10 +08:00 · 2026-08-24 22:26:30 +08:00 · 2026-07-24 16:59:48 +08:00 · 2026-07-24 16:30:07 +08:00 · 2025-08-01 14:45:57 +08:00
11 changed files with 1368 additions and 150 deletions
@@ -247,7 +247,7 @@ if coefficients is not None:
        coefficients, num_inference_steps, teacache_threshold, num_skip_start_steps=num_skip_start_steps, offload=teacache_offload
    )

-if cfg_skip_ratio is not None:
+if cfg_skip_ratio is not None and cfg_skip_ratio > 0:
    print(f"Enable cfg_skip_ratio {cfg_skip_ratio}.")
    pipeline.transformer.enable_cfg_skip(cfg_skip_ratio, num_inference_steps)

@@ -0,0 +1,319 @@
+import os
+import sys
+import argparse
+import numpy as np
+import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+from omegaconf import OmegaConf
+from PIL import Image
+from transformers import AutoTokenizer
+
+# 添加项目根目录到系统路径
+current_file_path = os.path.abspath(__file__)
+project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
+for project_root in project_roots:
+    sys.path.insert(0, project_root) if project_root not in sys.path else None
+
+from videox_fun.dist import set_multi_gpus_devices, shard_model
+from videox_fun.models import (AutoencoderKLWan, AutoTokenizer, CLIPModel,
+                              WanT5EncoderModel, WanTransformer3DModel)
+from videox_fun.models.cache_utils import get_teacache_coefficients
+from videox_fun.pipeline import WanI2VPipeline
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8, replace_parameters_by_name,
+                                              convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                   save_videos_grid, timer_record)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+
+def parse_args():
+    # 解析命令行参数
+    parser = argparse.ArgumentParser()
+
+    # 基础模型与推理参数
+    parser.add_argument("--GPU_memory_mode", type=str, default="sequential_cpu_offload")
+    parser.add_argument("--ulysses_degree", type=int, default=1)
+    parser.add_argument("--ring_degree", type=int, default=1)
+    parser.add_argument("--fsdp_dit", action="store_true")
+    parser.add_argument("--fsdp_text_encoder", action="store_true")
+    parser.add_argument("--compile_dit", action="store_true")
+
+    # TeaCache 参数
+    parser.add_argument("--enable_teacache", action="store_true")
+    parser.add_argument("--teacache_threshold", type=float, default=0.10)
+    parser.add_argument("--num_skip_start_steps", type=int, default=5)
+    parser.add_argument("--teacache_offload", action="store_true")
+
+    # CFG Skip 参数
+    parser.add_argument("--cfg_skip_ratio", type=float, default=0.0)
+
+    # Riflex 参数
+    parser.add_argument("--enable_riflex", action="store_true")
+    parser.add_argument("--riflex_k", type=int, default=6)
+
+    # 模型路径和配置
+    parser.add_argument("--config_path", type=str, default="config/wan2.1/wan_civitai.yaml")
+    parser.add_argument("--model_name", type=str, default="models/Diffusion_Transformer/Wan2.1-I2V-14B-480P")
+    parser.add_argument("--transformer_path", type=str, default=None)
+    parser.add_argument("--vae_path", type=str, default=None)
+    parser.add_argument("--lora_path", type=str, default=None)
+    parser.add_argument("--sample_size", nargs='+', type=int, default=[480, 832])
+    parser.add_argument("--video_length", type=int, default=81)
+    parser.add_argument("--fps", type=int, default=16)
+    parser.add_argument("--weight_dtype", type=str, default="bfloat16")
+
+    # 输入图像相关
+    parser.add_argument("--validation_image_start", type=str, default="asset/1.png")
+    parser.add_argument("--validation_image_end", type=str, default=None)
+
+    # 推理参数
+    parser.add_argument("--prompt", type=str, default="一只棕色的狗摇着头，坐在舒适房间里的浅色沙发上。在狗的后面，架子上有一幅镶框的画，周围是粉红色的花朵。房间里柔和温暖的灯光营造出舒适的氛围。")
+    parser.add_argument("--negative_prompt", type=str, default="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走")
+    parser.add_argument("--guidance_scale", type=float, default=6.0)
+    parser.add_argument("--seed", type=int, default=43)
+    parser.add_argument("--num_inference_steps", type=int, default=50)
+    parser.add_argument("--lora_weight", type=float, default=0.55)
+    parser.add_argument("--save_path", type=str, default="samples/wan-videos-i2v")
+
+    # 采样器设置
+    parser.add_argument("--sampler_name", type=str, choices=["Flow", "Flow_Unipc", "Flow_DPM++"], default="Flow_Unipc")
+    parser.add_argument("--shift", type=float, default=3.0)
+    return parser.parse_args()
+
+args = parse_args()
+
+# 获取参数
+GPU_memory_mode         = args.GPU_memory_mode
+ulysses_degree          = args.ulysses_degree
+ring_degree             = args.ring_degree
+fsdp_dit                = args.fsdp_dit
+fsdp_text_encoder       = args.fsdp_text_encoder
+compile_dit             = args.compile_dit
+enable_teacache         = args.enable_teacache
+teacache_threshold      = args.teacache_threshold
+num_skip_start_steps    = args.num_skip_start_steps
+teacache_offload        = args.teacache_offload
+cfg_skip_ratio          = args.cfg_skip_ratio
+enable_riflex           = args.enable_riflex
+riflex_k                = args.riflex_k
+config_path             = args.config_path
+model_name              = args.model_name
+transformer_path        = args.transformer_path
+vae_path                = args.vae_path
+lora_path               = args.lora_path
+sample_size             = args.sample_size
+video_length            = args.video_length
+fps                     = args.fps
+weight_dtype            = torch.bfloat16 if args.weight_dtype == "bfloat16" else torch.float16
+prompt                  = args.prompt
+negative_prompt         = args.negative_prompt
+guidance_scale          = args.guidance_scale
+seed                    = args.seed
+num_inference_steps     = args.num_inference_steps
+lora_weight             = args.lora_weight
+save_path               = args.save_path
+sampler_name            = args.sampler_name
+shift                   = args.shift
+validation_image_start  = args.validation_image_start
+validation_image_end    = args.validation_image_end
+
+device = set_multi_gpus_devices(ulysses_degree, ring_degree)
+config = OmegaConf.load(config_path)
+
+# 初始化模型组件
+transformer = WanTransformer3DModel.from_pretrained(
+    os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
+    transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
+    low_cpu_mem_usage=True,
+    torch_dtype=weight_dtype,
+)
+
+if transformer_path is not None:
+    print(f"From checkpoint: {transformer_path}")
+    if transformer_path.endswith("safetensors"):
+        from safetensors.torch import load_file, safe_open
+        state_dict = load_file(transformer_path)
+    else:
+        state_dict = torch.load(transformer_path, map_location="cpu")
+    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
+
+    m, u = transformer.load_state_dict(state_dict, strict=False)
+    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
+
+# 获取VAE
+vae = AutoencoderKLWan.from_pretrained(
+    os.path.join(model_name, config['vae_kwargs'].get('vae_subpath', 'vae')),
+    additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
+).to(weight_dtype)
+
+if vae_path is not None:
+    print(f"From checkpoint: {vae_path}")
+    if vae_path.endswith("safetensors"):
+        from safetensors.torch import load_file, safe_open
+        state_dict = load_file(vae_path)
+    else:
+        state_dict = torch.load(vae_path, map_location="cpu")
+    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
+
+    m, u = vae.load_state_dict(state_dict, strict=False)
+    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
+
+# 获取分词器
+tokenizer = AutoTokenizer.from_pretrained(
+    os.path.join(model_name, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')),
+)
+
+# 获取文本编码器
+text_encoder = WanT5EncoderModel.from_pretrained(
+    os.path.join(model_name, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
+    additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
+    low_cpu_mem_usage=True,
+    torch_dtype=weight_dtype,
+)
+text_encoder = text_encoder.eval()
+
+# 获取CLIP图像编码器
+clip_image_encoder = CLIPModel.from_pretrained(
+    os.path.join(model_name, config['image_encoder_kwargs'].get('image_encoder_subpath', 'image_encoder')),
+).to(weight_dtype)
+clip_image_encoder = clip_image_encoder.eval()
+
+# 获取调度器
+scheduler_dict = {
+    "Flow": FlowMatchEulerDiscreteScheduler,
+    "Flow_Unipc": FlowUniPCMultistepScheduler,
+    "Flow_DPM++": FlowDPMSolverMultistepScheduler,
+}
+Choosen_Scheduler = scheduler_dict[sampler_name]
+
+if sampler_name in ["Flow_Unipc", "Flow_DPM++"]:
+    config['scheduler_kwargs']['shift'] = 1
+scheduler = Choosen_Scheduler(
+    **filter_kwargs(Choosen_Scheduler, OmegaConf.to_container(config['scheduler_kwargs']))
+)
+
+# 创建Pipeline
+pipeline = WanI2VPipeline(
+    transformer=transformer,
+    vae=vae,
+    tokenizer=tokenizer,
+    text_encoder=text_encoder,
+    scheduler=scheduler,
+    clip_image_encoder=clip_image_encoder
+)
+
+# 分布式设置
+if ulysses_degree > 1 or ring_degree > 1:
+    from functools import partial
+    transformer.enable_multi_gpus_inference()
+    if fsdp_dit:
+        shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
+        pipeline.transformer = shard_fn(pipeline.transformer)
+        print("Add FSDP DIT")
+    if fsdp_text_encoder:
+        shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
+        pipeline.text_encoder = shard_fn(pipeline.text_encoder)
+        print("Add FSDP TEXT ENCODER")
+
+# 编译优化
+if compile_dit:
+    for i in range(len(pipeline.transformer.blocks)):
+        pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
+    print("Add Compile")
+
+# 内存优化策略
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
+
+for i in range(2):
+    # TeaCache配置
+    coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
+    if coefficients is not None:
+        print(f"Enable TeaCache with threshold {teacache_threshold} and skip the first {num_skip_start_steps} steps.")
+        pipeline.transformer.enable_teacache(
+            coefficients, num_inference_steps, teacache_threshold, num_skip_start_steps=num_skip_start_steps, offload=teacache_offload
+        )
+
+    # CFG跳过配置
+    if cfg_skip_ratio is not None and cfg_skip_ratio > 0:
+        print(f"Enable cfg_skip_ratio {cfg_skip_ratio}.")
+        pipeline.transformer.enable_cfg_skip(cfg_skip_ratio, num_inference_steps)
+
+    # 随机种子
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    # LoRA加载
+    if lora_path is not None:
+        pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
+
+    # 生成视频
+    with torch.no_grad():
+        video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
+        latent_frames = (video_length - 1) // vae.config.temporal_compression_ratio + 1
+
+        if enable_riflex:
+            pipeline.transformer.enable_riflex(k=riflex_k, L_test=latent_frames)
+
+        # 输入图像处理
+        input_video, input_video_mask, clip_image = get_image_to_video_latent(validation_image_start, validation_image_end, video_length=video_length, sample_size=sample_size)
+
+        # 执行推理
+        sample = pipeline(
+            prompt, 
+            num_frames=video_length,
+            negative_prompt=negative_prompt,
+            height=sample_size[0],
+            width=sample_size[1],
+            generator=generator,
+            guidance_scale=guidance_scale,
+            num_inference_steps=num_inference_steps,
+            video=input_video,
+            mask_video=input_video_mask,
+            clip_image=clip_image,
+            shift=shift,
+        ).videos
+
+    # LoRA卸载
+    if lora_path is not None:
+        pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
+
+    # 保存结果
+    def save_results():
+        if not os.path.exists(save_path):
+            os.makedirs(save_path, exist_ok=True)
+
+        index = len([path for path in os.listdir(save_path)]) + 1
+        prefix = str(index).zfill(8)
+        if video_length == 1:
+            video_path = os.path.join(save_path, prefix + ".png")
+            image = sample[0, :, 0]
+            image = image.transpose(0, 1).transpose(1, 2)
+            image = (image * 255).numpy().astype(np.uint8)
+            image = Image.fromarray(image)
+            image.save(video_path)
+        else:
+            video_path = os.path.join(save_path, prefix + ".mp4")
+            save_videos_grid(sample, video_path, fps=fps)
+
+    # 分布式保存
+    if ulysses_degree * ring_degree > 1:
+        import torch.distributed as dist
+        if dist.get_rank() == 0:
+            save_results()
+    else:
+        save_results()
@@ -0,0 +1,53 @@
+export EXCEL_FILE="./speed.xlsx"
+
+export DIT_EXCEL_COL=0 VAE_EXCEL_COL=1 TOTAL_EXCEL_COL=2
+
+# 14B 720P
+export DIT_EXCEL_ROW=1 VAE_EXCEL_ROW=1 TOTAL_EXCEL_ROW=1
+python examples/wan2.1/predict_i2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-I2V-14B-720P" \
+    --GPU_memory_mode="model_full_load_and_qfloat8" --ulysses_degree=1 --ring_degree=1 --fsdp_text_encoder --compile_dit \
+    --enable_teacache --teacache_threshold=0.30 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=2 VAE_EXCEL_ROW=2 TOTAL_EXCEL_ROW=2
+torchrun --nproc-per-node=2 examples/wan2.1/predict_i2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-I2V-14B-720P" \
+    --GPU_memory_mode="model_full_load_and_qfloat8" --ulysses_degree=2 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.30 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=3 VAE_EXCEL_ROW=3 TOTAL_EXCEL_ROW=3
+torchrun --nproc-per-node=4 examples/wan2.1/predict_i2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-I2V-14B-720P" \
+    --GPU_memory_mode="model_full_load_and_qfloat8" --ulysses_degree=4 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.30 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=4 VAE_EXCEL_ROW=4 TOTAL_EXCEL_ROW=4
+torchrun --nproc-per-node=8 examples/wan2.1/predict_i2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-I2V-14B-720P" \
+    --GPU_memory_mode="model_full_load_and_qfloat8" --ulysses_degree=8 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.30 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+# 14B 480P
+export DIT_EXCEL_ROW=5 VAE_EXCEL_ROW=5 TOTAL_EXCEL_ROW=5
+python examples/wan2.1/predict_i2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-I2V-14B-720P" \
+    --GPU_memory_mode="model_full_load_and_qfloat8" --ulysses_degree=1 --ring_degree=1 --fsdp_text_encoder --compile_dit \
+    --enable_teacache --teacache_threshold=0.30 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=6 VAE_EXCEL_ROW=6 TOTAL_EXCEL_ROW=6
+torchrun --nproc-per-node=2 examples/wan2.1/predict_i2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-I2V-14B-720P" \
+    --GPU_memory_mode="model_full_load_and_qfloat8" --ulysses_degree=2 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.30 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=7 VAE_EXCEL_ROW=7 TOTAL_EXCEL_ROW=7
+torchrun --nproc-per-node=4 examples/wan2.1/predict_i2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-I2V-14B-720P" \
+    --GPU_memory_mode="model_full_load_and_qfloat8" --ulysses_degree=4 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.30 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=8 VAE_EXCEL_ROW=8 TOTAL_EXCEL_ROW=8
+torchrun --nproc-per-node=8 examples/wan2.1/predict_i2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-I2V-14B-720P" \
+    --GPU_memory_mode="model_full_load_and_qfloat8" --ulysses_degree=8 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.30 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
@@ -234,7 +234,7 @@ if coefficients is not None:
        coefficients, num_inference_steps, teacache_threshold, num_skip_start_steps=num_skip_start_steps, offload=teacache_offload
    )

-if cfg_skip_ratio is not None:
+if cfg_skip_ratio is not None and cfg_skip_ratio > 0:
    print(f"Enable cfg_skip_ratio {cfg_skip_ratio}.")
    pipeline.transformer.enable_cfg_skip(cfg_skip_ratio, num_inference_steps)

@@ -0,0 +1,303 @@
+import os
+import sys
+
+import numpy as np
+import torch
+import argparse
+from diffusers import FlowMatchEulerDiscreteScheduler
+from omegaconf import OmegaConf
+from PIL import Image
+
+current_file_path = os.path.abspath(__file__)
+project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
+for project_root in project_roots:
+    sys.path.insert(0, project_root) if project_root not in sys.path else None
+
+from videox_fun.dist import set_multi_gpus_devices, shard_model
+from videox_fun.models import (AutoencoderKLWan, WanT5EncoderModel, AutoTokenizer,
+                              WanTransformer3DModel)
+from videox_fun.models.cache_utils import get_teacache_coefficients
+from videox_fun.pipeline import WanPipeline
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8, replace_parameters_by_name,
+                                              convert_weight_dtype_wrapper)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                   save_videos_grid, timer_record)
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Video Generation with Wan2.1-Fun")
+    parser.add_argument("--GPU_memory_mode", type=str, default="sequential_cpu_offload",
+                        choices=["model_full_load", "model_full_load_and_qfloat8", "model_cpu_offload", 
+                                 "model_cpu_offload_and_qfloat8", "sequential_cpu_offload"],
+                        help="GPU memory optimization mode.")
+    parser.add_argument("--ulysses_degree", type=int, default=1,
+                        help="Ulysses parallelism degree.")
+    parser.add_argument("--ring_degree", type=int, default=1,
+                        help="Ring parallelism degree.")
+    parser.add_argument("--fsdp_dit", action="store_true",
+                        help="Use FSDP for transformer to save GPU memory.")
+    parser.add_argument("--fsdp_text_encoder", action="store_true",
+                        help="Use FSDP for text encoder to save GPU memory.")
+    parser.add_argument("--compile_dit", action="store_true",
+                        help="Compile transformer for fixed resolution speedup.")
+    parser.add_argument("--enable_teacache", action="store_true",
+                        help="Enable TeaCache optimization.")
+    parser.add_argument("--teacache_threshold", type=float, default=0.10,
+                        help="TeaCache threshold for step caching.")
+    parser.add_argument("--num_skip_start_steps", type=int, default=5,
+                        help="Number of steps to skip TeaCache at inference start.")
+    parser.add_argument("--teacache_offload", action="store_true",
+                        help="Offload TeaCache tensors to CPU.")
+    parser.add_argument("--cfg_skip_ratio", type=float, default=0.0,
+                        help="CFG skip ratio for inference.")
+    parser.add_argument("--enable_riflex", action="store_true",
+                        help="Enable Riflex frequency optimization.")
+    parser.add_argument("--riflex_k", type=int, default=6,
+                        help="Intrinsic frequency index for Riflex.")
+    parser.add_argument("--config_path", type=str, default="config/wan2.1/wan_civitai.yaml",
+                        help="Path to model config file.")
+    parser.add_argument("--model_name", type=str, default="models/Diffusion_Transformer/Wan2.1-T2V-1.3B",
+                        help="Path to model directory.")
+    parser.add_argument("--sampler_name", type=str, default="Flow_Unipc",
+                        choices=["Flow", "Flow_Unipc", "Flow_DPM++"],
+                        help="Sampler type for video generation.")
+    parser.add_argument("--shift", type=float, default=3.0,
+                        help="Noise schedule shift parameter for Flow_Unipc/Flow_DPM++.")
+    parser.add_argument("--transformer_path", type=str, default=None,
+                        help="Path to pre-trained transformer checkpoint.")
+    parser.add_argument("--vae_path", type=str, default=None,
+                        help="Path to pre-trained VAE checkpoint.")
+    parser.add_argument("--lora_path", type=str, default=None,
+                        help="Path to LoRA weights.")
+    parser.add_argument("--sample_size", nargs=2, type=int, default=[480, 832],
+                        help="Sample size [height, width].")
+    parser.add_argument("--video_length", type=int, default=81,
+                        help="Number of frames in the video.")
+    parser.add_argument("--fps", type=int, default=16,
+                        help="Frames per second for output video.")
+    parser.add_argument("--weight_dtype", type=str, default="bfloat16",
+                        choices=["float16", "bfloat16"],
+                        help="Weight data type (float16 or bfloat16).")
+    parser.add_argument("--prompt", type=str, default="一只棕色的狗摇着头，坐在舒适房间里的浅色沙发上。在狗的后面，架子上有一幅镶框的画，周围是粉红色的花朵。房间里柔和温暖的灯光营造出舒适的氛围。",
+                        help="Text prompt for video generation.")
+    parser.add_argument("--negative_prompt", type=str, default="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
+                        help="Negative prompt for video generation.")
+    parser.add_argument("--guidance_scale", type=float, default=6.0,
+                        help="Classifier-free guidance scale.")
+    parser.add_argument("--seed", type=int, default=43,
+                        help="Random seed for reproducibility.")
+    parser.add_argument("--num_inference_steps", type=int, default=50,
+                        help="Number of inference steps.")
+    parser.add_argument("--lora_weight", type=float, default=0.55,
+                        help="LoRA weight scaling factor.")
+    parser.add_argument("--save_path", type=str, default="samples/wan-videos-t2v",
+                        help="Directory to save generated videos.")
+    return parser.parse_args()
+
+args = parse_args()
+
+# 将 argparse 参数映射到原有变量
+GPU_memory_mode     = args.GPU_memory_mode
+ulysses_degree      = args.ulysses_degree
+ring_degree         = args.ring_degree
+fsdp_dit            = args.fsdp_dit
+fsdp_text_encoder   = args.fsdp_text_encoder
+compile_dit         = args.compile_dit
+enable_teacache     = args.enable_teacache
+teacache_threshold  = args.teacache_threshold
+num_skip_start_steps = args.num_skip_start_steps
+teacache_offload    = args.teacache_offload
+cfg_skip_ratio      = args.cfg_skip_ratio
+enable_riflex       = args.enable_riflex
+riflex_k            = args.riflex_k
+config_path         = args.config_path
+model_name          = args.model_name
+sampler_name        = args.sampler_name
+shift               = args.shift
+transformer_path    = args.transformer_path
+vae_path            = args.vae_path
+lora_path           = args.lora_path
+sample_size         = args.sample_size
+video_length        = args.video_length
+fps                 = args.fps
+weight_dtype        = torch.bfloat16 if args.weight_dtype == "bfloat16" else torch.float16
+prompt              = args.prompt
+negative_prompt     = args.negative_prompt
+guidance_scale      = args.guidance_scale
+seed                = args.seed
+num_inference_steps = args.num_inference_steps
+lora_weight         = args.lora_weight
+save_path           = args.save_path
+
+device = set_multi_gpus_devices(ulysses_degree, ring_degree)
+config = OmegaConf.load(config_path)
+
+transformer = WanTransformer3DModel.from_pretrained(
+    os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_subpath', 'transformer')),
+    transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
+    low_cpu_mem_usage=True if not fsdp_dit else False,
+    torch_dtype=weight_dtype,
+)
+
+if transformer_path is not None:
+    print(f"From checkpoint: {transformer_path}")
+    if transformer_path.endswith("safetensors"):
+        from safetensors.torch import load_file, safe_open
+        state_dict = load_file(transformer_path)
+    else:
+        state_dict = torch.load(transformer_path, map_location="cpu")
+    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
+
+    m, u = transformer.load_state_dict(state_dict, strict=False)
+    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
+
+# Get Vae
+vae = AutoencoderKLWan.from_pretrained(
+    os.path.join(model_name, config['vae_kwargs'].get('vae_subpath', 'vae')),
+    additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
+).to(weight_dtype)
+
+if vae_path is not None:
+    print(f"From checkpoint: {vae_path}")
+    if vae_path.endswith("safetensors"):
+        from safetensors.torch import load_file, safe_open
+        state_dict = load_file(vae_path)
+    else:
+        state_dict = torch.load(vae_path, map_location="cpu")
+    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
+
+    m, u = vae.load_state_dict(state_dict, strict=False)
+    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
+
+# Get Tokenizer
+tokenizer = AutoTokenizer.from_pretrained(
+    os.path.join(model_name, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')),
+)
+
+# Get Text encoder
+text_encoder = WanT5EncoderModel.from_pretrained(
+    os.path.join(model_name, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
+    additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
+    low_cpu_mem_usage=True,
+    torch_dtype=weight_dtype,
+)
+
+# Get Scheduler
+Choosen_Scheduler = scheduler_dict = {
+    "Flow": FlowMatchEulerDiscreteScheduler,
+    "Flow_Unipc": FlowUniPCMultistepScheduler,
+    "Flow_DPM++": FlowDPMSolverMultistepScheduler,
+}[sampler_name]
+if sampler_name == "Flow_Unipc" or sampler_name == "Flow_DPM++":
+    config['scheduler_kwargs']['shift'] = 1
+scheduler = Choosen_Scheduler(
+    **filter_kwargs(Choosen_Scheduler, OmegaConf.to_container(config['scheduler_kwargs']))
+)
+
+# Get Pipeline
+pipeline = WanPipeline(
+    transformer=transformer,
+    vae=vae,
+    tokenizer=tokenizer,
+    text_encoder=text_encoder,
+    scheduler=scheduler,
+)
+if ulysses_degree > 1 or ring_degree > 1:
+    from functools import partial
+    transformer.enable_multi_gpus_inference()
+    if fsdp_dit:
+        shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
+        pipeline.transformer = shard_fn(pipeline.transformer)
+        print("Add FSDP DIT")
+    if fsdp_text_encoder:
+        shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
+        pipeline.text_encoder = shard_fn(pipeline.text_encoder)
+        print("Add FSDP TEXT ENCODER")
+
+if compile_dit:
+    for i in range(len(pipeline.transformer.blocks)):
+        pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
+    print("Add Compile")
+
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
+
+for i in range(2):
+    coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
+    if coefficients is not None:
+        print(f"Enable TeaCache with threshold {teacache_threshold} and skip the first {num_skip_start_steps} steps.")
+        pipeline.transformer.enable_teacache(
+            coefficients, num_inference_steps, teacache_threshold, num_skip_start_steps=num_skip_start_steps, offload=teacache_offload
+        )
+
+    if cfg_skip_ratio is not None and cfg_skip_ratio > 0:
+        print(f"Enable cfg_skip_ratio {cfg_skip_ratio}.")
+        pipeline.transformer.enable_cfg_skip(cfg_skip_ratio, num_inference_steps)
+
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    if lora_path is not None:
+        pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
+
+    with torch.no_grad():
+        video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
+        latent_frames = (video_length - 1) // vae.config.temporal_compression_ratio + 1
+
+        if enable_riflex:
+            pipeline.transformer.enable_riflex(k = riflex_k, L_test = latent_frames)
+
+        sample = pipeline(
+            prompt, 
+            num_frames = video_length,
+            negative_prompt = negative_prompt,
+            height      = sample_size[0],
+            width       = sample_size[1],
+            generator   = generator,
+            guidance_scale = guidance_scale,
+            num_inference_steps = num_inference_steps,
+            shift = shift,
+        ).videos
+
+    if lora_path is not None:
+        pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
+
+    def save_results():
+        if not os.path.exists(save_path):
+            os.makedirs(save_path, exist_ok=True)
+
+        index = len([path for path in os.listdir(save_path)]) + 1
+        prefix = str(index).zfill(8)
+        if video_length == 1:
+            video_path = os.path.join(save_path, prefix + ".png")
+
+            image = sample[0, :, 0]
+            image = image.transpose(0, 1).transpose(1, 2)
+            image = (image * 255).numpy().astype(np.uint8)
+            image = Image.fromarray(image)
+            image.save(video_path)
+        else:
+            video_path = os.path.join(save_path, prefix + ".mp4")
+            save_videos_grid(sample, video_path, fps=fps)
+
+    if ulysses_degree * ring_degree > 1:
+        import torch.distributed as dist
+        if dist.get_rank() == 0:
+            save_results()
+    else:
+        save_results()
@@ -0,0 +1,103 @@
+export EXCEL_FILE="./speed.xlsx"
+
+export DIT_EXCEL_COL=0 VAE_EXCEL_COL=1 TOTAL_EXCEL_COL=2
+
+# 1.3B 720P
+export DIT_EXCEL_ROW=1 VAE_EXCEL_ROW=1 TOTAL_EXCEL_ROW=1
+python examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=1 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=2 VAE_EXCEL_ROW=2 TOTAL_EXCEL_ROW=2
+torchrun --nproc-per-node=2 examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=2 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=3 VAE_EXCEL_ROW=3 TOTAL_EXCEL_ROW=3
+torchrun --nproc-per-node=4 examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=4 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=4 VAE_EXCEL_ROW=4 TOTAL_EXCEL_ROW=4
+torchrun --nproc-per-node=8 examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=4 --ring_degree=2 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+# 1.3B 480P
+export DIT_EXCEL_ROW=5 VAE_EXCEL_ROW=5 TOTAL_EXCEL_ROW=5
+python examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=1 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=6 VAE_EXCEL_ROW=6 TOTAL_EXCEL_ROW=6
+torchrun --nproc-per-node=2 examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=2 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=7 VAE_EXCEL_ROW=7 TOTAL_EXCEL_ROW=7
+torchrun --nproc-per-node=4 examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=4 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=8 VAE_EXCEL_ROW=8 TOTAL_EXCEL_ROW=8
+torchrun --nproc-per-node=8 examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-1.3B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=4 --ring_degree=2 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
+
+# 14B 720P
+export DIT_EXCEL_ROW=9 VAE_EXCEL_ROW=9 TOTAL_EXCEL_ROW=9
+python examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=1 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.15 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=10 VAE_EXCEL_ROW=10 TOTAL_EXCEL_ROW=10
+torchrun --nproc-per-node=2 examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=2 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.15 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=11 VAE_EXCEL_ROW=11 TOTAL_EXCEL_ROW=11
+torchrun --nproc-per-node=4 examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=4 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.15 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=12 VAE_EXCEL_ROW=12 TOTAL_EXCEL_ROW=12
+torchrun --nproc-per-node=8 examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=8 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.15 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+# 14B 480P
+export DIT_EXCEL_ROW=13 VAE_EXCEL_ROW=13 TOTAL_EXCEL_ROW=13
+python examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=1 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.15 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=14 VAE_EXCEL_ROW=14 TOTAL_EXCEL_ROW=14
+torchrun --nproc-per-node=2 examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=2 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.15 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=15 VAE_EXCEL_ROW=15 TOTAL_EXCEL_ROW=15
+torchrun --nproc-per-node=4 examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=4 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.15 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=16 VAE_EXCEL_ROW=16 TOTAL_EXCEL_ROW=16
+torchrun --nproc-per-node=8 examples/wan2.1/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.1-T2V-14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=8 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.15 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
@@ -0,0 +1,352 @@
+import argparse
+import os
+import sys
+
+import numpy as np
+import torch
+from diffusers import FlowMatchEulerDiscreteScheduler
+from omegaconf import OmegaConf
+from PIL import Image
+
+current_file_path = os.path.abspath(__file__)
+project_roots = [os.path.dirname(current_file_path), os.path.dirname(os.path.dirname(current_file_path)), os.path.dirname(os.path.dirname(os.path.dirname(current_file_path)))]
+for project_root in project_roots:
+    sys.path.insert(0, project_root) if project_root not in sys.path else None
+
+from videox_fun.dist import set_multi_gpus_devices, shard_model
+from videox_fun.models import (AutoencoderKLWan, AutoTokenizer,
+                               Wan2_2Transformer3DModel, WanT5EncoderModel,
+                               WanTransformer3DModel)
+from videox_fun.models.cache_utils import get_teacache_coefficients
+from videox_fun.pipeline import Wan2_2Pipeline, WanPipeline
+from videox_fun.utils.fm_solvers import FlowDPMSolverMultistepScheduler
+from videox_fun.utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from videox_fun.utils.fp8_optimization import (convert_model_weight_to_float8,
+                                               convert_weight_dtype_wrapper,
+                                               replace_parameters_by_name)
+from videox_fun.utils.lora_utils import merge_lora, unmerge_lora
+from videox_fun.utils.utils import (filter_kwargs, get_image_to_video_latent,
+                                    save_videos_grid, timer_record)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Video Generation with Wan2.2-Fun")
+    parser.add_argument("--GPU_memory_mode", type=str, default="sequential_cpu_offload",
+                        choices=["model_full_load", "model_full_load_and_qfloat8", "model_cpu_offload", 
+                                 "model_cpu_offload_and_qfloat8", "sequential_cpu_offload"],
+                        help="GPU memory optimization mode.")
+    parser.add_argument("--ulysses_degree", type=int, default=1,
+                        help="Ulysses parallelism degree.")
+    parser.add_argument("--ring_degree", type=int, default=1,
+                        help="Ring parallelism degree.")
+    parser.add_argument("--fsdp_dit", action="store_true",
+                        help="Use FSDP for transformer to save GPU memory.")
+    parser.add_argument("--fsdp_text_encoder", action="store_true",
+                        help="Use FSDP for text encoder to save GPU memory.")
+    parser.add_argument("--compile_dit", action="store_true",
+                        help="Compile transformer for fixed resolution speedup.")
+    parser.add_argument("--enable_teacache", action="store_true",
+                        help="Enable TeaCache optimization.")
+    parser.add_argument("--teacache_threshold", type=float, default=0.10,
+                        help="TeaCache threshold for step caching.")
+    parser.add_argument("--num_skip_start_steps", type=int, default=5,
+                        help="Number of steps to skip TeaCache at inference start.")
+    parser.add_argument("--teacache_offload", action="store_true",
+                        help="Offload TeaCache tensors to CPU.")
+    parser.add_argument("--cfg_skip_ratio", type=float, default=0.0,
+                        help="CFG skip ratio for inference.")
+    parser.add_argument("--enable_riflex", action="store_true",
+                        help="Enable Riflex frequency optimization.")
+    parser.add_argument("--riflex_k", type=int, default=6,
+                        help="Intrinsic frequency index for Riflex.")
+    parser.add_argument("--config_path", type=str, default="config/wan2.2/wan_civitai_t2v.yaml",
+                        help="Path to model config file.")
+    parser.add_argument("--model_name", type=str, default="models/Diffusion_Transformer/Wan2.2-T2V-A14B",
+                        help="Path to model directory.")
+    parser.add_argument("--sampler_name", type=str, default="Flow_Unipc",
+                        choices=["Flow", "Flow_Unipc", "Flow_DPM++"],
+                        help="Sampler type for video generation.")
+    parser.add_argument("--shift", type=float, default=3.0,
+                        help="Noise schedule shift parameter for Flow_Unipc/Flow_DPM++.")
+    parser.add_argument("--transformer_path", type=str, default=None,
+                        help="Path to pre-trained transformer checkpoint.")
+    parser.add_argument("--transformer_high_path", type=str, default=None,
+                        help="Path to pre-trained high noise transformer checkpoint.")
+    parser.add_argument("--vae_path", type=str, default=None,
+                        help="Path to pre-trained VAE checkpoint.")
+    parser.add_argument("--lora_path", type=str, default=None,
+                        help="Path to LoRA weights.")
+    parser.add_argument("--lora_high_path", type=str, default=None,
+                        help="Path to high noise LoRA weights.")
+    parser.add_argument("--sample_size", nargs=2, type=int, default=[480, 832],
+                        help="Sample size [height, width].")
+    parser.add_argument("--video_length", type=int, default=81,
+                        help="Number of frames in the video.")
+    parser.add_argument("--fps", type=int, default=16,
+                        help="Frames per second for output video.")
+    parser.add_argument("--weight_dtype", type=str, default="bfloat16",
+                        choices=["float16", "bfloat16"],
+                        help="Weight data type (float16 or bfloat16).")
+    parser.add_argument("--prompt", type=str, default="一只棕色的狗摇着头，坐在舒适房间里的浅色沙发上。在狗的后面，架子上有一幅镶框的画，周围是粉红色的花朵。房间里柔和温暖的灯光营造出舒适的氛围。",
+                        help="Text prompt for video generation.")
+    parser.add_argument("--negative_prompt", type=str, default="色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走",
+                        help="Negative prompt for video generation.")
+    parser.add_argument("--guidance_scale", type=float, default=6.0,
+                        help="Classifier-free guidance scale.")
+    parser.add_argument("--seed", type=int, default=43,
+                        help="Random seed for reproducibility.")
+    parser.add_argument("--num_inference_steps", type=int, default=50,
+                        help="Number of inference steps.")
+    parser.add_argument("--lora_weight", type=float, default=0.55,
+                        help="LoRA weight scaling factor.")
+    parser.add_argument("--lora_high_weight", type=float, default=0.55,
+                        help="High noise LoRA weight scaling factor.")
+    parser.add_argument("--save_path", type=str, default="samples/wan-videos-t2v",
+                        help="Directory to save generated videos.")
+    return parser.parse_args()
+
+args = parse_args()
+
+# 将 argparse 参数映射到原有变量
+GPU_memory_mode     = args.GPU_memory_mode
+ulysses_degree      = args.ulysses_degree
+ring_degree         = args.ring_degree
+fsdp_dit            = args.fsdp_dit
+fsdp_text_encoder   = args.fsdp_text_encoder
+compile_dit         = args.compile_dit
+enable_teacache     = args.enable_teacache
+teacache_threshold  = args.teacache_threshold
+num_skip_start_steps = args.num_skip_start_steps
+teacache_offload    = args.teacache_offload
+cfg_skip_ratio      = args.cfg_skip_ratio
+enable_riflex       = args.enable_riflex
+riflex_k            = args.riflex_k
+config_path         = args.config_path
+model_name          = args.model_name
+sampler_name        = args.sampler_name
+shift               = args.shift
+transformer_path    = args.transformer_path
+transformer_high_path = args.transformer_high_path
+vae_path            = args.vae_path
+lora_path           = args.lora_path
+lora_high_path      = args.lora_high_path
+sample_size         = args.sample_size
+video_length        = args.video_length
+fps                 = args.fps
+weight_dtype        = torch.bfloat16 if args.weight_dtype == "bfloat16" else torch.float16
+prompt              = args.prompt
+negative_prompt     = args.negative_prompt
+guidance_scale      = args.guidance_scale
+seed                = args.seed
+num_inference_steps = args.num_inference_steps
+lora_weight         = args.lora_weight
+lora_high_weight    = args.lora_high_weight
+save_path           = args.save_path
+
+device = set_multi_gpus_devices(ulysses_degree, ring_degree)
+config = OmegaConf.load(config_path)
+boundary = config['transformer_additional_kwargs'].get('boundary', 0.875)
+
+transformer = Wan2_2Transformer3DModel.from_pretrained(
+    os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_low_noise_model_subpath', 'transformer')),
+    transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
+    low_cpu_mem_usage=True,
+    torch_dtype=weight_dtype,
+)
+
+transformer_2 = Wan2_2Transformer3DModel.from_pretrained(
+    os.path.join(model_name, config['transformer_additional_kwargs'].get('transformer_high_noise_model_subpath', 'transformer')),
+    transformer_additional_kwargs=OmegaConf.to_container(config['transformer_additional_kwargs']),
+    low_cpu_mem_usage=True,
+    torch_dtype=weight_dtype,
+)
+
+if transformer_path is not None:
+    print(f"From checkpoint: {transformer_path}")
+    if transformer_path.endswith("safetensors"):
+        from safetensors.torch import load_file, safe_open
+        state_dict = load_file(transformer_path)
+    else:
+        state_dict = torch.load(transformer_path, map_location="cpu")
+    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
+
+    m, u = transformer.load_state_dict(state_dict, strict=False)
+    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
+
+if transformer_high_path is not None:
+    print(f"From checkpoint: {transformer_high_path}")
+    if transformer_high_path.endswith("safetensors"):
+        from safetensors.torch import load_file, safe_open
+        state_dict = load_file(transformer_high_path)
+    else:
+        state_dict = torch.load(transformer_high_path, map_location="cpu")
+    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
+
+    m, u = transformer_2.load_state_dict(state_dict, strict=False)
+    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
+
+# Get Vae
+vae = AutoencoderKLWan.from_pretrained(
+    os.path.join(model_name, config['vae_kwargs'].get('vae_subpath', 'vae')),
+    additional_kwargs=OmegaConf.to_container(config['vae_kwargs']),
+).to(weight_dtype)
+
+if vae_path is not None:
+    print(f"From checkpoint: {vae_path}")
+    if vae_path.endswith("safetensors"):
+        from safetensors.torch import load_file, safe_open
+        state_dict = load_file(vae_path)
+    else:
+        state_dict = torch.load(vae_path, map_location="cpu")
+    state_dict = state_dict["state_dict"] if "state_dict" in state_dict else state_dict
+
+    m, u = vae.load_state_dict(state_dict, strict=False)
+    print(f"missing keys: {len(m)}, unexpected keys: {len(u)}")
+
+# Get Tokenizer
+tokenizer = AutoTokenizer.from_pretrained(
+    os.path.join(model_name, config['text_encoder_kwargs'].get('tokenizer_subpath', 'tokenizer')),
+)
+
+# Get Text encoder
+text_encoder = WanT5EncoderModel.from_pretrained(
+    os.path.join(model_name, config['text_encoder_kwargs'].get('text_encoder_subpath', 'text_encoder')),
+    additional_kwargs=OmegaConf.to_container(config['text_encoder_kwargs']),
+    low_cpu_mem_usage=True,
+    torch_dtype=weight_dtype,
+)
+
+# Get Scheduler
+Choosen_Scheduler = scheduler_dict = {
+    "Flow": FlowMatchEulerDiscreteScheduler,
+    "Flow_Unipc": FlowUniPCMultistepScheduler,
+    "Flow_DPM++": FlowDPMSolverMultistepScheduler,
+}[sampler_name]
+if sampler_name == "Flow_Unipc" or sampler_name == "Flow_DPM++":
+    config['scheduler_kwargs']['shift'] = 1
+scheduler = Choosen_Scheduler(
+    **filter_kwargs(Choosen_Scheduler, OmegaConf.to_container(config['scheduler_kwargs']))
+)
+
+# Get Pipeline
+pipeline = Wan2_2Pipeline(
+    transformer=transformer,
+    transformer_2=transformer_2,
+    vae=vae,
+    tokenizer=tokenizer,
+    text_encoder=text_encoder,
+    scheduler=scheduler,
+)
+if ulysses_degree > 1 or ring_degree > 1:
+    from functools import partial
+    transformer.enable_multi_gpus_inference()
+    transformer_2.enable_multi_gpus_inference()
+    if fsdp_dit:
+        shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
+        pipeline.transformer = shard_fn(pipeline.transformer)
+        pipeline.transformer_2 = shard_fn(pipeline.transformer_2)
+        print("Add FSDP DIT")
+    if fsdp_text_encoder:
+        shard_fn = partial(shard_model, device_id=device, param_dtype=weight_dtype)
+        pipeline.text_encoder = shard_fn(pipeline.text_encoder)
+        print("Add FSDP TEXT ENCODER")
+
+if compile_dit:
+    for i in range(len(pipeline.transformer.blocks)):
+        pipeline.transformer.blocks[i] = torch.compile(pipeline.transformer.blocks[i])
+    for i in range(len(pipeline.transformer_2.blocks)):
+        pipeline.transformer_2.blocks[i] = torch.compile(pipeline.transformer_2.blocks[i])
+    print("Add Compile")
+
+if GPU_memory_mode == "sequential_cpu_offload":
+    replace_parameters_by_name(transformer, ["modulation",], device=device)
+    replace_parameters_by_name(transformer_2, ["modulation",], device=device)
+    transformer.freqs = transformer.freqs.to(device=device)
+    transformer_2.freqs = transformer_2.freqs.to(device=device)
+    pipeline.enable_sequential_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_cpu_offload":
+    pipeline.enable_model_cpu_offload(device=device)
+elif GPU_memory_mode == "model_full_load_and_qfloat8":
+    convert_model_weight_to_float8(transformer, exclude_module_name=["modulation",], device=device)
+    convert_model_weight_to_float8(transformer_2, exclude_module_name=["modulation",], device=device)
+    convert_weight_dtype_wrapper(transformer, weight_dtype)
+    convert_weight_dtype_wrapper(transformer_2, weight_dtype)
+    pipeline.to(device=device)
+else:
+    pipeline.to(device=device)
+
+for i in range(2):
+    coefficients = get_teacache_coefficients(model_name) if enable_teacache else None
+    if coefficients is not None:
+        print(f"Enable TeaCache with threshold {teacache_threshold} and skip the first {num_skip_start_steps} steps.")
+        pipeline.transformer.enable_teacache(
+            coefficients, num_inference_steps, teacache_threshold, num_skip_start_steps=num_skip_start_steps, offload=teacache_offload
+        )
+        pipeline.transformer_2.share_teacache(transformer=pipeline.transformer)
+
+    if cfg_skip_ratio is not None:
+        print(f"Enable cfg_skip_ratio {cfg_skip_ratio}.")
+        pipeline.transformer.enable_cfg_skip(cfg_skip_ratio, num_inference_steps)
+        pipeline.transformer_2.share_cfg_skip(transformer=pipeline.transformer)
+
+    generator = torch.Generator(device=device).manual_seed(seed)
+
+    if lora_path is not None:
+        pipeline = merge_lora(pipeline, lora_path, lora_weight, device=device)
+        pipeline = merge_lora(pipeline, lora_high_path, lora_high_weight, device=device, sub_transformer_name="transformer_2")
+
+    with torch.no_grad():
+        video_length = int((video_length - 1) // vae.config.temporal_compression_ratio * vae.config.temporal_compression_ratio) + 1 if video_length != 1 else 1
+        latent_frames = (video_length - 1) // vae.config.temporal_compression_ratio + 1
+
+        if enable_riflex:
+            pipeline.transformer.enable_riflex(k = riflex_k, L_test = latent_frames)
+            pipeline.transformer_2.enable_riflex(k = riflex_k, L_test = latent_frames)
+
+        sample = pipeline(
+            prompt, 
+            num_frames = video_length,
+            negative_prompt = negative_prompt,
+            height      = sample_size[0],
+            width       = sample_size[1],
+            generator   = generator,
+            guidance_scale = guidance_scale,
+            num_inference_steps = num_inference_steps,
+            boundary = boundary,
+            shift = shift,
+        ).videos
+
+    if lora_path is not None:
+        pipeline = unmerge_lora(pipeline, lora_path, lora_weight, device=device)
+        pipeline = unmerge_lora(pipeline, lora_high_path, lora_high_weight, device=device, sub_transformer_name="transformer_2")
+
+    def save_results():
+        if not os.path.exists(save_path):
+            os.makedirs(save_path, exist_ok=True)
+
+        index = len([path for path in os.listdir(save_path)]) + 1
+        prefix = str(index).zfill(8)
+        if video_length == 1:
+            video_path = os.path.join(save_path, prefix + ".png")
+
+            image = sample[0, :, 0]
+            image = image.transpose(0, 1).transpose(1, 2)
+            image = (image * 255).numpy().astype(np.uint8)
+            image = Image.fromarray(image)
+            image.save(video_path)
+        else:
+            video_path = os.path.join(save_path, prefix + ".mp4")
+            save_videos_grid(sample, video_path, fps=fps)
+
+    if ulysses_degree * ring_degree > 1:
+        import torch.distributed as dist
+        if dist.get_rank() == 0:
+            save_results()
+    else:
+        save_results()
@@ -0,0 +1,53 @@
+export EXCEL_FILE="./speed.xlsx"
+
+export DIT_EXCEL_COL=0 VAE_EXCEL_COL=1 TOTAL_EXCEL_COL=2
+
+# 14B 720P
+export DIT_EXCEL_ROW=1 VAE_EXCEL_ROW=1 TOTAL_EXCEL_ROW=1
+python examples/wan2.2/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.2-T2V-A14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=1 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=2 VAE_EXCEL_ROW=2 TOTAL_EXCEL_ROW=2
+torchrun --nproc-per-node=2 examples/wan2.2/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.2-T2V-A14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=2 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=3 VAE_EXCEL_ROW=3 TOTAL_EXCEL_ROW=3
+torchrun --nproc-per-node=4 examples/wan2.2/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.2-T2V-A14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=4 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=4 VAE_EXCEL_ROW=4 TOTAL_EXCEL_ROW=4
+torchrun --nproc-per-node=8 examples/wan2.2/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.2-T2V-A14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=4 --ring_degree=2 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=5 \
+    --sample_size 720 1280 --num_inference_steps=40 
+
+# 14B 480P
+export DIT_EXCEL_ROW=5 VAE_EXCEL_ROW=5 TOTAL_EXCEL_ROW=5
+python examples/wan2.2/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.2-T2V-A14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=1 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=6 VAE_EXCEL_ROW=6 TOTAL_EXCEL_ROW=6
+torchrun --nproc-per-node=2 examples/wan2.2/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.2-T2V-A14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=2 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=7 VAE_EXCEL_ROW=7 TOTAL_EXCEL_ROW=7
+torchrun --nproc-per-node=4 examples/wan2.2/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.2-T2V-A14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=4 --ring_degree=1 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
+
+export DIT_EXCEL_ROW=8 VAE_EXCEL_ROW=8 TOTAL_EXCEL_ROW=8
+torchrun --nproc-per-node=8 examples/wan2.2/predict_t2v_speed.py --model_name="models/Diffusion_Transformer/Wan2.2-T2V-A14B" \
+    --GPU_memory_mode="model_full_load" --ulysses_degree=4 --ring_degree=2 --fsdp_text_encoder --fsdp_dit \
+    --enable_teacache --teacache_threshold=0.10 --num_skip_start_steps=2 --cfg_skip_ratio=0.25 --shift=3 \
+    --sample_size 480 832 --num_inference_steps=40 
@@ -17,6 +17,7 @@ from ..models import (AutoencoderKLWan, AutoTokenizer,
 from ..utils.fm_solvers import (FlowDPMSolverMultistepScheduler,
                                get_sampling_sigmas)
 from ..utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from ..utils.utils import timer_record

 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name

@@ -381,6 +382,7 @@ class WanPipeline(DiffusionPipeline):
    def interrupt(self):
        return self._interrupt

+    @timer_record("TOTAL")
    @torch.no_grad()
    @replace_example_docstring(EXAMPLE_DOC_STRING)
    def __call__(
@@ -513,57 +515,67 @@ class WanPipeline(DiffusionPipeline):
        # 7. Denoising loop
        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
        self.transformer.num_inference_steps = num_inference_steps
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                self.transformer.current_steps = i

-                if self.interrupt:
-                    continue
+        @timer_record("DIT")
+        def dit_forward(latents):
+            with self.progress_bar(total=num_inference_steps) as progress_bar:
+                for i, t in enumerate(timesteps):
+                    self.transformer.current_steps = i

-                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)
+                    if self.interrupt:
+                        continue

-                # 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), torch.cuda.device(device=device):
-                    noise_pred = self.transformer(
-                        x=latent_model_input,
-                        context=in_prompt_embeds,
-                        t=timestep,
-                        seq_len=seq_len,
-                    )
+                    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)

-                # 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)
+                    # 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), torch.cuda.device(device=device):
+                        noise_pred = self.transformer(
+                            x=latent_model_input,
+                            context=in_prompt_embeds,
+                            t=timestep,
+                            seq_len=seq_len,
+                        )

-                # 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]
+                    # 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)

-                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)
+                    # 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]

-                    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 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)

-                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)
+                        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 output_type == "pil":
-            video = self.decode_latents(latents)
-            video = torch.from_numpy(video)
-        else:
-            video = latents
+                    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)
+            return latents
+
+        latents = dit_forward(latents)
+
+        @timer_record("VAE")
+        def vae_forward(latents):
+            if output_type == "pil":
+                video = self.decode_latents(latents)
+                video = torch.from_numpy(video)
+            else:
+                video = latents
+            return video
+        video = vae_forward(latents)

        # Offload all models
        self.maybe_free_model_hooks()
@@ -17,6 +17,7 @@ from ..models import (AutoencoderKLWan, AutoTokenizer,
 from ..utils.fm_solvers import (FlowDPMSolverMultistepScheduler,
                                get_sampling_sigmas)
 from ..utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from ..utils.utils import timer_record

 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name

@@ -383,6 +384,7 @@ class Wan2_2Pipeline(DiffusionPipeline):
    def interrupt(self):
        return self._interrupt

+    @timer_record("TOTAL")
    @torch.no_grad()
    @replace_example_docstring(EXAMPLE_DOC_STRING)
    def __call__(
@@ -516,69 +518,79 @@ class Wan2_2Pipeline(DiffusionPipeline):
        # 7. Denoising loop
        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
        self.transformer.num_inference_steps = num_inference_steps
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                self.transformer.current_steps = i

-                if self.interrupt:
-                    continue
+        @timer_record("DIT")
+        def dit_forward(latents):
+            with self.progress_bar(total=num_inference_steps) as progress_bar:
+                for i, t in enumerate(timesteps):
+                    self.transformer.current_steps = i

-                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)
+                    if self.interrupt:
+                        continue

-                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
-                timestep = t.expand(latent_model_input.shape[0])
-                
-                if self.transformer_2 is not None:
-                    if t >= boundary * self.scheduler.config.num_train_timesteps:
-                        local_transformer = self.transformer_2
+                    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])
+                    
+                    if self.transformer_2 is not None:
+                        if t >= boundary * self.scheduler.config.num_train_timesteps:
+                            local_transformer = self.transformer_2
+                        else:
+                            local_transformer = self.transformer
                    else:
                        local_transformer = self.transformer
-                else:
-                    local_transformer = self.transformer

-                # predict noise model_output
-                with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=device):
-                    noise_pred = local_transformer(
-                        x=latent_model_input,
-                        context=in_prompt_embeds,
-                        t=timestep,
-                        seq_len=seq_len,
-                    )
+                    # predict noise model_output
+                    with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=device):
+                        noise_pred = local_transformer(
+                            x=latent_model_input,
+                            context=in_prompt_embeds,
+                            t=timestep,
+                            seq_len=seq_len,
+                        )

-                # perform guidance
-                if do_classifier_free_guidance:
-                    if self.transformer_2 is not None and (isinstance(self.guidance_scale, (list, tuple))):
-                        sample_guide_scale = self.guidance_scale[1] if t >= self.transformer_2.config.boundary * self.scheduler.config.num_train_timesteps else self.guidance_scale[0]
-                    else:
-                        sample_guide_scale = self.guidance_scale
-                    noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
-                    noise_pred = noise_pred_uncond + sample_guide_scale * (noise_pred_text - noise_pred_uncond)
+                    # perform guidance
+                    if do_classifier_free_guidance:
+                        if self.transformer_2 is not None and (isinstance(self.guidance_scale, (list, tuple))):
+                            sample_guide_scale = self.guidance_scale[1] if t >= self.transformer_2.config.boundary * self.scheduler.config.num_train_timesteps else self.guidance_scale[0]
+                        else:
+                            sample_guide_scale = self.guidance_scale
+                        noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+                        noise_pred = noise_pred_uncond + sample_guide_scale * (noise_pred_text - noise_pred_uncond)

-                # compute the previous noisy sample x_t -> x_t-1
-                latents = self.scheduler.step(noise_pred, t, latents, **extra_step_kwargs, return_dict=False)[0]
+                    # 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)
+                    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)
+                        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 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)
+            return latents

-        if output_type == "pil":
-            video = self.decode_latents(latents)
-            video = torch.from_numpy(video)
-        else:
-            video = latents
+        latents = dit_forward(latents)
+
+        @timer_record("VAE")
+        def vae_forward(latents):
+            if output_type == "pil":
+                video = self.decode_latents(latents)
+                video = torch.from_numpy(video)
+            else:
+                video = latents
+            return video
+        video = vae_forward(latents)

        # Offload all models
        self.maybe_free_model_hooks()
@@ -24,6 +24,7 @@ from ..models import (AutoencoderKLWan, AutoTokenizer, CLIPModel,
 from ..utils.fm_solvers import (FlowDPMSolverMultistepScheduler,
                                get_sampling_sigmas)
 from ..utils.fm_solvers_unipc import FlowUniPCMultistepScheduler
+from ..utils.utils import timer_record

 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name

@@ -466,6 +467,7 @@ class WanFunInpaintPipeline(DiffusionPipeline):
    def interrupt(self):
        return self._interrupt

+    @timer_record("TOTAL")
    @torch.no_grad()
    @replace_example_docstring(EXAMPLE_DOC_STRING)
    def __call__(
@@ -658,71 +660,80 @@ class WanFunInpaintPipeline(DiffusionPipeline):
        # 7. Denoising loop
        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
        self.transformer.num_inference_steps = num_inference_steps
-        with self.progress_bar(total=num_inference_steps) as progress_bar:
-            for i, t in enumerate(timesteps):
-                self.transformer.current_steps = i

-                if self.interrupt:
-                    continue
+        @timer_record("DIT")
+        def dit_forward(latents):
+            with self.progress_bar(total=num_inference_steps) as progress_bar:
+                for i, t in enumerate(timesteps):
+                    self.transformer.current_steps = i

-                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)
+                    if self.interrupt:
+                        continue

-                if init_video is not None:
-                    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) 
+                    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)

-                clip_context_input = (
-                    torch.cat([clip_context] * 2) if do_classifier_free_guidance else clip_context
-                )
+                    if init_video is not None:
+                        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) 

-                # 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), torch.cuda.device(device=device):
-                    noise_pred = self.transformer(
-                        x=latent_model_input,
-                        context=in_prompt_embeds,
-                        t=timestep,
-                        seq_len=seq_len,
-                        y=y,
-                        clip_fea=clip_context_input,
+                    clip_context_input = (
+                        torch.cat([clip_context] * 2) if do_classifier_free_guidance else 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)
+                    # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                    timestep = t.expand(latent_model_input.shape[0])
                    
-                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)
+                    # predict noise model_output
+                    with torch.cuda.amp.autocast(dtype=weight_dtype), torch.cuda.device(device=device):
+                        noise_pred = self.transformer(
+                            x=latent_model_input,
+                            context=in_prompt_embeds,
+                            t=timestep,
+                            seq_len=seq_len,
+                            y=y,
+                            clip_fea=clip_context_input,
+                        )

-        if output_type == "pil":
-            video = self.decode_latents(latents)
-            video = torch.from_numpy(video)
-        else:
-            video = latents
+                    # 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)
+            return latents
+
+        latents = dit_forward(latents)
+
+        @timer_record("VAE")
+        def vae_forward(latents):
+            if output_type == "pil":
+                video = self.decode_latents(latents)
+                video = torch.from_numpy(video)
+            else:
+                video = latents
+            return video
+        video = vae_forward(latents)
        # Offload all models
        self.maybe_free_model_hooks()
Author	SHA1	Message	Date
bubbliiiing	d1c581bb17	Merge branch 'main' into speed	2026-09-04 21:47:41 +08:00
bubbliiiing	d31de92872	Merge branch 'main' into speed	2026-08-26 17:22:10 +08:00
bubbliiiing	35bfe679dc	Merge branch 'main' into speed	2026-08-24 22:26:30 +08:00
bubbliiiing	d555ea056b	Update speed test	2026-07-24 16:59:48 +08:00
bubbliiiing	b6e5a32b3e	Merge branch 'main' into speed	2026-07-24 16:30:07 +08:00
bubbliiiing	67ab552885	Merge branch 'main' into speed	2025-08-01 14:45:57 +08:00
bubbliiiing	413f831ced	Update Wan2.2 Speed	2025-07-31 19:28:17 +08:00
bubbliiiing	f4ffd1c64b	Update Wan2.2 Speed	2025-07-31 15:55:06 +08:00
bubbliiiing	ed42563dd3	Update Wan2.2 Speed	2025-07-31 15:00:46 +08:00
bubbliiiing	8ee48e420c	Update Wan2.2 Speed	2025-07-31 14:31:08 +08:00
bubbliiiing	b757b8edae	Merge branch 'main' into speed	2025-07-31 14:17:20 +08:00
bubbliiiing	924dd8528a	Update auto_tile_batch_size args	2025-07-30 19:39:31 +08:00
bubbliiiing	e258d4158b	Update ui	2025-07-30 14:20:47 +08:00
bubbliiiing	48c8288323	Update api and ui	2025-07-30 12:53:45 +08:00
bubbliiiing	1aacfe6bca	Update Predict Code	2025-07-30 11:04:49 +08:00
bubbliiiing	288e88eb46	Update Wan2.2 UI	2025-07-30 10:54:08 +08:00
bubbliiiing	dfb8ca04a7	Update Wan2.2 UI	2025-07-30 10:52:02 +08:00
bubbliiiing	499eef12c3	Update Wan2.2 UI	2025-07-30 10:48:59 +08:00
bubbliiiing	fa8623b22b	Update speed sh	2025-06-03 02:42:30 +00:00
bubbliiiing	61fb59833f	Update i2v speed	2025-06-02 16:24:49 +00:00
bubbliiiing	4b0b009fd3	Update i2v speed	2025-06-02 15:46:37 +00:00
bubbliiiing	518c2bdc7e	Merge branch 'main' into speed	2025-05-29 03:35:06 +00:00
bubbliiiing	6188e66fc4	Merge branch 'main' into speed	2025-05-26 10:41:24 +00:00
bubbliiiing	7c9d655822	Merge branch 'main' into speed	2025-05-19 03:13:34 +00:00
bubbliiiing	873f622dae	Speed Test	2025-05-16 02:42:35 +00:00