add wan22 test

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

csrc/attn/vmoba_attn/setup.py

@@ -20,5 +20,7 @@ setup(
         "License :: OSI Approved :: Apache Software License",
     ],
     python_requires='>=3.12',
-    install_requires=[]
+    install_requires=[
+        "flash-attn >= 2.7.1",
+    ]
 )

csrc/attn/vmoba_attn/vmoba/vmoba.py

@@ -6,8 +6,16 @@ import time
 import os
 import torch
 from typing import Tuple
-from flash_attn import flash_attn_varlen_func  # Use the new flash attention function
-from flash_attn.flash_attn_interface import _flash_attn_varlen_forward, _flash_attn_varlen_backward
+try:
+    from flash_attn import flash_attn_varlen_func  # Use the new flash attention function
+    from flash_attn.flash_attn_interface import _flash_attn_varlen_forward, _flash_attn_varlen_backward
+except ImportError:
+    def _unsupported(*args, **kwargs):
+        raise ImportError("flash-attn is not installed. Please install it, e.g., `pip install flash-attn`.")
+    _flash_attn_varlen_forward = _unsupported
+    _flash_attn_varlen_backward = _unsupported
+    flash_attn_varlen_func = _unsupported
+    
 from functools import lru_cache
 from einops import rearrange
 

fastvideo/attention/backends/vmoba.py

@@ -5,7 +5,6 @@ from dataclasses import dataclass
 
 import torch
 from einops import rearrange
-from flash_attn.bert_padding import pad_input
 
 from csrc.attn.vmoba_attn.vmoba import (moba_attn_varlen, process_moba_input,
                                         process_moba_output)
@@ -134,6 +133,8 @@ class VMOBAAttentionImpl(AttentionImpl):
                  **extra_impl_args) -> None:
         self.prefix = prefix
         self.layer_idx = self._get_layer_idx(prefix)
+        from flash_attn.bert_padding import pad_input
+        self.pad_input = pad_input
 
     def _get_layer_idx(self, prefix: str) -> int | None:
         match = re.search(r"blocks\.(\d+)", prefix)
@@ -169,7 +170,6 @@ class VMOBAAttentionImpl(AttentionImpl):
             moba_chunk_size = attn_metadata.st_chunk_size
             moba_topk = attn_metadata.st_topk
 
-        # torch.distributed.breakpoint()
         query, chunk_size = process_moba_input(query,
                                                attn_metadata.patch_resolution,
                                                moba_chunk_size)
@@ -205,8 +205,8 @@ class VMOBAAttentionImpl(AttentionImpl):
             simsum_threshold=attn_metadata.moba_threshold,
             threshold_type=attn_metadata.moba_threshold_type,
         )
-        hidden_states = pad_input(hidden_states, indices_q, batch_size,
-                                  sequence_length)
+        hidden_states = self.pad_input(hidden_states, indices_q, batch_size,
+                                       sequence_length)
         hidden_states = process_moba_output(hidden_states,
                                             attn_metadata.patch_resolution,
                                             moba_chunk_size)

fastvideo/pipelines/stages/denoising.py

@@ -205,13 +205,13 @@ class DenoisingStage(PipelineStage):
 
         # (Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
         if fastvideo_args.pipeline_config.dit_config.boundary_ratio is not None:
-            boundary_timestep = fastvideo_args.pipeline_config.dit_config.boundary_ratio
-            if batch.boundary_timestep is not None:
-                logger.info("Overriding boundary timestep from %s to %s",
-                            boundary_timestep, batch.boundary_timestep)
-                boundary_timestep = batch.boundary_timestep
+            boundary_ratio = fastvideo_args.pipeline_config.dit_config.boundary_ratio
+            if batch.boundary_ratio is not None:
+                logger.info("Overriding boundary ratio from %s to %s",
+                            boundary_ratio, batch.boundary_ratio)
+                boundary_ratio = batch.boundary_ratio
 
-            boundary_timestep *= self.scheduler.num_train_timesteps
+            boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
         else:
             boundary_timestep = None
         latent_model_input = latents.to(target_dtype)

fastvideo/tests/ssim/test_inference_similarity.py

@@ -28,14 +28,14 @@ HUNYUAN_PARAMS = {
     "width": 1280,
     "num_frames": 45,
     "num_inference_steps": 6,
-    "guidance_scale": 1,
-    "embedded_cfg_scale": 6,
-    "flow_shift": 17,
+    # "guidance_scale": 1,
+    # "embedded_cfg_scale": 6,
+    # "flow_shift": 17,
     "seed": 1024,
     "sp_size": 2,
     "tp_size": 1,
-    "vae_sp": True,
-    "fps": 24,
+    # "vae_sp": True,
+    # "fps": 24,
 }
 
 WAN_T2V_PARAMS = {
@@ -45,14 +45,14 @@ WAN_T2V_PARAMS = {
     "width": 832,
     "num_frames": 45,
     "num_inference_steps": 20,
-    "guidance_scale": 3,
-    "embedded_cfg_scale": 6,
-    "flow_shift": 7.0,
+    # "guidance_scale": 3,
+    # "embedded_cfg_scale": 6,
+    # "flow_shift": 7.0,
     "seed": 1024,
     "sp_size": 2,
     "tp_size": 1,
-    "vae_sp": True,
-    "fps": 24,
+    # "vae_sp": True,
+    # "fps": 24,
     "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
     "text-encoder-precision": ("fp32",)
 }
@@ -64,18 +64,33 @@ WAN_I2V_PARAMS = {
     "width": 832,
     "num_frames": 45,
     "num_inference_steps": 6,
-    "guidance_scale": 5.0,
-    "embedded_cfg_scale": 6,
-    "flow_shift": 7.0,
+    # "guidance_scale": 5.0,
+    # "embedded_cfg_scale": 6,
+    # "flow_shift": 7.0,
     "seed": 1024,
     "sp_size": 2,
     "tp_size": 1,
-    "vae_sp": True,
-    "fps": 24,
-    "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
+    # "vae_sp": True,
+    # "fps": 24,
+    # "neg_prompt": "Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, misshapen limbs, fused fingers, still picture, messy background, three legs, many people in the background, walking backwards",
     "text-encoder-precision": ("fp32",)
 }
 
+WAN2_2_I2V_PARAMS = {
+    "num_gpus": 2,
+    "model_path": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+    "height": 480,
+    "width": 832,
+    "num_frames": 45,
+    "num_inference_steps": 20,
+    # "guidance_scale": 5.0,
+    # "embedded_cfg_scale": 6,
+    # "flow_shift": 7.0,
+    "seed": 1024,
+    "sp_size": 2,
+    "tp_size": 1,
+}
+
 MODEL_TO_PARAMS = {
     "FastHunyuan-diffusers": HUNYUAN_PARAMS,
     "Wan2.1-T2V-1.3B-Diffusers": WAN_T2V_PARAMS,
@@ -83,6 +98,7 @@ MODEL_TO_PARAMS = {
 
 I2V_MODEL_TO_PARAMS = {
     "Wan2.1-I2V-14B-480P-Diffusers": WAN_I2V_PARAMS,
+    "Wan2.2-I2V-A14B-Diffusers": WAN2_2_I2V_PARAMS,
 }
 
 TEST_PROMPTS = [
@@ -125,7 +141,6 @@ def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
 
     init_kwargs = {
         "num_gpus": BASE_PARAMS["num_gpus"],
-        "flow_shift": BASE_PARAMS["flow_shift"],
         "sp_size": BASE_PARAMS["sp_size"],
         "tp_size": BASE_PARAMS["tp_size"],
     }
@@ -142,10 +157,7 @@ def test_i2v_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
         "height": BASE_PARAMS["height"],
         "width": BASE_PARAMS["width"],
         "num_frames": BASE_PARAMS["num_frames"],
-        "guidance_scale": BASE_PARAMS["guidance_scale"],
-        "embedded_cfg_scale": BASE_PARAMS["embedded_cfg_scale"],
         "seed": BASE_PARAMS["seed"],
-        "fps": BASE_PARAMS["fps"],
     }
     if "neg_prompt" in BASE_PARAMS:
         generation_kwargs["neg_prompt"] = BASE_PARAMS["neg_prompt"]
@@ -225,7 +237,6 @@ def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
 
     init_kwargs = {
         "num_gpus": BASE_PARAMS["num_gpus"],
-        "flow_shift": BASE_PARAMS["flow_shift"],
         "sp_size": BASE_PARAMS["sp_size"],
         "tp_size": BASE_PARAMS["tp_size"],
         "dit_cpu_offload": True,
@@ -242,10 +253,7 @@ def test_inference_similarity(prompt, ATTENTION_BACKEND, model_id):
         "height": BASE_PARAMS["height"],
         "width": BASE_PARAMS["width"],
         "num_frames": BASE_PARAMS["num_frames"],
-        "guidance_scale": BASE_PARAMS["guidance_scale"],
-        "embedded_cfg_scale": BASE_PARAMS["embedded_cfg_scale"],
         "seed": BASE_PARAMS["seed"],
-        "fps": BASE_PARAMS["fps"],
     }
     if "neg_prompt" in BASE_PARAMS:
         generation_kwargs["neg_prompt"] = BASE_PARAMS["neg_prompt"]