[diffusion] model: Support TurboWan2.2-I2V SLA && add CI test for TurboWan (#16536)
This commit is contained in:
@@ -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"""
|
||||
|
||||
@@ -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 =============
|
||||
# =============================================
|
||||
|
||||
@@ -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 Attention(SLA)
|
||||
|:-----------------------------|:--------------------------------------------------|:--------------------|:--------:|:-----------------:|:---------:|:----------------------------:|:----------------------------:|
|
||||
| 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.
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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":
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user