From 33d1aeb07f2da7298273cd85b99b250a3660a2dc Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 12 Nov 2025 12:39:34 +0800 Subject: [PATCH] diffusion: refactor task type of models (#13118) --- .../multimodal_gen/configs/pipelines/base.py | 22 +++-- .../multimodal_gen/configs/pipelines/flux.py | 8 +- .../configs/pipelines/qwen_image.py | 6 +- .../multimodal_gen/configs/pipelines/wan.py | 7 +- .../entrypoints/diffusion_generator.py | 2 +- .../runtime/entrypoints/utils.py | 93 +++++++++---------- .../runtime/pipelines/stages/denoising.py | 22 +++-- .../pipelines/stages/input_validation.py | 5 +- .../multimodal_gen/runtime/server_args.py | 2 +- 9 files changed, 96 insertions(+), 71 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/pipelines/base.py b/python/sglang/multimodal_gen/configs/pipelines/base.py index 09a01f526..8e6039bd7 100644 --- a/python/sglang/multimodal_gen/configs/pipelines/base.py +++ b/python/sglang/multimodal_gen/configs/pipelines/base.py @@ -4,7 +4,7 @@ import json from collections.abc import Callable from dataclasses import asdict, dataclass, field, fields -from enum import Enum +from enum import Enum, auto from typing import Any import torch @@ -28,6 +28,19 @@ from sglang.multimodal_gen.utils import ( logger = init_logger(__name__) +# NOTE: possible duplication with DataType, WorkloadType +# this may focus on the model's original ability +class ModelTaskType(Enum): + I2V = auto() # Image to Video + T2V = auto() # Text to Video + TI2V = auto() # Text and Image to Video + T2I = auto() # Text to Image + I2I = auto() # Image to Image + + def is_image_task(self): + return self == ModelTaskType.T2I or self == ModelTaskType.I2I + + class STA_Mode(str, Enum): """STA (Sliding Tile Attention) modes.""" @@ -51,11 +64,11 @@ def postprocess_text(output: BaseEncoderOutput, _text_inputs) -> torch.tensor: class PipelineConfig: """Base configuration for all pipeline architectures.""" + task_type: ModelTaskType + model_path: str = "" pipeline_config_path: str | None = None - is_image_gen: bool = False - # generation parameters # controls the timestep embedding generation should_use_guidance: bool = True @@ -113,9 +126,6 @@ class PipelineConfig: dmd_denoising_steps: list[int] | None = field(default=None) # Wan2.2 TI2V parameters - ti2v_task: bool = False - i2v_task: bool = False - ti2i_task: bool = False boundary_ratio: float | None = None # Compilation diff --git a/python/sglang/multimodal_gen/configs/pipelines/flux.py b/python/sglang/multimodal_gen/configs/pipelines/flux.py index a5348ec25..a187c4f15 100644 --- a/python/sglang/multimodal_gen/configs/pipelines/flux.py +++ b/python/sglang/multimodal_gen/configs/pipelines/flux.py @@ -13,7 +13,11 @@ from sglang.multimodal_gen.configs.models.encoders import ( T5Config, ) from sglang.multimodal_gen.configs.models.vaes.flux import FluxVAEConfig -from sglang.multimodal_gen.configs.pipelines.base import PipelineConfig, preprocess_text +from sglang.multimodal_gen.configs.pipelines.base import ( + ModelTaskType, + PipelineConfig, + preprocess_text, +) from sglang.multimodal_gen.configs.pipelines.hunyuan import ( clip_postprocess_text, clip_preprocess_text, @@ -29,7 +33,7 @@ class FluxPipelineConfig(PipelineConfig): # FIXME: duplicate with SamplingParams.guidance_scale? embedded_cfg_scale: float = 3.5 - is_image_gen: bool = True + task_type: ModelTaskType = ModelTaskType.T2I vae_tiling: bool = False diff --git a/python/sglang/multimodal_gen/configs/pipelines/qwen_image.py b/python/sglang/multimodal_gen/configs/pipelines/qwen_image.py index 1b16c6fd4..eda4cb20e 100644 --- a/python/sglang/multimodal_gen/configs/pipelines/qwen_image.py +++ b/python/sglang/multimodal_gen/configs/pipelines/qwen_image.py @@ -10,7 +10,7 @@ from sglang.multimodal_gen.configs.models import DiTConfig, EncoderConfig, VAECo from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig from sglang.multimodal_gen.configs.models.encoders.qwen_image import Qwen2_5VLConfig from sglang.multimodal_gen.configs.models.vaes.qwenimage import QwenImageVAEConfig -from sglang.multimodal_gen.configs.pipelines.base import PipelineConfig +from sglang.multimodal_gen.configs.pipelines.base import ModelTaskType, PipelineConfig def _extract_masked_hidden(hidden_states: torch.Tensor, mask: torch.Tensor): @@ -64,7 +64,7 @@ def _pack_latents(latents, batch_size, num_channels_latents, height, width): class QwenImagePipelineConfig(PipelineConfig): should_use_guidance: bool = False - is_image_gen: bool = True + task_type: ModelTaskType = ModelTaskType.T2I vae_tiling: bool = False @@ -194,7 +194,7 @@ class QwenImagePipelineConfig(PipelineConfig): class QwenImageEditPipelineConfig(QwenImagePipelineConfig): - ti2i_task = True + task_type: ModelTaskType = ModelTaskType.I2I def prepare_pos_cond_kwargs(self, batch, device, rotary_emb, dtype): # TODO: lots of duplications here diff --git a/python/sglang/multimodal_gen/configs/pipelines/wan.py b/python/sglang/multimodal_gen/configs/pipelines/wan.py index d98e5fe86..d5efa2d64 100644 --- a/python/sglang/multimodal_gen/configs/pipelines/wan.py +++ b/python/sglang/multimodal_gen/configs/pipelines/wan.py @@ -14,7 +14,7 @@ from sglang.multimodal_gen.configs.models.encoders import ( T5Config, ) from sglang.multimodal_gen.configs.models.vaes import WanVAEConfig -from sglang.multimodal_gen.configs.pipelines.base import PipelineConfig +from sglang.multimodal_gen.configs.pipelines.base import ModelTaskType, PipelineConfig def t5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tensor: @@ -37,6 +37,7 @@ def t5_postprocess_text(outputs: BaseEncoderOutput, _text_inputs) -> torch.Tenso class WanT2V480PConfig(PipelineConfig): """Base configuration for Wan T2V 1.3B pipeline architecture.""" + task_type: ModelTaskType = ModelTaskType.T2V # WanConfig-specific parameters with defaults # DiT dit_config: DiTConfig = field(default_factory=WanVideoConfig) @@ -84,7 +85,7 @@ class WanI2V480PConfig(WanT2V480PConfig): """Base configuration for Wan I2V 14B 480P pipeline architecture.""" # WanConfig-specific parameters with defaults - i2v_task: bool = True + task_type: ModelTaskType = ModelTaskType.I2V # Precision for each component image_encoder_config: EncoderConfig = field(default_factory=CLIPVisionConfig) image_encoder_precision: str = "fp32" @@ -129,7 +130,7 @@ class FastWan2_1_T2V_480P_Config(WanT2V480PConfig): @dataclass class Wan2_2_TI2V_5B_Config(WanT2V480PConfig): flow_shift: float | None = 5.0 - ti2v_task: bool = True + task_type: ModelTaskType = ModelTaskType.TI2V expand_timesteps: bool = True # ti2v, 5B vae_stride = (4, 16, 16) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index c59335465..962a880ad 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -259,7 +259,7 @@ class DiffGenerator: # TODO: simplify data_type = ( DataType.IMAGE - if self.server_args.pipeline_config.is_image_gen + if self.server_args.pipeline_config.task_type.is_image_task() or pretrained_sampling_params.num_frames == 1 else DataType.VIDEO ) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index 835f860f4..62e5c842d 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -50,60 +50,59 @@ def prepare_sampling_params( f"num_frames={sampling_params.num_frames}" ) - temporal_scale_factor = ( - pipeline_config.vae_config.arch_config.temporal_compression_ratio - ) - - # settle num_frames - if server_args.pipeline_config.is_image_gen: + if pipeline_config.task_type.is_image_task(): + # settle num_frames logger.debug(f"Setting num_frames to 1 because this is a image-gen model") sampling_params.num_frames = 1 - - num_frames = sampling_params.num_frames - num_gpus = server_args.num_gpus - use_temporal_scaling_frames = pipeline_config.vae_config.use_temporal_scaling_frames - - # Adjust number of frames based on number of GPUs - if use_temporal_scaling_frames: - orig_latent_num_frames = (num_frames - 1) // temporal_scale_factor + 1 - else: # stepvideo only - orig_latent_num_frames = sampling_params.num_frames // 17 * 3 - - if orig_latent_num_frames % server_args.num_gpus != 0: - # Adjust latent frames to be divisible by number of GPUs - if sampling_params.num_frames_round_down: - # Ensure we have at least 1 batch per GPU - new_latent_num_frames = ( - max(1, (orig_latent_num_frames // num_gpus)) * num_gpus - ) - else: - new_latent_num_frames = ( - math.ceil(orig_latent_num_frames / num_gpus) * num_gpus - ) + sampling_params.data_type = DataType.IMAGE + else: + # Adjust number of frames based on number of GPUs for video task + use_temporal_scaling_frames = ( + pipeline_config.vae_config.use_temporal_scaling_frames + ) + num_frames = sampling_params.num_frames + num_gpus = server_args.num_gpus + temporal_scale_factor = ( + pipeline_config.vae_config.arch_config.temporal_compression_ratio + ) if use_temporal_scaling_frames: - # Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor - new_num_frames = (new_latent_num_frames - 1) * temporal_scale_factor + 1 + orig_latent_num_frames = (num_frames - 1) // temporal_scale_factor + 1 else: # stepvideo only - # Find the least common multiple of 3 and num_gpus - divisor = math.lcm(3, num_gpus) - # Round up to the nearest multiple of this LCM - new_latent_num_frames = ( - (new_latent_num_frames + divisor - 1) // divisor - ) * divisor - # Convert back to actual frames using the StepVideo formula - new_num_frames = new_latent_num_frames // 3 * 17 + orig_latent_num_frames = sampling_params.num_frames // 17 * 3 - logger.info( - "Adjusting number of frames from %s to %s based on number of GPUs (%s)", - sampling_params.num_frames, - new_num_frames, - server_args.num_gpus, - ) - sampling_params.num_frames = new_num_frames + if orig_latent_num_frames % server_args.num_gpus != 0: + # Adjust latent frames to be divisible by number of GPUs + if sampling_params.num_frames_round_down: + # Ensure we have at least 1 batch per GPU + new_latent_num_frames = ( + max(1, (orig_latent_num_frames // num_gpus)) * num_gpus + ) + else: + new_latent_num_frames = ( + math.ceil(orig_latent_num_frames / num_gpus) * num_gpus + ) - if pipeline_config.is_image_gen: - sampling_params.data_type = DataType.IMAGE + if use_temporal_scaling_frames: + # Convert back to number of frames, ensuring num_frames-1 is a multiple of temporal_scale_factor + new_num_frames = (new_latent_num_frames - 1) * temporal_scale_factor + 1 + else: # stepvideo only + # Find the least common multiple of 3 and num_gpus + divisor = math.lcm(3, num_gpus) + # Round up to the nearest multiple of this LCM + new_latent_num_frames = ( + (new_latent_num_frames + divisor - 1) // divisor + ) * divisor + # Convert back to actual frames using the StepVideo formula + new_num_frames = new_latent_num_frames // 3 * 17 + + logger.info( + "Adjusting number of frames from %s to %s based on number of GPUs (%s)", + sampling_params.num_frames, + new_num_frames, + server_args.num_gpus, + ) + sampling_params.num_frames = new_num_frames sampling_params.set_output_file_ext() sampling_params.log(server_args=server_args) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py index 2e3922b11..a151fef86 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/denoising.py @@ -19,7 +19,7 @@ import torch.profiler from einops import rearrange from tqdm.auto import tqdm -from sglang.multimodal_gen.configs.pipelines.base import STA_Mode +from sglang.multimodal_gen.configs.pipelines.base import ModelTaskType, STA_Mode from sglang.multimodal_gen.runtime.distributed import ( cfg_model_parallel_all_reduce, get_local_torch_device, @@ -269,7 +269,10 @@ class DenoisingStage(PipelineStage): None, ) # FIXME: should probably move to latent preparation stage, to handle with offload - if server_args.pipeline_config.ti2v_task and batch.pil_image is not None: + if ( + server_args.pipeline_config.task_type == ModelTaskType.TI2V + and batch.pil_image is not None + ): # Wan2.2 TI2V directly replaces the first frame of the latent with # the image latent instead of appending along the channel dim assert batch.image_latent is None, "TI2V task should not have image latents" @@ -334,7 +337,7 @@ class DenoisingStage(PipelineStage): # Shard z and reserved_frames_mask for TI2V if SP is enabled if ( - server_args.pipeline_config.ti2v_task + server_args.pipeline_config.task_type == ModelTaskType.TI2V and batch.pil_image is not None and get_sp_world_size() > 1 ): @@ -681,7 +684,10 @@ class DenoisingStage(PipelineStage): ): bsz = batch.raw_latent_shape[0] # expand timestep - if server_args.pipeline_config.ti2v_task and batch.pil_image is not None: + if ( + server_args.pipeline_config.task_type == ModelTaskType.TI2V + and batch.pil_image is not None + ): # Explicitly cast t_device to the target float type at the beginning. # This ensures any precision-based rounding (e.g., float32(999.0) -> bfloat16(1000.0)) # is applied consistently *before* it's used by any rank. @@ -723,7 +729,10 @@ class DenoisingStage(PipelineStage): """ For Wan2.2 ti2v task, global first frame should be replaced with encoded image after each timestep """ - if server_args.pipeline_config.ti2v_task and batch.pil_image is not None: + if ( + server_args.pipeline_config.task_type == ModelTaskType.TI2V + and batch.pil_image is not None + ): # Apply TI2V mask blending with SP-aware z and reserved_frames_mask. # This ensures the first frame is always the condition image after each step. # This is only applied on rank 0, where z is not None. @@ -814,7 +823,8 @@ class DenoisingStage(PipelineStage): latent_model_input = latents.to(target_dtype) if batch.image_latent is not None: assert ( - not server_args.pipeline_config.ti2v_task + not server_args.pipeline_config.task_type + == ModelTaskType.TI2V ), "image latents should not be provided for TI2V task" latent_model_input = torch.cat( [latent_model_input, batch.image_latent], dim=1 diff --git a/python/sglang/multimodal_gen/runtime/pipelines/stages/input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines/stages/input_validation.py index 1fd09b3c5..49a618055 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/stages/input_validation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/stages/input_validation.py @@ -10,6 +10,7 @@ import torchvision.transforms.functional as TF from PIL import Image from sglang.multimodal_gen.configs.pipelines import WanI2V480PConfig +from sglang.multimodal_gen.configs.pipelines.base import ModelTaskType from sglang.multimodal_gen.configs.pipelines.qwen_image import ( QwenImageEditPipelineConfig, ) @@ -128,8 +129,8 @@ class InputValidationStage(PipelineStage): batch.width = width batch.height = height elif ( - server_args.pipeline_config.ti2v_task - or server_args.pipeline_config.ti2i_task + server_args.pipeline_config.task_type == ModelTaskType.TI2V + or server_args.pipeline_config.task_type == ModelTaskType.I2I ) and batch.pil_image is not None: # further processing for ti2v task img = batch.pil_image diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 02b1dea0b..1f4367411 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -800,7 +800,7 @@ class ServerArgs: def check_server_sp_args(self): - if self.pipeline_config.is_image_gen: + if self.pipeline_config.task_type.is_image_task(): if ( (self.sp_degree and self.sp_degree > 1) or (self.ulysses_degree and self.ulysses_degree > 1)