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)