diffusion: refactor task type of models (#13118)

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

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
 

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

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)

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
         )

python/sglang/multimodal_gen/runtime/entrypoints/utils.py

@@ -50,20 +50,22 @@ def prepare_sampling_params(
             f"num_frames={sampling_params.num_frames}"
         )
 
+    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
+        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
         )
 
-    # settle num_frames
-    if server_args.pipeline_config.is_image_gen:
-        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
@@ -102,9 +104,6 @@ def prepare_sampling_params(
             )
             sampling_params.num_frames = new_num_frames
 
-    if pipeline_config.is_image_gen:
-        sampling_params.data_type = DataType.IMAGE
-
     sampling_params.set_output_file_ext()
     sampling_params.log(server_args=server_args)
     return sampling_params

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

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
+            server_args.pipeline_config.task_type == ModelTaskType.TI2V
-            or server_args.pipeline_config.ti2i_task
+            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

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)