[diffusion] refactor: reduce redundancy and improve stage api (#19060)

docs/diffusion/support_new_models.md

@@ -45,7 +45,6 @@ You can build your custom `ComposedPipeline` by combining the following availabl
 | `TextEncodingStage`              | Encodes text prompts into embeddings using one or more text encoders.                                   |
 | `ImageEncodingStage`             | Encodes input images into embeddings, often used in image-to-image tasks.                               |
 | `ImageVAEEncodingStage`          | Specifically encodes an input image into the latent space using a Variational Autoencoder (VAE).        |
-| `ConditioningStage`              | Prepares the conditioning tensors (e.g., from text or image embeddings) for the denoising loop.         |
 | `TimestepPreparationStage`       | Prepares the scheduler's timesteps for the diffusion process.                                           |
 | `LatentPreparationStage`         | Creates the initial noisy latent tensor that will be denoised.                                          |
 | `DenoisingStage`                 | Executes the main denoising loop, iteratively applying the model (e.g., UNet) to refine the latents.    |
@@ -88,15 +87,13 @@ To illustrate the process, let's look at how `Qwen-Image-Edit` is implemented. T
         _required_config_modules = ["processor", "scheduler", "text_encoder", "tokenizer", "transformer", "vae"]
 
         def create_pipeline_stages(self, server_args: ServerArgs):
-            """Set up pipeline stages sequentially."""
+            self.add_stage(InputValidationStage())
-            self.add_stage(stage_name="input_validation_stage", stage=InputValidationStage())
+            self.add_stage(ImageEncodingStage(...))
-            self.add_stage(stage_name="prompt_encoding_stage_primary", stage=ImageEncodingStage(...))
+            self.add_stage(ImageVAEEncodingStage(...))
-            self.add_stage(stage_name="image_encoding_stage_primary", stage=ImageVAEEncodingStage(...))
+            self.add_stage(TimestepPreparationStage(...))
-            self.add_stage(stage_name="timestep_preparation_stage", stage=TimestepPreparationStage(...))
+            self.add_stage(LatentPreparationStage(...))
-            self.add_stage(stage_name="latent_preparation_stage", stage=LatentPreparationStage(...))
+            self.add_stage(DenoisingStage(...))
-            self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
+            self.add_stage(DecodingStage(...))
-            self.add_stage(stage_name="denoising_stage", stage=DenoisingStage(...))
-            self.add_stage(stage_name="decoding_stage", stage=DecodingStage(...))
     ```
     The pipeline is constructed by adding stages in order. `Qwen-Image-Edit` uses `ImageEncodingStage` (for prompt and image processing) and `ImageVAEEncodingStage` (for latent extraction) before standard denoising and decoding.
 

python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py

@@ -666,24 +666,19 @@ class ComfyUIFluxPipeline(LoRAPipeline, ComposedPipelineBase):
             "ComfyUIFluxPipeline.create_pipeline_stages() called - creating latent_preparation_stage and denoising_stage"
         )
 
-        # Add ComfyUILatentPreparationStage to handle latents properly for SP
+        self.add_stages(
-        # This stage includes device mismatch fix for ComfyUI pipelines in multi-GPU scenarios
+            [
-        self.add_stage(
+                ComfyUILatentPreparationStage(
-            stage_name="latent_preparation_stage",
-            stage=ComfyUILatentPreparationStage(
                     scheduler=self.get_module("scheduler"),
                     transformer=self.get_module("transformer"),
                 ),
+                DenoisingStage(
+                    transformer=self.get_module("transformer"),
+                    scheduler=self.get_module("scheduler"),
+                ),
+            ]
         )
 
-        # Add DenoisingStage for the actual denoising process
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-            ),
-        )
         logger.info(
             f"ComfyUIFluxPipeline stages created: {list(self._stage_name_mapping.keys())}"
         )

python/sglang/multimodal_gen/runtime/pipelines/comfyui_qwen_image_pipeline.py

@@ -293,24 +293,19 @@ class ComfyUIQwenImagePipelineBase(LoRAPipeline, ComposedPipelineBase):
             f"{self.__class__.__name__}.create_pipeline_stages() called - creating latent_preparation_stage and denoising_stage"
         )
 
-        # Add ComfyUILatentPreparationStage to handle latents properly for SP
+        self.add_stages(
-        # This stage includes device mismatch fix for ComfyUI pipelines in multi-GPU scenarios
+            [
-        self.add_stage(
+                ComfyUILatentPreparationStage(
-            stage_name="latent_preparation_stage",
-            stage=ComfyUILatentPreparationStage(
                     scheduler=self.get_module("scheduler"),
                     transformer=self.get_module("transformer"),
                 ),
+                DenoisingStage(
+                    transformer=self.get_module("transformer"),
+                    scheduler=self.get_module("scheduler"),
+                ),
+            ]
         )
 
-        # Add DenoisingStage for the actual denoising process
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-            ),
-        )
         logger.info(
             f"{self.__class__.__name__} stages created: {list(self._stage_name_mapping.keys())}"
         )

python/sglang/multimodal_gen/runtime/pipelines/comfyui_zimage_pipeline.py

@@ -381,24 +381,19 @@ class ComfyUIZImagePipeline(LoRAPipeline, ComposedPipelineBase):
             "ComfyUIZImagePipeline.create_pipeline_stages() called - creating latent_preparation_stage and denoising_stage"
         )
 
-        # Add ComfyUILatentPreparationStage to handle latents properly for SP
+        self.add_stages(
-        # This stage includes device mismatch fix for ComfyUI pipelines in multi-GPU scenarios
+            [
-        self.add_stage(
+                ComfyUILatentPreparationStage(
-            stage_name="latent_preparation_stage",
-            stage=ComfyUILatentPreparationStage(
                     scheduler=self.get_module("scheduler"),
                     transformer=self.get_module("transformer"),
                 ),
+                DenoisingStage(
+                    transformer=self.get_module("transformer"),
+                    scheduler=self.get_module("scheduler"),
+                ),
+            ]
         )
 
-        # Add DenoisingStage for the actual denoising process
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-            ),
-        )
         logger.info(
             f"ComfyUIZImagePipeline stages created: {list(self._stage_name_mapping.keys())}"
         )

python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py

@@ -604,8 +604,7 @@ class DiffusersPipeline(ComposedPipelineBase):
     def create_pipeline_stages(self, server_args: ServerArgs):
         """Create the execution stage wrapping the diffusers pipeline."""
         self.add_stage(
-            stage_name="diffusers_execution",
+            DiffusersExecutionStage(self.diffusers_pipe), "diffusers_execution"
-            stage=DiffusersExecutionStage(self.diffusers_pipe),
         )
 
     def initialize_pipeline(self, server_args: ServerArgs):
@@ -620,11 +619,18 @@ class DiffusersPipeline(ComposedPipelineBase):
         self.initialize_pipeline(self.server_args)
         self.create_pipeline_stages(self.server_args)
 
-    def add_stage(self, stage_name: str, stage: PipelineStage):
+    def add_stage(
+        self, stage: PipelineStage, stage_name: str | None = None
+    ) -> "DiffusersPipeline":
         """Add a stage to the pipeline."""
+        if stage_name is None:
+            stage_name = self._infer_stage_name(stage)
+        if stage_name in self._stage_name_mapping:
+            raise ValueError(f"Duplicate stage name detected: {stage_name}")
+
         self._stages.append(stage)
         self._stage_name_mapping[stage_name] = stage
-        setattr(self, stage_name, stage)
+        return self
 
     @property
     def stages(self) -> list[PipelineStage]:

python/sglang/multimodal_gen/runtime/pipelines/flux.py

@@ -8,13 +8,8 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import
 )
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
 from sglang.multimodal_gen.runtime.pipelines_core.stages import (
-    ConditioningStage,
-    DecodingStage,
-    DenoisingStage,
     InputValidationStage,
-    LatentPreparationStage,
     TextEncodingStage,
-    TimestepPreparationStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -70,15 +65,10 @@ class FluxPipeline(LoRAPipeline, ComposedPipelineBase):
     ]
 
     def create_pipeline_stages(self, server_args: ServerArgs):
-        """Set up pipeline stages with proper dependency injection."""
+        self.add_stage(InputValidationStage())
 
         self.add_stage(
-            stage_name="input_validation_stage", stage=InputValidationStage()
+            TextEncodingStage(
-        )
-
-        self.add_stage(
-            stage_name="prompt_encoding_stage_primary",
-            stage=TextEncodingStage(
                 text_encoders=[
                     self.get_module("text_encoder"),
                     self.get_module("text_encoder_2"),
@@ -88,37 +78,13 @@ class FluxPipeline(LoRAPipeline, ComposedPipelineBase):
                     self.get_module("tokenizer_2"),
                 ],
             ),
+            "prompt_encoding_stage_primary",
         )
 
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
+        self.add_standard_timestep_preparation_stage(prepare_extra_kwargs=[prepare_mu])
-
+        self.add_standard_latent_preparation_stage()
-        self.add_stage(
+        self.add_standard_denoising_stage()
-            stage_name="timestep_preparation_stage",
+        self.add_standard_decoding_stage()
-            stage=TimestepPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                prepare_extra_set_timesteps_kwargs=[prepare_mu],
-            ),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
-        )
 
 
 EntryClass = FluxPipeline

python/sglang/multimodal_gen/runtime/pipelines/flux_2.py

@@ -7,16 +7,6 @@ from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline, Req
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
-from sglang.multimodal_gen.runtime.pipelines_core.stages import (
-    ConditioningStage,
-    DecodingStage,
-    DenoisingStage,
-    ImageVAEEncodingStage,
-    InputValidationStage,
-    LatentPreparationStage,
-    TextEncodingStage,
-    TimestepPreparationStage,
-)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
@@ -55,69 +45,17 @@ class Flux2Pipeline(LoRAPipeline, ComposedPipelineBase):
     ]
 
     def create_pipeline_stages(self, server_args: ServerArgs):
-        """Set up pipeline stages with proper dependency injection."""
+        vae_image_processor = VaeImageProcessor(
-
-        self.add_stage(
-            stage_name="input_validation_stage",
-            stage=InputValidationStage(
-                vae_image_processor=VaeImageProcessor(
             vae_scale_factor=server_args.pipeline_config.vae_config.arch_config.vae_scale_factor
             * 2
-                ),
-            ),
         )
 
-        self.add_stage(
+        self.add_standard_ti2i_stages(
-            stage_name="prompt_encoding_stage_primary",
+            include_input_validation=True,
-            stage=TextEncodingStage(
+            vae_image_processor=vae_image_processor,
-                text_encoders=[
+            prompt_encoding="text",
-                    self.get_module("text_encoder"),
+            image_vae_stage_kwargs={"vae_image_processor": vae_image_processor},
-                ],
+            prepare_extra_timestep_kwargs=[compute_empirical_mu],
-                tokenizers=[
-                    self.get_module("tokenizer"),
-                ],
-            ),
-        )
-
-        self.add_stage(
-            stage_name="image_encoding_stage_primary",
-            stage=ImageVAEEncodingStage(
-                vae_image_processor=VaeImageProcessor(
-                    vae_scale_factor=server_args.pipeline_config.vae_config.arch_config.vae_scale_factor
-                    * 2
-                ),
-                vae=self.get_module("vae"),
-            ),
-        )
-
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="timestep_preparation_stage",
-            stage=TimestepPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                prepare_extra_set_timesteps_kwargs=[compute_empirical_mu],
-            ),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
         )
 
 

python/sglang/multimodal_gen/runtime/pipelines/glm_image.py

@@ -2,10 +2,7 @@ from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
-from sglang.multimodal_gen.runtime.pipelines_core.stages import (
+from sglang.multimodal_gen.runtime.pipelines_core.stages import DenoisingStage
-    DecodingStage,
-    DenoisingStage,
-)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.glm_image import (
     GlmImageBeforeDenoisingStage,
 )
@@ -30,8 +27,7 @@ class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase):
 
     def create_pipeline_stages(self, server_args: ServerArgs):
         self.add_stage(
-            stage_name="GlmImageBeforeDenoisingStage",
+            GlmImageBeforeDenoisingStage(
-            stage=GlmImageBeforeDenoisingStage(
                 vae=self.get_module("vae"),
                 text_encoder=self.get_module("text_encoder"),
                 tokenizer=self.get_module("tokenizer"),
@@ -40,19 +36,17 @@ class GlmImagePipeline(LoRAPipeline, ComposedPipelineBase):
                 scheduler=self.get_module("scheduler"),
                 vision_language_encoder=self.get_module("vision_language_encoder"),
             ),
+            "glm_image_before_denoising_stage",
         )
 
         self.add_stage(
-            stage_name="denoising_stage",
+            DenoisingStage(
-            stage=DenoisingStage(
                 transformer=self.get_module("transformer"),
                 scheduler=self.get_module("scheduler"),
             ),
         )
 
-        self.add_stage(
+        self.add_standard_decoding_stage()
-            stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
-        )
 
 
 EntryClass = [GlmImagePipeline]

python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py

@@ -12,13 +12,8 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import
     ComposedPipelineBase,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages import (
-    ConditioningStage,
-    DecodingStage,
-    DenoisingStage,
     InputValidationStage,
-    LatentPreparationStage,
     TextEncodingStage,
-    TimestepPreparationStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -43,15 +38,9 @@ class HunyuanVideoPipeline(ComposedPipelineBase):
     ]
 
     def create_pipeline_stages(self, server_args: ServerArgs):
-        """Set up pipeline stages with proper dependency injection."""
+        self.add_stage(InputValidationStage())
-
         self.add_stage(
-            stage_name="input_validation_stage", stage=InputValidationStage()
+            TextEncodingStage(
-        )
-
-        self.add_stage(
-            stage_name="prompt_encoding_stage_primary",
-            stage=TextEncodingStage(
                 text_encoders=[
                     self.get_module("text_encoder"),
                     self.get_module("text_encoder_2"),
@@ -61,34 +50,12 @@ class HunyuanVideoPipeline(ComposedPipelineBase):
                     self.get_module("tokenizer_2"),
                 ],
             ),
+            "prompt_encoding_stage_primary",
         )
-
+        self.add_standard_timestep_preparation_stage()
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
+        self.add_standard_latent_preparation_stage()
-
+        self.add_standard_denoising_stage()
-        self.add_stage(
+        self.add_standard_decoding_stage()
-            stage_name="timestep_preparation_stage",
-            stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
-        )
 
 
 EntryClass = HunyuanVideoPipeline

python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py

@@ -13,7 +13,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages import (
     LTX2AVLatentPreparationStage,
     LTX2TextConnectorStage,
     TextEncodingStage,
-    TimestepPreparationStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -118,73 +117,39 @@ class LTX2Pipeline(ComposedPipelineBase):
     ]
 
     def create_pipeline_stages(self, server_args: ServerArgs):
-        """Set up pipeline stages with proper dependency injection."""
+        self.add_stages(
-
+            [
-        # 1. Input Validation
+                InputValidationStage(),
-        self.add_stage(
+                TextEncodingStage(
-            stage_name="input_validation_stage", stage=InputValidationStage()
+                    text_encoders=[self.get_module("text_encoder")],
-        )
+                    tokenizers=[self.get_module("tokenizer")],
-
-        # 2. Text Encoding
-        self.add_stage(
-            stage_name="text_encoding_stage",
-            stage=TextEncodingStage(
-                # LTX-2 needs two contexts (video/audio). We reuse the same
-                # underlying Gemma encoder/tokenizer twice.
-                text_encoders=[
-                    self.get_module("text_encoder"),
-                ],
-                tokenizers=[
-                    self.get_module("tokenizer"),
-                ],
                 ),
+                LTX2TextConnectorStage(connectors=self.get_module("connectors")),
+            ]
         )
 
-        # 3. connector stage
+        self.add_standard_timestep_preparation_stage(prepare_extra_kwargs=[prepare_mu])
-        self.add_stage(
-            stage_name="text_connector_stage",
-            stage=LTX2TextConnectorStage(connectors=self.get_module("connectors")),
-        )
 
-        # 4. Timestep Preparation
+        self.add_stages(
-        self.add_stage(
+            [
-            stage_name="timestep_preparation_stage",
+                LTX2AVLatentPreparationStage(
-            stage=TimestepPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                prepare_extra_set_timesteps_kwargs=[prepare_mu],
-            ),
-        )
-
-        # 4. Latent Preparation
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LTX2AVLatentPreparationStage(
                     scheduler=self.get_module("scheduler"),
                     transformer=self.get_module("transformer"),
                     audio_vae=self.get_module("audio_vae"),
                 ),
-        )
+                LTX2AVDenoisingStage(
-
-        # 5. Denoising
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=LTX2AVDenoisingStage(
                     transformer=self.get_module("transformer"),
                     scheduler=self.get_module("scheduler"),
                     vae=self.get_module("vae"),
                     audio_vae=self.get_module("audio_vae"),
                 ),
-        )
+                LTX2AVDecodingStage(
-
-        # 6. Decoding
-        self.add_stage(
-            stage_name="decoding_stage",
-            stage=LTX2AVDecodingStage(
                     vae=self.get_module("vae"),
                     audio_vae=self.get_module("audio_vae"),
                     vocoder=self.get_module("vocoder"),
                     pipeline=self,
                 ),
+            ]
         )
 
 

python/sglang/multimodal_gen/runtime/pipelines/mova_pipeline.py

@@ -11,10 +11,8 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import
     ComposedPipelineBase,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages import (
-    ConditioningStage,
     ImageVAEEncodingStage,
     InputValidationStage,
-    TextEncodingStage,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.mova import (
     MOVADecodingStage,
@@ -62,53 +60,41 @@ class MOVAPipeline(ComposedPipelineBase):
             )
 
     def create_pipeline_stages(self, server_args: ServerArgs) -> None:
-        self.add_stage(
+        self.add_stage(InputValidationStage())
-            stage_name="input_validation_stage", stage=InputValidationStage()
+        self.add_standard_text_encoding_stage()
-        )
-        self.add_stage(
-            stage_name="prompt_encoding_stage",
-            stage=TextEncodingStage(
-                text_encoders=[self.get_module("text_encoder")],
-                tokenizers=[self.get_module("tokenizer")],
-            ),
-        )
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
         if getattr(self.get_module("video_dit"), "require_vae_embedding", True):
+            self.add_stage(ImageVAEEncodingStage(vae=self.get_module("video_vae")))
         self.add_stage(
-                stage_name="image_latent_preparation_stage",
+            MOVALatentPreparationStage(
-                stage=ImageVAEEncodingStage(vae=self.get_module("video_vae")),
-            )
-        self.add_stage(
-            stage_name="mova_latent_preparation_stage",
-            stage=MOVALatentPreparationStage(
                 audio_vae=self.get_module("audio_vae"),
                 require_vae_embedding=getattr(
                     self.get_module("video_dit"), "require_vae_embedding", True
                 ),
             ),
+            "mova_latent_preparation_stage",
         )
         self.add_stage(
-            stage_name="mova_timestep_preparation_stage",
+            MOVATimestepPreparationStage(
-            stage=MOVATimestepPreparationStage(
                 scheduler=self.get_module("scheduler"),
             ),
+            "mova_timestep_preparation_stage",
         )
         self.add_stage(
-            stage_name="mova_denoising_stage",
+            MOVADenoisingStage(
-            stage=MOVADenoisingStage(
                 video_dit=self.get_module("video_dit"),
                 video_dit_2=self.get_module("video_dit_2"),
                 audio_dit=self.get_module("audio_dit"),
                 dual_tower_bridge=self.get_module("dual_tower_bridge"),
                 scheduler=self.get_module("scheduler"),
             ),
+            "mova_denoising_stage",
         )
         self.add_stage(
-            stage_name="mova_decoding_stage",
+            MOVADecodingStage(
-            stage=MOVADecodingStage(
                 video_vae=self.get_module("video_vae"),
                 audio_vae=self.get_module("audio_vae"),
             ),
+            "mova_decoding_stage",
         )
 
 

python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py

@@ -8,19 +8,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import
     ComposedPipelineBase,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
-from sglang.multimodal_gen.runtime.pipelines_core.stages import (
-    DecodingStage,
-    DenoisingStage,
-    ImageEncodingStage,
-    ImageVAEEncodingStage,
-    InputValidationStage,
-    LatentPreparationStage,
-    TextEncodingStage,
-    TimestepPreparationStage,
-)
-from sglang.multimodal_gen.runtime.pipelines_core.stages.conditioning import (
-    ConditioningStage,
-)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.qwen_image_layered import (
     QwenImageLayeredBeforeDenoisingStage,
 )
@@ -75,53 +62,7 @@ class QwenImagePipeline(LoRAPipeline, ComposedPipelineBase):
     ]
 
     def create_pipeline_stages(self, server_args: ServerArgs):
-        """Set up pipeline stages with proper dependency injection."""
+        self.add_standard_t2i_stages(prepare_extra_timestep_kwargs=[prepare_mu])
-
-        self.add_stage(
-            stage_name="input_validation_stage", stage=InputValidationStage()
-        )
-
-        self.add_stage(
-            stage_name="prompt_encoding_stage_primary",
-            stage=TextEncodingStage(
-                text_encoders=[
-                    self.get_module("text_encoder"),
-                ],
-                tokenizers=[
-                    self.get_module("tokenizer"),
-                ],
-            ),
-        )
-
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
-
-        self.add_stage(
-            stage_name="timestep_preparation_stage",
-            stage=TimestepPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                prepare_extra_set_timesteps_kwargs=[prepare_mu],
-            ),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
-        )
 
 
 class QwenImageEditPipeline(LoRAPipeline, ComposedPipelineBase):
@@ -137,61 +78,17 @@ class QwenImageEditPipeline(LoRAPipeline, ComposedPipelineBase):
     ]
 
     def create_pipeline_stages(self, server_args: ServerArgs):
-        """Set up pipeline stages with proper dependency injection."""
+        vae_image_processor = VaeImageProcessor(
-
-        self.add_stage(
-            stage_name="input_validation_stage",
-            stage=InputValidationStage(
-                vae_image_processor=VaeImageProcessor(
             vae_scale_factor=server_args.pipeline_config.vae_config.arch_config.vae_scale_factor
             * 2
         )
-            ),
-        )
 
-        self.add_stage(
+        self.add_standard_ti2i_stages(
-            stage_name="prompt_encoding_stage_primary",
+            vae_image_processor=vae_image_processor,
-            stage=ImageEncodingStage(
+            prompt_encoding="image_encoding",
-                image_processor=self.get_module("processor"),
+            image_processor_key="processor",
-                text_encoder=self.get_module("text_encoder"),
+            prompt_text_encoder_key="text_encoder",
-            ),
+            prepare_extra_timestep_kwargs=[prepare_mu],
-        )
-
-        self.add_stage(
-            stage_name="image_encoding_stage_primary",
-            stage=ImageVAEEncodingStage(
-                vae=self.get_module("vae"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="timestep_preparation_stage",
-            stage=TimestepPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                prepare_extra_set_timesteps_kwargs=[prepare_mu],
-            ),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer"),
-            ),
-        )
-
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
         )
 
 
@@ -217,39 +114,22 @@ class QwenImageLayeredPipeline(QwenImageEditPipeline):
     ]
 
     def create_pipeline_stages(self, server_args: ServerArgs):
-        """Set up pipeline stages with proper dependency injection."""
-
         self.add_stage(
-            stage_name="QwenImageLayeredBeforeDenoisingStage",
+            QwenImageLayeredBeforeDenoisingStage(
-            stage=QwenImageLayeredBeforeDenoisingStage(
                 vae=self.get_module("vae"),
                 tokenizer=self.get_module("tokenizer"),
                 processor=self.get_module("processor"),
                 transformer=self.get_module("transformer"),
                 scheduler=self.get_module("scheduler"),
                 model_path=self.model_path,
-            ),
+            )
         )
 
-        self.add_stage(
+        self.add_standard_timestep_preparation_stage(
-            stage_name="timestep_preparation_stage",
+            prepare_extra_kwargs=[prepare_mu_layered]
-            stage=TimestepPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                prepare_extra_set_timesteps_kwargs=[prepare_mu_layered],
-            ),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
         )
+        self.add_standard_denoising_stage()
+        self.add_standard_decoding_stage()
 
 
 EntryClass = [

python/sglang/multimodal_gen/runtime/pipelines/wan_causal_dmd_pipeline.py

@@ -14,12 +14,8 @@ from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipel
 
 # isort: off
 from sglang.multimodal_gen.runtime.pipelines_core.stages import (
-    ConditioningStage,
-    DecodingStage,
     CausalDMDDenoisingStage,
     InputValidationStage,
-    LatentPreparationStage,
-    TextEncodingStage,
 )
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
@@ -41,41 +37,18 @@ class WanCausalDMDPipeline(LoRAPipeline, ComposedPipelineBase):
     ]
 
     def create_pipeline_stages(self, server_args: ServerArgs) -> None:
-        """Set up pipeline stages with proper dependency injection."""
+        self.add_stage(InputValidationStage())
+        self.add_standard_text_encoding_stage()
+        self.add_standard_latent_preparation_stage()
 
         self.add_stage(
-            stage_name="input_validation_stage", stage=InputValidationStage()
+            CausalDMDDenoisingStage(
-        )
-
-        self.add_stage(
-            stage_name="prompt_encoding_stage",
-            stage=TextEncodingStage(
-                text_encoders=[self.get_module("text_encoder")],
-                tokenizers=[self.get_module("tokenizer")],
-            ),
-        )
-
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer", None),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=CausalDMDDenoisingStage(
                 transformer=self.get_module("transformer"),
                 scheduler=self.get_module("scheduler"),
             ),
         )
 
-        self.add_stage(
+        self.add_standard_decoding_stage()
-            stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
-        )
 
 
 EntryClass = WanCausalDMDPipeline

python/sglang/multimodal_gen/runtime/pipelines/wan_dmd_pipeline.py

@@ -20,13 +20,8 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
 # isort: off
 from sglang.multimodal_gen.runtime.pipelines_core.stages import (
-    ConditioningStage,
-    DecodingStage,
     DmdDenoisingStage,
     InputValidationStage,
-    LatentPreparationStage,
-    TextEncodingStage,
-    TimestepPreparationStage,
 )
 
 # isort: on
@@ -56,46 +51,27 @@ class WanDMDPipeline(LoRAPipeline, ComposedPipelineBase):
         )
 
     def create_pipeline_stages(self, server_args: ServerArgs) -> None:
-        """Set up pipeline stages with proper dependency injection."""
+        self.add_stages(
-
+            [
-        self.add_stage(
+                InputValidationStage(),
-            stage_name="input_validation_stage", stage=InputValidationStage()
+            ]
         )
 
-        self.add_stage(
+        self.add_standard_text_encoding_stage()
-            stage_name="prompt_encoding_stage",
-            stage=TextEncodingStage(
-                text_encoders=[self.get_module("text_encoder")],
-                tokenizers=[self.get_module("tokenizer")],
-            ),
-        )
 
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
+        self.add_standard_timestep_preparation_stage()
+        self.add_standard_latent_preparation_stage()
 
-        self.add_stage(
+        self.add_stages(
-            stage_name="timestep_preparation_stage",
+            [
-            stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
+                DmdDenoisingStage(
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer", None),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DmdDenoisingStage(
                     transformer=self.get_module("transformer"),
                     scheduler=self.get_module("scheduler"),
                 ),
+            ]
         )
 
-        self.add_stage(
+        self.add_standard_decoding_stage()
-            stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
-        )
 
 
 EntryClass = WanDMDPipeline

python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_dmd_pipeline.py

@@ -8,31 +8,17 @@ This module contains an implementation of the Wan video diffusion pipeline
 using the modular pipeline architecture.
 """
 
+from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import (
+    FlowMatchEulerDiscreteScheduler,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline
+from sglang.multimodal_gen.runtime.pipelines_core.stages import DmdDenoisingStage
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
-# isort: off
-from sglang.multimodal_gen.runtime.pipelines_core.stages import (
-    ImageEncodingStage,
-    ConditioningStage,
-    DecodingStage,
-    DmdDenoisingStage,
-    ImageVAEEncodingStage,
-    InputValidationStage,
-    LatentPreparationStage,
-    TextEncodingStage,
-    TimestepPreparationStage,
-)
-
-# isort: on
-from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_match_euler_discrete import (
-    FlowMatchEulerDiscreteScheduler,
-)
-
 logger = init_logger(__name__)
 
 
@@ -55,63 +41,14 @@ class WanImageToVideoDmdPipeline(LoRAPipeline, ComposedPipelineBase):
         )
 
     def create_pipeline_stages(self, server_args: ServerArgs):
-        """Set up pipeline stages with proper dependency injection."""
+        self.add_standard_ti2v_stages(
-
+            image_vae_encoding_position="after_latent",
-        self.add_stage(
+            denoising_stage_factory=lambda: DmdDenoisingStage(
-            stage_name="input_validation_stage", stage=InputValidationStage()
-        )
-
-        self.add_stage(
-            stage_name="prompt_encoding_stage",
-            stage=TextEncodingStage(
-                text_encoders=[self.get_module("text_encoder")],
-                tokenizers=[self.get_module("tokenizer")],
-            ),
-        )
-        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())
-
-        self.add_stage(
-            stage_name="timestep_preparation_stage",
-            stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="image_latent_preparation_stage",
-            stage=ImageVAEEncodingStage(vae=self.get_module("vae")),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DmdDenoisingStage(
                 transformer=self.get_module("transformer"),
                 scheduler=self.get_module("scheduler"),
                 transformer_2=self.get_module("transformer_2"),
             ),
         )
 
-        self.add_stage(
-            stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
-        )
-
 
 EntryClass = WanImageToVideoDmdPipeline

python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_pipeline.py

@@ -8,6 +8,9 @@ This module contains an implementation of the Wan video diffusion pipeline
 using the modular pipeline architecture.
 """
 
+from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_unipc_multistep import (
+    FlowUniPCMultistepScheduler,
+)
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
@@ -15,24 +18,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipel
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
-# isort: off
-from sglang.multimodal_gen.runtime.pipelines_core.stages import (
-    ImageEncodingStage,
-    ConditioningStage,
-    DecodingStage,
-    DenoisingStage,
-    ImageVAEEncodingStage,
-    InputValidationStage,
-    LatentPreparationStage,
-    TextEncodingStage,
-    TimestepPreparationStage,
-)
-
-# isort: on
-from sglang.multimodal_gen.runtime.models.schedulers.scheduling_flow_unipc_multistep import (
-    FlowUniPCMultistepScheduler,
-)
-
 logger = init_logger(__name__)
 
 
@@ -55,64 +40,7 @@ class WanImageToVideoPipeline(LoRAPipeline, ComposedPipelineBase):
         )
 
     def create_pipeline_stages(self, server_args: ServerArgs):
-        """Set up pipeline stages with proper dependency injection."""
+        self.add_standard_ti2v_stages()
-
-        self.add_stage(
-            stage_name="input_validation_stage", stage=InputValidationStage()
-        )
-
-        self.add_stage(
-            stage_name="prompt_encoding_stage",
-            stage=TextEncodingStage(
-                text_encoders=[self.get_module("text_encoder")],
-                tokenizers=[self.get_module("tokenizer")],
-            ),
-        )
-
-        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())
-
-        self.add_stage(
-            stage_name="timestep_preparation_stage",
-            stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="image_latent_preparation_stage",
-            stage=ImageVAEEncodingStage(vae=self.get_module("vae")),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                transformer_2=self.get_module("transformer_2"),
-                scheduler=self.get_module("scheduler"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
-        )
 
 
 EntryClass = WanImageToVideoPipeline

python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py

@@ -15,15 +15,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import
     ComposedPipelineBase,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline
-from sglang.multimodal_gen.runtime.pipelines_core.stages import (
-    ConditioningStage,
-    DecodingStage,
-    DenoisingStage,
-    InputValidationStage,
-    LatentPreparationStage,
-    TextEncodingStage,
-    TimestepPreparationStage,
-)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
@@ -52,50 +43,7 @@ class WanPipeline(LoRAPipeline, ComposedPipelineBase):
         )
 
     def create_pipeline_stages(self, server_args: ServerArgs) -> None:
-        """Set up pipeline stages with proper dependency injection."""
+        self.add_standard_t2i_stages()
-
-        self.add_stage(
-            stage_name="input_validation_stage", stage=InputValidationStage()
-        )
-
-        self.add_stage(
-            stage_name="prompt_encoding_stage",
-            stage=TextEncodingStage(
-                text_encoders=[self.get_module("text_encoder")],
-                tokenizers=[self.get_module("tokenizer")],
-            ),
-        )
-
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
-
-        self.add_stage(
-            stage_name="timestep_preparation_stage",
-            stage=TimestepPreparationStage(scheduler=self.get_module("scheduler")),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer", None),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                transformer_2=self.get_module("transformer_2", None),
-                scheduler=self.get_module("scheduler"),
-                vae=self.get_module("vae"),
-                pipeline=self,
-            ),
-        )
-
-        self.add_stage(
-            stage_name="decoding_stage",
-            stage=DecodingStage(vae=self.get_module("vae"), pipeline=self),
-        )
 
 
 EntryClass = WanPipeline

python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py

@@ -6,15 +6,6 @@ from sglang.multimodal_gen.runtime.pipelines_core import LoRAPipeline, Req
 from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import (
     ComposedPipelineBase,
 )
-from sglang.multimodal_gen.runtime.pipelines_core.stages import (
-    ConditioningStage,
-    DecodingStage,
-    DenoisingStage,
-    InputValidationStage,
-    LatentPreparationStage,
-    TextEncodingStage,
-    TimestepPreparationStage,
-)
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 
@@ -64,53 +55,7 @@ class ZImagePipeline(LoRAPipeline, ComposedPipelineBase):
     ]
 
     def create_pipeline_stages(self, server_args: ServerArgs):
-        """Set up pipeline stages with proper dependency injection."""
+        self.add_standard_t2i_stages(prepare_extra_timestep_kwargs=[prepare_mu])
-
-        self.add_stage(
-            stage_name="input_validation_stage", stage=InputValidationStage()
-        )
-
-        self.add_stage(
-            stage_name="prompt_encoding_stage_primary",
-            stage=TextEncodingStage(
-                text_encoders=[
-                    self.get_module("text_encoder"),
-                ],
-                tokenizers=[
-                    self.get_module("tokenizer"),
-                ],
-            ),
-        )
-
-        self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage())
-
-        self.add_stage(
-            stage_name="timestep_preparation_stage",
-            stage=TimestepPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                prepare_extra_set_timesteps_kwargs=[prepare_mu],
-            ),
-        )
-
-        self.add_stage(
-            stage_name="latent_preparation_stage",
-            stage=LatentPreparationStage(
-                scheduler=self.get_module("scheduler"),
-                transformer=self.get_module("transformer"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="denoising_stage",
-            stage=DenoisingStage(
-                transformer=self.get_module("transformer"),
-                scheduler=self.get_module("scheduler"),
-            ),
-        )
-
-        self.add_stage(
-            stage_name="decoding_stage", stage=DecodingStage(vae=self.get_module("vae"))
-        )
 
 
 EntryClass = ZImagePipeline

python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py

@@ -8,8 +8,9 @@ This module defines the base class for pipelines that are composed of multiple s
 """
 
 import os
+import re
 from abc import ABC, abstractmethod
-from typing import Any, cast
+from typing import Any, Callable, Literal, cast
 
 import torch
 from tqdm import tqdm
@@ -21,7 +22,17 @@ from sglang.multimodal_gen.runtime.pipelines_core.executors.pipeline_executor im
     PipelineExecutor,
 )
 from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req
-from sglang.multimodal_gen.runtime.pipelines_core.stages import PipelineStage
+from sglang.multimodal_gen.runtime.pipelines_core.stages import (
+    DecodingStage,
+    DenoisingStage,
+    ImageEncodingStage,
+    ImageVAEEncodingStage,
+    InputValidationStage,
+    LatentPreparationStage,
+    PipelineStage,
+    TextEncodingStage,
+    TimestepPreparationStage,
+)
 from sglang.multimodal_gen.runtime.platforms import current_platform
 from sglang.multimodal_gen.runtime.server_args import ServerArgs
 from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import (
@@ -109,9 +120,7 @@ class ComposedPipelineBase(ABC):
         self.create_pipeline_stages(self.server_args)
 
     def get_module(self, module_name: str, default_value: Any = None) -> Any:
-        if module_name not in self.modules:
+        return self.modules.get(module_name, default_value)
-            return default_value
-        return self.modules[module_name]
 
     def add_module(self, module_name: str, module: Any):
         self.modules[module_name] = module
@@ -320,11 +329,265 @@ class ComposedPipelineBase(ABC):
 
         return loaded_components
 
-    def add_stage(self, stage_name: str, stage: PipelineStage):
+    @staticmethod
+    def _infer_stage_name(stage: PipelineStage) -> str:
+        class_name = stage.__class__.__name__
+        # snake_case
+        name = re.sub(r"(?<!^)(?=[A-Z])", "_", class_name).lower()
+        if not name.endswith("_stage"):
+            name += "_stage"
+        return name
+
+    def add_stage(
+        self, stage: PipelineStage, stage_name: str | None = None
+    ) -> "ComposedPipelineBase":
+
         assert self.modules is not None, "No modules are registered"
+
+        if stage_name is None:
+            stage_name = self._infer_stage_name(stage)
+        if stage_name in self._stage_name_mapping:
+            raise ValueError(f"Duplicate stage name detected: {stage_name}")
+
         self._stages.append(stage)
         self._stage_name_mapping[stage_name] = stage
-        setattr(self, stage_name, stage)
+        return self
+
+    def add_stages(
+        self, stages: list[PipelineStage | tuple[PipelineStage, str]]
+    ) -> "ComposedPipelineBase":
+
+        for item in stages:
+            if isinstance(item, tuple):
+                stage, name = item
+                self.add_stage(stage, name)
+            else:
+                self.add_stage(item)
+        return self
+
+    def add_stage_if(
+        self,
+        condition: bool | Callable[[], bool],
+        stage: PipelineStage,
+    ) -> "ComposedPipelineBase":
+        should_add = condition() if callable(condition) else condition
+        if should_add:
+            self.add_stage(stage)
+        return self
+
+    def get_stage(self, stage_name: str) -> PipelineStage | None:
+        """Get a stage by name."""
+        return self._stage_name_mapping.get(stage_name)
+
+    def add_standard_text_encoding_stage(
+        self,
+        text_encoder_key: str = "text_encoder",
+        tokenizer_key: str = "tokenizer",
+    ) -> "ComposedPipelineBase":
+        return self.add_stage(
+            TextEncodingStage(
+                text_encoders=[self.get_module(text_encoder_key)],
+                tokenizers=[self.get_module(tokenizer_key)],
+            ),
+        )
+
+    def add_standard_timestep_preparation_stage(
+        self,
+        scheduler_key: str = "scheduler",
+        prepare_extra_kwargs: list[Callable] | None = [],
+    ) -> "ComposedPipelineBase":
+        return self.add_stage(
+            TimestepPreparationStage(
+                scheduler=self.get_module(scheduler_key),
+                prepare_extra_set_timesteps_kwargs=prepare_extra_kwargs,
+            ),
+        )
+
+    def add_standard_latent_preparation_stage(
+        self,
+        scheduler_key: str = "scheduler",
+        transformer_key: str = "transformer",
+    ) -> "ComposedPipelineBase":
+        return self.add_stage(
+            LatentPreparationStage(
+                scheduler=self.get_module(scheduler_key),
+                transformer=self.get_module(transformer_key),
+            ),
+        )
+
+    def add_standard_denoising_stage(
+        self,
+        transformer_key: str = "transformer",
+        transformer_2_key: str | None = "transformer_2",
+        scheduler_key: str = "scheduler",
+        vae_key: str | None = "vae",
+    ) -> "ComposedPipelineBase":
+
+        kwargs = {
+            "transformer": self.get_module(transformer_key),
+            "scheduler": self.get_module(scheduler_key),
+        }
+
+        if transformer_2_key:
+            transformer_2 = self.get_module(transformer_2_key, None)
+            if transformer_2 is not None:
+                kwargs["transformer_2"] = transformer_2
+
+        if vae_key:
+            vae = self.get_module(vae_key, None)
+            if vae is not None:
+                kwargs["vae"] = vae
+                kwargs["pipeline"] = self
+
+        return self.add_stage(DenoisingStage(**kwargs))
+
+    def add_standard_decoding_stage(
+        self,
+        vae_key: str = "vae",
+    ) -> "ComposedPipelineBase":
+
+        return self.add_stage(
+            DecodingStage(vae=self.get_module(vae_key), pipeline=self),
+        )
+
+    def add_standard_t2i_stages(
+        self,
+        include_input_validation: bool = True,
+        prepare_extra_timestep_kwargs: list[Callable] | None = [],
+    ) -> "ComposedPipelineBase":
+
+        if include_input_validation:
+            self.add_stage(InputValidationStage())
+
+        self.add_standard_text_encoding_stage()
+
+        self.add_standard_latent_preparation_stage()
+        self.add_standard_timestep_preparation_stage(
+            prepare_extra_kwargs=prepare_extra_timestep_kwargs
+        )
+        self.add_standard_denoising_stage()
+        self.add_standard_decoding_stage()
+
+        return self
+
+    def add_standard_ti2i_stages(
+        self,
+        *,
+        include_input_validation: bool = True,
+        vae_image_processor: Any | None = None,
+        prompt_encoding: Literal["text", "image_encoding"] = "text",
+        text_encoder_key: str = "text_encoder",
+        tokenizer_key: str = "tokenizer",
+        image_processor_key: str = "processor",
+        prompt_text_encoder_key: str = "text_encoder",
+        image_vae_key: str = "vae",
+        image_vae_stage_kwargs: dict[str, Any] | None = None,
+        prepare_extra_timestep_kwargs: list[Callable] | None = [],
+    ) -> "ComposedPipelineBase":
+        if include_input_validation:
+            self.add_stage(
+                InputValidationStage(vae_image_processor=vae_image_processor)
+            )
+
+        if prompt_encoding == "text":
+            self.add_standard_text_encoding_stage(
+                text_encoder_key=text_encoder_key,
+                tokenizer_key=tokenizer_key,
+            )
+        elif prompt_encoding == "image_encoding":
+            self.add_stage(
+                ImageEncodingStage(
+                    image_processor=self.get_module(image_processor_key),
+                    text_encoder=self.get_module(prompt_text_encoder_key),
+                ),
+            )
+        else:
+            raise ValueError(f"Unknown prompt_encoding: {prompt_encoding}")
+
+        self.add_stage(
+            ImageVAEEncodingStage(
+                vae=self.get_module(image_vae_key),
+                **(image_vae_stage_kwargs or {}),
+            ),
+        )
+
+        self.add_standard_latent_preparation_stage()
+
+        self.add_standard_timestep_preparation_stage(
+            prepare_extra_kwargs=prepare_extra_timestep_kwargs
+        )
+        self.add_standard_denoising_stage()
+        self.add_standard_decoding_stage()
+        return self
+
+    def add_standard_ti2v_stages(
+        self,
+        *,
+        include_input_validation: bool = True,
+        vae_image_processor: Any | None = None,
+        text_encoder_key: str = "text_encoder",
+        tokenizer_key: str = "tokenizer",
+        image_encoder_key: str = "image_encoder",
+        image_processor_key: str = "image_processor",
+        image_vae_key: str = "vae",
+        image_vae_stage_kwargs: dict[str, Any] | None = None,
+        image_vae_encoding_position: Literal[
+            "before_timestep", "after_latent"
+        ] = "before_timestep",
+        prepare_extra_timestep_kwargs: list[Callable] | None = [],
+        denoising_stage_factory: Callable[[], PipelineStage] | None = None,
+    ) -> "ComposedPipelineBase":
+        if include_input_validation:
+            self.add_stage(
+                InputValidationStage(vae_image_processor=vae_image_processor)
+            )
+
+        self.add_standard_text_encoding_stage(
+            text_encoder_key=text_encoder_key,
+            tokenizer_key=tokenizer_key,
+        )
+
+        image_encoder = self.get_module(image_encoder_key, None)
+        image_processor = self.get_module(image_processor_key, None)
+        self.add_stage_if(
+            image_encoder is not None and image_processor is not None,
+            ImageEncodingStage(
+                image_encoder=image_encoder,
+                image_processor=image_processor,
+            ),
+        )
+
+        if image_vae_encoding_position == "before_timestep":
+            self.add_stage(
+                ImageVAEEncodingStage(
+                    vae=self.get_module(image_vae_key),
+                    **(image_vae_stage_kwargs or {}),
+                )
+            )
+
+        self.add_standard_latent_preparation_stage()
+        self.add_standard_timestep_preparation_stage(
+            prepare_extra_kwargs=prepare_extra_timestep_kwargs
+        )
+        if image_vae_encoding_position == "after_latent":
+            self.add_stage(
+                ImageVAEEncodingStage(
+                    vae=self.get_module(image_vae_key),
+                    **(image_vae_stage_kwargs or {}),
+                )
+            )
+        elif image_vae_encoding_position != "before_timestep":
+            raise ValueError(
+                f"Unknown image_vae_encoding_position: {image_vae_encoding_position}"
+            )
+
+        if denoising_stage_factory is None:
+            self.add_standard_denoising_stage()
+        else:
+            self.add_stage(denoising_stage_factory())
+
+        self.add_standard_decoding_stage()
+        return self
 
     # TODO(will): don't hardcode no_grad
     @torch.no_grad()

python/sglang/multimodal_gen/runtime/pipelines_core/lora_format_adapter.py

@@ -21,11 +21,6 @@ class LoRAFormat(str, Enum):
     WAN = "wan"
 
 
-# ---------------------------------------------------------------------------
-# Helpers
-# ---------------------------------------------------------------------------
-
-
 def _sample_keys(keys: Iterable[str], k: int = 20) -> list[str]:
     out = []
     for i, key in enumerate(keys):
@@ -43,11 +38,6 @@ def _has_prefix_key(keys: Iterable[str], prefix: str) -> bool:
     return any(k.startswith(prefix) for k in keys)
 
 
-# ---------------------------------------------------------------------------
-# Format-specific heuristics
-# ---------------------------------------------------------------------------
-
-
 def _looks_like_xlabs_flux_key(k: str) -> bool:
     """XLabs FLUX-style keys under double_blocks/single_blocks with lora down/up."""
     if not (k.endswith(".down.weight") or k.endswith(".up.weight")):
@@ -114,11 +104,6 @@ def _looks_like_qwen_image(state_dict: Mapping[str, torch.Tensor]) -> bool:
     )
 
 
-# ---------------------------------------------------------------------------
-# Format detection
-# ---------------------------------------------------------------------------
-
-
 def detect_lora_format_from_state_dict(
     state_dict: Mapping[str, torch.Tensor],
 ) -> LoRAFormat:
@@ -150,11 +135,6 @@ def detect_lora_format_from_state_dict(
     return LoRAFormat.STANDARD
 
 
-# ---------------------------------------------------------------------------
-# Converters
-# ---------------------------------------------------------------------------
-
-
 def _convert_qwen_image_standard(
     state_dict: Mapping[str, torch.Tensor],
     log: logging.Logger,
@@ -329,11 +309,6 @@ def _convert_kohya_flux_via_diffusers(
     )
 
 
-# ---------------------------------------------------------------------------
-# Conversion dispatcher
-# ---------------------------------------------------------------------------
-
-
 def convert_lora_state_dict_by_format(
     state_dict: Mapping[str, torch.Tensor],
     fmt: LoRAFormat,
@@ -380,11 +355,6 @@ def convert_lora_state_dict_by_format(
     return dict(state_dict)
 
 
-# ---------------------------------------------------------------------------
-# Public entry point
-# ---------------------------------------------------------------------------
-
-
 def normalize_lora_state_dict(
     state_dict: Mapping[str, torch.Tensor],
     logger: Optional[logging.Logger] = None,

python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py

@@ -15,9 +15,6 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.causal_denoising import
 from sglang.multimodal_gen.runtime.pipelines_core.stages.comfyui_latent_preparation import (
     ComfyUILatentPreparationStage,
 )
-from sglang.multimodal_gen.runtime.pipelines_core.stages.conditioning import (
-    ConditioningStage,
-)
 from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage
 from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding_av import (
     LTX2AVDecodingStage,
@@ -60,7 +57,6 @@ __all__ = [
     "LatentPreparationStage",
     "ComfyUILatentPreparationStage",
     "LTX2AVLatentPreparationStage",
-    "ConditioningStage",
     "DenoisingStage",
     "DmdDenoisingStage",
     "LTX2AVDenoisingStage",

python/sglang/multimodal_gen/runtime/pipelines_core/stages/conditioning.py

@@ -1,105 +0,0 @@
-# Copied and adapted from: https://github.com/hao-ai-lab/FastVideo
-
-# SPDX-License-Identifier: Apache-2.0
-"""
-Conditioning stage for diffusion pipelines.
-"""
-
-import torch
-
-from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
-from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
-from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
-    StageValidators as V,
-)
-from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
-    VerificationResult,
-)
-from sglang.multimodal_gen.runtime.server_args import ServerArgs
-from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
-
-logger = init_logger(__name__)
-
-
-class ConditioningStage(PipelineStage):
-    """
-    Stage for applying conditioning to the diffusion process.
-
-    This stage handles the application of conditioning, such as classifier-free guidance,
-    to the diffusion process.
-    """
-
-    @torch.no_grad()
-    def forward(
-        self,
-        batch: Req,
-        server_args: ServerArgs,
-    ) -> Req:
-        """
-        Apply conditioning to the diffusion process.
-
-
-
-        Returns:
-            The batch with applied conditioning.
-        """
-        # TODO!!
-        if not batch.do_classifier_free_guidance:
-            return batch
-        else:
-            return batch
-
-        logger.info("batch.negative_prompt_embeds: %s", batch.negative_prompt_embeds)
-        logger.info(
-            "do_classifier_free_guidance: %s", batch.do_classifier_free_guidance
-        )
-        logger.info("cfg_scale: %s", batch.guidance_scale)
-
-        # Ensure negative prompt embeddings are available
-        assert (
-            batch.negative_prompt_embeds is not None
-        ), "Negative prompt embeddings are required for classifier-free guidance"
-
-        # Concatenate primary embeddings and masks
-        batch.prompt_embeds = torch.cat(
-            [batch.negative_prompt_embeds, batch.prompt_embeds]
-        )
-        if batch.attention_mask is not None:
-            batch.attention_mask = torch.cat(
-                [batch.negative_attention_mask, batch.attention_mask]
-            )
-
-        # Concatenate secondary embeddings and masks if present
-        if batch.prompt_embeds_2 is not None:
-            batch.prompt_embeds_2 = torch.cat(
-                [batch.negative_prompt_embeds_2, batch.prompt_embeds_2]
-            )
-        if batch.attention_mask_2 is not None:
-            batch.attention_mask_2 = torch.cat(
-                [batch.negative_attention_mask_2, batch.attention_mask_2]
-            )
-
-        return batch
-
-    def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
-        """Verify conditioning stage inputs."""
-        result = VerificationResult()
-        result.add_check(
-            "do_classifier_free_guidance",
-            batch.do_classifier_free_guidance,
-            V.bool_value,
-        )
-        result.add_check("guidance_scale", batch.guidance_scale, V.non_negative_float)
-        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
-        result.add_check(
-            "negative_prompt_embeds",
-            batch.negative_prompt_embeds,
-            lambda x: not batch.do_classifier_free_guidance or V.list_not_empty(x),
-        )
-        return result
-
-    def verify_output(self, batch: Req, server_args: ServerArgs) -> VerificationResult:
-        """Verify conditioning stage outputs."""
-        result = VerificationResult()
-        result.add_check("prompt_embeds", batch.prompt_embeds, V.list_not_empty)
-        return result

python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py

@@ -47,7 +47,9 @@ class TimestepPreparationStage(PipelineStage):
     ) -> None:
         super().__init__()
         self.scheduler = scheduler
-        self.prepare_extra_set_timesteps_kwargs = prepare_extra_set_timesteps_kwargs
+        self.prepare_extra_set_timesteps_kwargs = (
+            prepare_extra_set_timesteps_kwargs or []
+        )
 
     @property
     def parallelism_type(self) -> StageParallelismType:

python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json

@@ -9,7 +9,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.1,
                 "TextEncodingStage": 1609.27,
-                "ConditioningStage": 0.02,
                 "TimestepPreparationStage": 3.46,
                 "LatentPreparationStage": 0.39,
                 "DenoisingStage": 26324.0,

python/sglang/multimodal_gen/test/server/perf_baselines.json

@@ -37,7 +37,6 @@
         "qwen_image_t2i": {
             "notes": "Single-image generation using the default prompt",
             "stages_ms": {
-                "ConditioningStage": 0.02,
                 "DecodingStage": 51.86,
                 "TextEncodingStage": 611.83,
                 "InputValidationStage": 0.05,
@@ -105,7 +104,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.04,
                 "TextEncodingStage": 693.2,
-                "ConditioningStage": 0.02,
                 "TimestepPreparationStage": 2.84,
                 "LatentPreparationStage": 9.13,
                 "DenoisingStage": 24529.77,
@@ -169,7 +167,6 @@
         },
         "flux_image_t2i": {
             "stages_ms": {
-                "ConditioningStage": 0.01,
                 "DecodingStage": 32.72,
                 "TextEncodingStage": 51.96,
                 "InputValidationStage": 0.03,
@@ -239,7 +236,6 @@
                 "TimestepPreparationStage": 2.91,
                 "TextEncodingStage": 518.54,
                 "ImageVAEEncodingStage": 0.0,
-                "ConditioningStage": 0.02,
                 "InputValidationStage": 0.05,
                 "DenoisingStage": 24901.97,
                 "DecodingStage": 8.98
@@ -302,7 +298,6 @@
         },
         "flux_2_klein_image_t2i": {
             "stages_ms": {
-                "ConditioningStage": 0.01,
                 "DecodingStage": 9.27,
                 "TextEncodingStage": 92.17,
                 "InputValidationStage": 0.05,
@@ -326,7 +321,6 @@
                 "InputValidationStage": 0.06,
                 "TextEncodingStage": 513.58,
                 "ImageVAEEncodingStage": 0.0,
-                "ConditioningStage": 0.03,
                 "LatentPreparationStage": 0.46,
                 "TimestepPreparationStage": 2.38,
                 "DenoisingStage": 52187.62,
@@ -393,7 +387,6 @@
                 "InputValidationStage": 99.82,
                 "TextEncodingStage": 519.88,
                 "ImageVAEEncodingStage": 254.56,
-                "ConditioningStage": 0.01,
                 "LatentPreparationStage": 12.4,
                 "TimestepPreparationStage": 2.71,
                 "DenoisingStage": 54705.41,
@@ -459,7 +452,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.03,
                 "TextEncodingStage": 74.47,
-                "ConditioningStage": 0.01,
                 "TimestepPreparationStage": 2.23,
                 "LatentPreparationStage": 6.17,
                 "DenoisingStage": 8400.49,
@@ -525,7 +517,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.03,
                 "TextEncodingStage": 403.47,
-                "ConditioningStage": 0.01,
                 "TimestepPreparationStage": 1.41,
                 "LatentPreparationStage": 0.11,
                 "DenoisingStage": 756.21,
@@ -550,7 +541,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.04,
                 "TextEncodingStage": 413.69,
-                "ConditioningStage": 0.01,
                 "TimestepPreparationStage": 1.3,
                 "LatentPreparationStage": 0.11,
                 "DenoisingStage": 813.7,
@@ -575,7 +565,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.08,
                 "TextEncodingStage": 420.74,
-                "ConditioningStage": 0.01,
                 "TimestepPreparationStage": 1.5,
                 "LatentPreparationStage": 0.12,
                 "DenoisingStage": 1304.07,
@@ -602,7 +591,6 @@
                 "TimestepPreparationStage": 2.62,
                 "ImageEncodingStage": 1174.26,
                 "ImageVAEEncodingStage": 132.67,
-                "ConditioningStage": 0.01,
                 "InputValidationStage": 38.1,
                 "DenoisingStage": 38135.64,
                 "DecodingStage": 139.72
@@ -667,7 +655,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.05,
                 "TextEncodingStage": 675.95,
-                "ConditioningStage": 0.02,
                 "TimestepPreparationStage": 3.21,
                 "LatentPreparationStage": 0.2,
                 "DenoisingStage": 5248.83,
@@ -733,74 +720,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.07,
                 "TextEncodingStage": 2237.78,
-                "ConditioningStage": 0.01,
-                "TimestepPreparationStage": 2.1,
-                "LatentPreparationStage": 0.84,
-                "DenoisingStage": 13041.23,
-                "DecodingStage": 1274.63,
-                "per_frame_generation": null
-            },
-            "denoise_step_ms": {
-                "0": 240.71,
-                "1": 248.13,
-                "2": 246.48,
-                "3": 247.87,
-                "4": 249.38,
-                "5": 246.76,
-                "6": 250.42,
-                "7": 250.81,
-                "8": 250.98,
-                "9": 249.9,
-                "10": 246.72,
-                "11": 249.79,
-                "12": 250.46,
-                "13": 249.19,
-                "14": 247.55,
-                "15": 250.12,
-                "16": 247.57,
-                "17": 247.21,
-                "18": 247.32,
-                "19": 247.42,
-                "20": 248.21,
-                "21": 247.19,
-                "22": 247.72,
-                "23": 247.45,
-                "24": 247.9,
-                "25": 247.87,
-                "26": 247.18,
-                "27": 247.65,
-                "28": 246.91,
-                "29": 248.26,
-                "30": 247.82,
-                "31": 247.73,
-                "32": 247.38,
-                "33": 247.84,
-                "34": 247.46,
-                "35": 247.52,
-                "36": 247.94,
-                "37": 248.76,
-                "38": 248.01,
-                "39": 247.45,
-                "40": 247.84,
-                "41": 248.33,
-                "42": 247.41,
-                "43": 248.16,
-                "44": 248.18,
-                "45": 248.44,
-                "46": 248.65,
-                "47": 247.73,
-                "48": 247.48,
-                "49": 247.54
-            },
-            "expected_e2e_ms": 18382.19,
-            "expected_avg_denoise_ms": 260.76,
-            "expected_median_denoise_ms": 247.84
-        },
-        "wan2_1_t2v_1.3b_teacache_enabled": {
-            "stages_ms": {
-                "InputValidationStage": 0.07,
-                "TextEncodingStage": 2237.78,
-                "ConditioningStage": 0.01,
                 "TimestepPreparationStage": 2.1,
                 "LatentPreparationStage": 0.84,
                 "DenoisingStage": 13041.23,
@@ -867,7 +786,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.07,
                 "TextEncodingStage": 2237.78,
-                "ConditioningStage": 0.01,
                 "TimestepPreparationStage": 2.1,
                 "LatentPreparationStage": 0.84,
                 "DenoisingStage": 13041.23,
@@ -934,7 +852,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.09,
                 "TextEncodingStage": 2480.54,
-                "ConditioningStage": 0.07,
                 "TimestepPreparationStage": 3.73,
                 "LatentPreparationStage": 1.34,
                 "DenoisingStage": 12514.88,
@@ -1001,7 +918,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.08,
                 "TextEncodingStage": 2700.44,
-                "ConditioningStage": 0.02,
                 "TimestepPreparationStage": 2.82,
                 "LatentPreparationStage": 2.0,
                 "DenoisingStage": 11640.75,
@@ -1068,7 +984,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.06,
                 "TextEncodingStage": 2508.95,
-                "ConditioningStage": 0.04,
                 "TimestepPreparationStage": 73.51,
                 "LatentPreparationStage": 1.34,
                 "DmdDenoisingStage": 1285.25,
@@ -1089,7 +1004,6 @@
             "stages_ms": {
                 "InputValidationStage": 96.27,
                 "TextEncodingStage": 2238.81,
-                "ConditioningStage": 0.02,
                 "TimestepPreparationStage": 2.39,
                 "LatentPreparationStage": 27.62,
                 "DenoisingStage": 134069.79,
@@ -1159,7 +1073,6 @@
                 "ImageVAEEncodingStage": 304.56,
                 "TimestepPreparationStage": 2.94,
                 "LatentPreparationStage": 0.2,
-                "ConditioningStage": 0.01,
                 "DenoisingStage": 50724.5,
                 "DecodingStage": 601.02
             },
@@ -1276,7 +1189,6 @@
             "stages_ms": {
                 "InputValidationStage": 300.00,
                 "TextEncodingStage": 843.86,
-                "ConditioningStage": 0.01,
                 "TimestepPreparationStage": 58.66,
                 "LatentPreparationStage": 28.55,
                 "DmdDenoisingStage": 499.34,
@@ -1296,7 +1208,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.34,
                 "TextEncodingStage": 550.63,
-                "ConditioningStage": 0.02,
                 "TimestepPreparationStage": 44.28,
                 "LatentPreparationStage": 0.29,
                 "DenoisingStage": 9154.39,
@@ -1319,7 +1230,6 @@
             "stages_ms": {
                 "InputValidationStage": 18.45,
                 "TextEncodingStage": 3337.77,
-                "ConditioningStage": 0.03,
                 "TimestepPreparationStage": 2.9,
                 "LatentPreparationStage": 1.25,
                 "ImageVAEEncodingStage": 1655.89,
@@ -1377,7 +1287,6 @@
             "stages_ms": {
                 "InputValidationStage": 25.01,
                 "TextEncodingStage": 5198.6,
-                "ConditioningStage": 0.04,
                 "TimestepPreparationStage": 56.26,
                 "LatentPreparationStage": 1.4,
                 "ImageVAEEncodingStage": 1001.89,
@@ -1400,7 +1309,6 @@
                 "InputValidationStage": 38.23,
                 "TextEncodingStage": 3550.36,
                 "ImageEncodingStage": 3462.55,
-                "ConditioningStage": 0.01,
                 "TimestepPreparationStage": 2.6,
                 "LatentPreparationStage": 9.73,
                 "ImageVAEEncodingStage": 2290.98,
@@ -1469,7 +1377,6 @@
                 "InputValidationStage": 53.67,
                 "TextEncodingStage": 2838,
                 "ImageEncodingStage": 3123.99,
-                "ConditioningStage": 0.01,
                 "TimestepPreparationStage": 3.39,
                 "LatentPreparationStage": 8.41,
                 "ImageVAEEncodingStage": 2261.05,
@@ -1536,7 +1443,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.07,
                 "TextEncodingStage": 2575.3,
-                "ConditioningStage": 0.01,
                 "TimestepPreparationStage": 1.99,
                 "LatentPreparationStage": 1.26,
                 "DenoisingStage": 156678.8406,
@@ -1593,7 +1499,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.05,
                 "TextEncodingStage": 2310.34,
-                "ConditioningStage": 0.02,
                 "TimestepPreparationStage": 2.42,
                 "LatentPreparationStage": 27.7,
                 "DenoisingStage": 803631.52,
@@ -1660,7 +1565,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.09,
                 "TextEncodingStage": 2552.97,
-                "ConditioningStage": 0.03,
                 "TimestepPreparationStage": 1.99,
                 "LatentPreparationStage": 1.29,
                 "DenoisingStage": 154340.69,
@@ -1717,7 +1621,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.06,
                 "TextEncodingStage": 2467.44,
-                "ConditioningStage": 0.02,
                 "TimestepPreparationStage": 2.96,
                 "LatentPreparationStage": 1.87,
                 "DenoisingStage": 14859.47,
@@ -1785,7 +1688,6 @@
                 "InputValidationStage": 23.97,
                 "TextEncodingStage": 2485.39,
                 "ImageEncodingStage": 2372.07,
-                "ConditioningStage": 0.01,
                 "TimestepPreparationStage": 2.6,
                 "LatentPreparationStage": 0.18,
                 "ImageVAEEncodingStage": 2500.13,
@@ -1854,7 +1756,6 @@
                 "InputValidationStage": 0.05,
                 "TextEncodingStage": 518.88,
                 "ImageVAEEncodingStage": 0.0,
-                "ConditioningStage": 0.01,
                 "LatentPreparationStage": 0.45,
                 "TimestepPreparationStage": 3.41,
                 "DenoisingStage": 26377.63,
@@ -1923,7 +1824,6 @@
                 "ImageVAEEncodingStage": 88.06,
                 "TimestepPreparationStage": 2.12,
                 "LatentPreparationStage": 0.14,
-                "ConditioningStage": 0.01,
                 "DenoisingStage": 23869.32,
                 "DecodingStage": 108.23
             },
@@ -1977,7 +1877,6 @@
             "stages_ms": {
                 "InputValidationStage": 0.04,
                 "TextEncodingStage": 411.12,
-                "ConditioningStage": 0.01,
                 "TimestepPreparationStage": 1.44,
                 "LatentPreparationStage": 0.1,
                 "DenoisingStage": 1569.61,

python/sglang/multimodal_gen/test/server/testcase_configs.py

@@ -10,7 +10,7 @@ pytest python/sglang/multimodal_gen/test/server/test_server_a.py -k qwen_image_t
 
 To add a new testcase:
 1. add your testcase with case-id: `my_new_test_case_id` to DIFFUSION_CASES
-2. run `SGLANG_GEN_BASELINE=1 pytest -s python/sglang/multimodal_gen/test/server/test_server_a.py -k my_new_test_case_id`
+2. run `SGLANG_GEN_BASELINE=1 pytest -s python/sglang/multimodal_gen/test/server/ -k my_new_test_case_id`
 3. insert or override the corresponding scenario in `scenarios` section of perf_baselines.json with the output baseline of step-2