[diffusion] model: Support TurboWan2.2-I2V SLA && add CI test for TurboWan (#16536)

This commit is contained in:
HuangJi
2026-01-12 13:55:38 +08:00
committed by GitHub
parent 38b30c7b56
commit feb39f7768
9 changed files with 247 additions and 25 deletions

View File

@@ -148,6 +148,18 @@ class WanI2V720PConfig(WanI2V480PConfig):
flow_shift: float | None = 5.0
@dataclass
class TurboWanI2V720Config(WanI2V720PConfig):
flow_shift: float | None = 8.0
dmd_denoising_steps: list[int] | None = field(
default_factory=lambda: [996, 932, 852, 608]
)
boundary_ratio: float | None = 0.9
def __post_init__(self) -> None:
self.dit_config.boundary_ratio = self.boundary_ratio
@dataclass
class FastWan2_1_T2V_480P_Config(WanT2V480PConfig):
"""Base configuration for FastWan T2V 1.3B 480P pipeline architecture with DMD"""

View File

@@ -273,6 +273,14 @@ class Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParams):
)
@dataclass
class Turbo_Wan2_2_I2V_A14B_SamplingParam(Wan2_2_Base_SamplingParams):
guidance_scale: float = 3.5 # high_noise
guidance_scale_2: float = 3.5 # low_noise
num_inference_steps: int = 4
fps: int = 16
# =============================================
# ============= Causal Self-Forcing =============
# =============================================

View File

@@ -16,19 +16,23 @@ default parameters when initializing and generating videos.
### Video Generation Models
| Model Name | Hugging Face Model ID | Resolutions | TeaCache | Sliding Tile Attn | Sage Attn | Video Sparse Attention (VSA) |
|:-----------------------------|:--------------------------------------------------|:--------------------|:--------:|:-----------------:|:---------:|:----------------------------:|
| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480p | ⭕ | ⭕ | ⭕ | ✅ |
| FastWan2.2 TI2V 5B Full Attn | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720p | ⭕ | ⭕ | ⭕ | ✅ |
| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720p | ⭕ | ⭕ | ✅ | ⭕ |
| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480p<br>720p | ❌ | ❌ | ✅ | ⭕ |
| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480p<br>720p | ❌ | ❌ | ✅ | ⭕ |
| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720×1280<br>544×960 | ❌ | ✅ | ✅ | ⭕ |
| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720×1280<br>544×960 | ❌ | ✅ | ✅ | ⭕ |
| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ |
| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480p, 720p | ✅ | ✅ | ✅ | ⭕ |
| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ |
| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720p | ✅ | ✅ | ✅ | ⭕ |
| Model Name | Hugging Face Model ID | Resolutions | TeaCache | Sliding Tile Attn | Sage Attn | Video Sparse Attention (VSA) | Sparse Linear AttentionSLA
|:-----------------------------|:--------------------------------------------------|:--------------------|:--------:|:-----------------:|:---------:|:----------------------------:|:----------------------------:|
| FastWan2.1 T2V 1.3B | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers` | 480p | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
| FastWan2.2 TI2V 5B Full Attn | `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` | 720p | ⭕ | ⭕ | ⭕ | ✅ | ❌ |
| Wan2.2 TI2V 5B | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 720p | ⭕ | ⭕ | ✅ | ⭕ | ❌ |
| Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480p<br>720p | ❌ | ❌ | ✅ | ⭕ | ❌ |
| Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480p<br>720p | ❌ | ❌ | ✅ | ⭕ | ❌ |
| HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720×1280<br>544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ |
| FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720×1280<br>544×960 | ❌ | ✅ | ✅ | ⭕ | ❌ |
| Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ |
| Wan2.1 T2V 14B | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | 480p, 720p | ✅ | ✅ | ✅ | ⭕ | ❌ |
| Wan2.1 I2V 480P | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ |
| Wan2.1 I2V 720P | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | 720p | ✅ | ✅ | ✅ | ⭕ | ❌ |
| TurboWan2.1 T2V 1.3B | `IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ |
| TurboWan2.1 T2V 14B | `IPostYellow/TurboWan2.1-T2V-14B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ |
| TurboWan2.1 T2V 14B 720P | `IPostYellow/TurboWan2.1-T2V-14B-720P-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ |
| TurboWan2.2 I2V A14B | `IPostYellow/TurboWan2.2-I2V-A14B-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ |
**Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.

View File

@@ -49,6 +49,7 @@ from sglang.multimodal_gen.configs.pipeline_configs.qwen_image import (
from sglang.multimodal_gen.configs.pipeline_configs.wan import (
FastWan2_1_T2V_480P_Config,
FastWan2_2_TI2V_5B_Config,
TurboWanI2V720Config,
TurboWanT2V480PConfig,
Wan2_2_I2V_A14B_Config,
Wan2_2_T2V_A14B_Config,
@@ -67,6 +68,7 @@ from sglang.multimodal_gen.configs.sample.qwenimage import (
)
from sglang.multimodal_gen.configs.sample.wan import (
FastWanT2V480PConfig,
Turbo_Wan2_2_I2V_A14B_SamplingParam,
Wan2_1_Fun_1_3B_InP_SamplingParams,
Wan2_2_I2V_A14B_SamplingParam,
Wan2_2_T2V_A14B_SamplingParam,
@@ -426,6 +428,7 @@ def _register_configs():
pipeline_config_cls=TurboWanT2V480PConfig,
hf_model_paths=[
"IPostYellow/TurboWan2.1-T2V-14B-Diffusers",
"IPostYellow/TurboWan2.1-T2V-14B-720P-Diffusers",
],
)
register_configs(
@@ -443,6 +446,13 @@ def _register_configs():
"Wan-AI/Wan2.1-I2V-14B-720P-Diffusers",
],
)
register_configs(
sampling_param_cls=Turbo_Wan2_2_I2V_A14B_SamplingParam,
pipeline_config_cls=TurboWanI2V720Config,
hf_model_paths=[
"IPostYellow/TurboWan2.2-I2V-A14B-Diffusers",
],
)
register_configs(
sampling_param_cls=Wan2_1_Fun_1_3B_InP_SamplingParams,
pipeline_config_cls=WanI2V480PConfig,

View File

@@ -68,14 +68,17 @@ class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase):
tokenizers=[self.get_module("tokenizer")],
),
)
self.add_stage(
stage_name="image_encoding_stage",
stage=ImageEncodingStage(
image_encoder=self.get_module("image_encoder"),
image_processor=self.get_module("image_processor"),
),
)
if (
self.get_module("image_encoder") is not None
and self.get_module("image_processor") is not None
):
self.add_stage(
stage_name="image_encoding_stage",
stage=ImageEncodingStage(
image_encoder=self.get_module("image_encoder"),
image_processor=self.get_module("image_processor"),
),
)
self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
@@ -102,6 +105,7 @@ class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase):
stage=DmdDenoisingStage(
transformer=self.get_module("transformer"),
scheduler=self.get_module("scheduler"),
transformer_2=self.get_module("transformer_2"),
),
)

View File

@@ -26,8 +26,10 @@ class DmdDenoisingStage(DenoisingStage):
Denoising stage for DMD.
"""
def __init__(self, transformer, scheduler) -> None:
super().__init__(transformer, scheduler)
def __init__(self, transformer, scheduler, transformer_2=None) -> None:
super().__init__(
transformer=transformer, scheduler=scheduler, transformer_2=transformer_2
)
self.scheduler = FlowMatchEulerDiscreteScheduler(shift=8.0)
def _preprocess_sp_latents(self, batch: Req, server_args: ServerArgs):
@@ -103,6 +105,20 @@ class DmdDenoisingStage(DenoisingStage):
timings=batch.timings,
perf_dump_path_provided=batch.perf_dump_path is not None,
):
t_int = int(t.item())
if self.transformer_2 is not None:
current_model, current_guidance_scale = (
self._select_and_manage_model(
t_int=t_int,
boundary_timestep=self._handle_boundary_ratio(
server_args, batch
),
server_args=server_args,
batch=batch,
)
)
else:
current_model = self.transformer
# Expand latents for I2V
noise_latents = latents.clone()
latent_model_input = latents.to(target_dtype)
@@ -143,7 +159,7 @@ class DmdDenoisingStage(DenoisingStage):
forward_batch=batch,
):
# Run transformer
pred_noise = self.transformer(
pred_noise = current_model(
hidden_states=latent_model_input.permute(0, 2, 1, 3, 4),
timestep=t_expand,
guidance=guidance_expand,
@@ -202,3 +218,76 @@ class DmdDenoisingStage(DenoisingStage):
)
return batch
def _select_and_manage_model(
self,
t_int: int,
boundary_timestep: float | None,
server_args: ServerArgs,
batch: Req,
):
if boundary_timestep is None or t_int >= boundary_timestep:
# High-noise stage
current_model = self.transformer
model_to_offload = self.transformer_2
current_guidance_scale = batch.guidance_scale
else:
# Low-noise stage
current_model = self.transformer_2
model_to_offload = self.transformer
current_guidance_scale = batch.guidance_scale_2
self._manage_device_placement(current_model, model_to_offload, server_args)
assert current_model is not None, "The model for the current step is not set."
return current_model, current_guidance_scale
def _manage_device_placement(
self,
model_to_use: torch.nn.Module,
model_to_offload: torch.nn.Module | None,
server_args: ServerArgs,
):
"""
Manages the offload / load behavior of dit
"""
if not server_args.dit_cpu_offload:
return
# Offload the unused model if it's on CUDA
if (
model_to_offload is not None
and next(model_to_offload.parameters()).device.type == "cuda"
):
model_to_offload.to("cpu")
# Load the model to use if it's on CPU
if (
model_to_use is not None
and next(model_to_use.parameters()).device.type == "cpu"
):
model_to_use.to(get_local_torch_device())
def _handle_boundary_ratio(
self,
server_args,
batch,
):
"""
(Wan2.2) Calculate timestep to switch from high noise expert to low noise expert
"""
boundary_ratio = server_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
if boundary_ratio is not None:
boundary_timestep = boundary_ratio * self.scheduler.num_train_timesteps
else:
boundary_timestep = None
return boundary_timestep

View File

@@ -804,6 +804,27 @@
"expected_avg_denoise_ms": 260.76,
"expected_median_denoise_ms": 247.84
},
"turbo_wan2_1_t2v_1.3b": {
"stages_ms": {
"InputValidationStage": 0.06,
"TextEncodingStage": 2508.95,
"ConditioningStage": 0.04,
"TimestepPreparationStage": 73.51,
"LatentPreparationStage": 1.34,
"DmdDenoisingStage": 1285.25,
"DecodingStage": 805.04,
"per_frame_generation": null
},
"denoise_step_ms": {
"0": 897.62,
"1": 126.04,
"2": 126.52,
"3": 128.26
},
"expected_e2e_ms": 4686.66,
"expected_avg_denoise_ms": 319.61,
"expected_median_denoise_ms": 127.39
},
"wan2_2_ti2v_5b": {
"stages_ms": {
"InputValidationStage": 96.27,
@@ -1092,6 +1113,28 @@
"expected_avg_denoise_ms": 2831.00,
"expected_median_denoise_ms": 1600.09
},
"turbo_wan2_2_i2v_a14b_2gpu": {
"stages_ms": {
"InputValidationStage": 25.01,
"TextEncodingStage": 5198.6,
"ConditioningStage": 0.04,
"TimestepPreparationStage": 56.26,
"LatentPreparationStage": 1.4,
"ImageVAEEncodingStage": 1001.89,
"DmdDenoisingStage": 4487.79,
"DecodingStage": 821.01,
"per_frame_generation": null
},
"denoise_step_ms": {
"0": 3042.56,
"1": 485.88,
"2": 477.59,
"3": 475.58
},
"expected_e2e_ms": 11605.97,
"expected_avg_denoise_ms": 1120.4,
"expected_median_denoise_ms": 481.74
},
"wan2_1_i2v_14b_480P_2gpu": {
"stages_ms": {
"InputValidationStage": 38.23,

View File

@@ -1078,7 +1078,11 @@ def get_generate_fn(
prompt=sampling_params.prompt,
size=sampling_params.output_size,
seconds=video_seconds,
extra_body={"reference_url": sampling_params.image_path},
extra_body={
"reference_url": sampling_params.image_path,
"fps": sampling_params.fps,
"num_frames": sampling_params.num_frames,
},
)
def generate_text_image_to_video(case_id, client) -> str:
@@ -1102,6 +1106,10 @@ def get_generate_fn(
size=output_size,
seconds=video_seconds,
input_reference=fh,
extra_body={
"fps": sampling_params.fps,
"num_frames": sampling_params.num_frames,
},
)
if modality == "video":

View File

@@ -25,6 +25,7 @@ from dataclasses import dataclass
from pathlib import Path
from typing import Sequence
from sglang.multimodal_gen.runtime.platforms import current_platform
from sglang.multimodal_gen.runtime.utils.perf_logger import RequestPerfRecord
@@ -300,6 +301,15 @@ TI2V_sampling_params = DiffusionSamplingParams(
direct_url_test=True,
)
TURBOWAN_I2V_sampling_params = DiffusionSamplingParams(
prompt="The man in the picture slowly turns his head, his expression enigmatic and otherworldly. The camera performs a slow, cinematic dolly out, focusing on his face. Moody lighting, neon signs glowing in the background, shallow depth of field.",
image_path="https://is1-ssl.mzstatic.com/image/thumb/Music114/v4/5f/fa/56/5ffa56c2-ea1f-7a17-6bad-192ff9b6476d/825646124206.jpg/600x600bb.jpg",
direct_url_test=True,
output_size="960x960",
num_frames=4,
fps=4,
)
# All test cases with clean default values
# To test different models, simply add more DiffusionCase entries
ONE_GPU_CASES_A: list[DiffusionTestCase] = [
@@ -496,6 +506,23 @@ ONE_GPU_CASES_B: list[DiffusionTestCase] = [
),
]
# Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536.
if not current_platform.is_hip():
ONE_GPU_CASES_B.append(
DiffusionTestCase(
"turbo_wan2_1_t2v_1.3b",
DiffusionServerArgs(
model_path="IPostYellow/TurboWan2.1-T2V-1.3B-Diffusers",
modality="video",
warmup=0,
custom_validator="video",
),
DiffusionSamplingParams(
prompt=T2V_PROMPT,
),
)
)
TWO_GPU_CASES_A = [
DiffusionTestCase(
"wan2_2_i2v_a14b_2gpu",
@@ -551,6 +578,23 @@ TWO_GPU_CASES_A = [
),
]
# Skip turbowan because Triton requires 81920 shared memory, but AMD only has 65536.
if not current_platform.is_hip():
TWO_GPU_CASES_A.append(
DiffusionTestCase(
"turbo_wan2_2_i2v_a14b_2gpu",
DiffusionServerArgs(
model_path="IPostYellow/TurboWan2.2-I2V-A14B-Diffusers",
modality="video",
warmup=0,
custom_validator="video",
num_gpus=2,
tp_size=2,
),
TURBOWAN_I2V_sampling_params,
)
)
TWO_GPU_CASES_B = [
DiffusionTestCase(
"wan2_1_i2v_14b_480P_2gpu",