From b89ca65789129567321a865165fae5bd56a8acda Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 21 Feb 2026 16:35:47 +0800 Subject: [PATCH] [diffusion] refactor: reduce redundancy and improve stage api (#19060) --- docs/diffusion/support_new_models.md | 17 +- .../pipelines/comfyui_flux_pipeline.py | 27 +- .../pipelines/comfyui_qwen_image_pipeline.py | 27 +- .../pipelines/comfyui_zimage_pipeline.py | 27 +- .../runtime/pipelines/diffusers_pipeline.py | 14 +- .../multimodal_gen/runtime/pipelines/flux.py | 48 +-- .../runtime/pipelines/flux_2.py | 80 +---- .../runtime/pipelines/glm_image.py | 16 +- .../runtime/pipelines/hunyuan_pipeline.py | 47 +-- .../runtime/pipelines/ltx_2_pipeline.py | 95 ++---- .../runtime/pipelines/mova_pipeline.py | 36 +-- .../runtime/pipelines/qwen_image.py | 152 +--------- .../pipelines/wan_causal_dmd_pipeline.py | 37 +-- .../runtime/pipelines/wan_dmd_pipeline.py | 58 ++-- .../runtime/pipelines/wan_i2v_dmd_pipeline.py | 77 +---- .../runtime/pipelines/wan_i2v_pipeline.py | 80 +---- .../runtime/pipelines/wan_pipeline.py | 54 +--- .../runtime/pipelines/zimage_pipeline.py | 57 +--- .../pipelines_core/composed_pipeline_base.py | 277 +++++++++++++++++- .../pipelines_core/lora_format_adapter.py | 30 -- .../runtime/pipelines_core/stages/__init__.py | 4 - .../pipelines_core/stages/conditioning.py | 105 ------- .../stages/timestep_preparation.py | 4 +- .../server/ascend/perf_baselines_npu.json | 1 - .../test/server/perf_baselines.json | 101 ------- .../test/server/testcase_configs.py | 2 +- 26 files changed, 444 insertions(+), 1029 deletions(-) delete mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/conditioning.py diff --git a/docs/diffusion/support_new_models.md b/docs/diffusion/support_new_models.md index b0f763243..3141d5aff 100644 --- a/docs/diffusion/support_new_models.md +++ b/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(stage_name="input_validation_stage", stage=InputValidationStage()) - self.add_stage(stage_name="prompt_encoding_stage_primary", stage=ImageEncodingStage(...)) - self.add_stage(stage_name="image_encoding_stage_primary", stage=ImageVAEEncodingStage(...)) - self.add_stage(stage_name="timestep_preparation_stage", stage=TimestepPreparationStage(...)) - self.add_stage(stage_name="latent_preparation_stage", stage=LatentPreparationStage(...)) - self.add_stage(stage_name="conditioning_stage", stage=ConditioningStage()) - self.add_stage(stage_name="denoising_stage", stage=DenoisingStage(...)) - self.add_stage(stage_name="decoding_stage", stage=DecodingStage(...)) + self.add_stage(InputValidationStage()) + self.add_stage(ImageEncodingStage(...)) + self.add_stage(ImageVAEEncodingStage(...)) + self.add_stage(TimestepPreparationStage(...)) + self.add_stage(LatentPreparationStage(...)) + self.add_stage(DenoisingStage(...)) + self.add_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. diff --git a/python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py index 233456306..f69101877 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py +++ b/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 - # This stage includes device mismatch fix for ComfyUI pipelines in multi-GPU scenarios - self.add_stage( - stage_name="latent_preparation_stage", - stage=ComfyUILatentPreparationStage( - scheduler=self.get_module("scheduler"), - transformer=self.get_module("transformer"), - ), + self.add_stages( + [ + 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())}" ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/comfyui_qwen_image_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/comfyui_qwen_image_pipeline.py index 0e2c031f7..138f74b28 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/comfyui_qwen_image_pipeline.py +++ b/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 - # This stage includes device mismatch fix for ComfyUI pipelines in multi-GPU scenarios - self.add_stage( - stage_name="latent_preparation_stage", - stage=ComfyUILatentPreparationStage( - scheduler=self.get_module("scheduler"), - transformer=self.get_module("transformer"), - ), + self.add_stages( + [ + 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())}" ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/comfyui_zimage_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/comfyui_zimage_pipeline.py index b58f86839..efe5e7430 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/comfyui_zimage_pipeline.py +++ b/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 - # This stage includes device mismatch fix for ComfyUI pipelines in multi-GPU scenarios - self.add_stage( - stage_name="latent_preparation_stage", - stage=ComfyUILatentPreparationStage( - scheduler=self.get_module("scheduler"), - transformer=self.get_module("transformer"), - ), + self.add_stages( + [ + 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())}" ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py index baa61562b..a45e6d6ae 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py +++ b/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", - stage=DiffusersExecutionStage(self.diffusers_pipe), + DiffusersExecutionStage(self.diffusers_pipe), "diffusers_execution" ) 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]: diff --git a/python/sglang/multimodal_gen/runtime/pipelines/flux.py b/python/sglang/multimodal_gen/runtime/pipelines/flux.py index d738c8ba6..371e1c44e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/flux.py +++ b/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() - ) - - self.add_stage( - stage_name="prompt_encoding_stage_primary", - stage=TextEncodingStage( + 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_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")) - ) + self.add_standard_timestep_preparation_stage(prepare_extra_kwargs=[prepare_mu]) + self.add_standard_latent_preparation_stage() + self.add_standard_denoising_stage() + self.add_standard_decoding_stage() EntryClass = FluxPipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines/flux_2.py b/python/sglang/multimodal_gen/runtime/pipelines/flux_2.py index 58ce257af..78beb90e3 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/flux_2.py +++ b/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.""" - - 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 - ), - ), + vae_image_processor = VaeImageProcessor( + vae_scale_factor=server_args.pipeline_config.vae_config.arch_config.vae_scale_factor + * 2 ) - 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="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")) + self.add_standard_ti2i_stages( + include_input_validation=True, + vae_image_processor=vae_image_processor, + prompt_encoding="text", + image_vae_stage_kwargs={"vae_image_processor": vae_image_processor}, + prepare_extra_timestep_kwargs=[compute_empirical_mu], ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/glm_image.py b/python/sglang/multimodal_gen/runtime/pipelines/glm_image.py index 4668eaf8b..df3e58e36 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/glm_image.py +++ b/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 ( - DecodingStage, - DenoisingStage, -) +from sglang.multimodal_gen.runtime.pipelines_core.stages import 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", - stage=GlmImageBeforeDenoisingStage( + 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", - stage=DenoisingStage( + 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_decoding_stage() EntryClass = [GlmImagePipeline] diff --git a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py index 9a829ab14..301c759f4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py +++ b/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() - ) - - self.add_stage( - stage_name="prompt_encoding_stage_primary", - stage=TextEncodingStage( + 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_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="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_timestep_preparation_stage() + self.add_standard_latent_preparation_stage() + self.add_standard_denoising_stage() + self.add_standard_decoding_stage() EntryClass = HunyuanVideoPipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py index 29f3195aa..e5049e4d7 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/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.""" - - # 1. Input Validation - self.add_stage( - stage_name="input_validation_stage", stage=InputValidationStage() + self.add_stages( + [ + InputValidationStage(), + TextEncodingStage( + text_encoders=[self.get_module("text_encoder")], + tokenizers=[self.get_module("tokenizer")], + ), + LTX2TextConnectorStage(connectors=self.get_module("connectors")), + ] ) - # 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"), - ], - ), - ) + self.add_standard_timestep_preparation_stage(prepare_extra_kwargs=[prepare_mu]) - # 3. connector stage - self.add_stage( - stage_name="text_connector_stage", - stage=LTX2TextConnectorStage(connectors=self.get_module("connectors")), - ) - - # 4. Timestep Preparation - self.add_stage( - stage_name="timestep_preparation_stage", - 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"), - ), - ) - - # 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"), - ), - ) - - # 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, - ), + self.add_stages( + [ + LTX2AVLatentPreparationStage( + scheduler=self.get_module("scheduler"), + transformer=self.get_module("transformer"), + audio_vae=self.get_module("audio_vae"), + ), + LTX2AVDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + vae=self.get_module("vae"), + audio_vae=self.get_module("audio_vae"), + ), + LTX2AVDecodingStage( + vae=self.get_module("vae"), + audio_vae=self.get_module("audio_vae"), + vocoder=self.get_module("vocoder"), + pipeline=self, + ), + ] ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/mova_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/mova_pipeline.py index e41353fcc..57e1c5029 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/mova_pipeline.py +++ b/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( - 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(InputValidationStage()) + self.add_standard_text_encoding_stage() if getattr(self.get_module("video_dit"), "require_vae_embedding", True): - self.add_stage( - stage_name="image_latent_preparation_stage", - stage=ImageVAEEncodingStage(vae=self.get_module("video_vae")), - ) + self.add_stage(ImageVAEEncodingStage(vae=self.get_module("video_vae"))) self.add_stage( - stage_name="mova_latent_preparation_stage", - stage=MOVALatentPreparationStage( + 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", - stage=MOVATimestepPreparationStage( + MOVATimestepPreparationStage( scheduler=self.get_module("scheduler"), ), + "mova_timestep_preparation_stage", ) self.add_stage( - stage_name="mova_denoising_stage", - stage=MOVADenoisingStage( + 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", - stage=MOVADecodingStage( + MOVADecodingStage( video_vae=self.get_module("video_vae"), audio_vae=self.get_module("audio_vae"), ), + "mova_decoding_stage", ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py b/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py index b69944f14..81695190a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py +++ b/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_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")) - ) + self.add_standard_t2i_stages(prepare_extra_timestep_kwargs=[prepare_mu]) 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.""" - - 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 - ) - ), + vae_image_processor = VaeImageProcessor( + vae_scale_factor=server_args.pipeline_config.vae_config.arch_config.vae_scale_factor + * 2 ) - self.add_stage( - stage_name="prompt_encoding_stage_primary", - stage=ImageEncodingStage( - image_processor=self.get_module("processor"), - text_encoder=self.get_module("text_encoder"), - ), - ) - - 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")) + self.add_standard_ti2i_stages( + vae_image_processor=vae_image_processor, + prompt_encoding="image_encoding", + image_processor_key="processor", + prompt_text_encoder_key="text_encoder", + prepare_extra_timestep_kwargs=[prepare_mu], ) @@ -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", - stage=QwenImageLayeredBeforeDenoisingStage( + 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( - stage_name="timestep_preparation_stage", - 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_timestep_preparation_stage( + prepare_extra_kwargs=[prepare_mu_layered] ) + self.add_standard_denoising_stage() + self.add_standard_decoding_stage() EntryClass = [ diff --git a/python/sglang/multimodal_gen/runtime/pipelines/wan_causal_dmd_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/wan_causal_dmd_pipeline.py index b103ee0a5..97f1f050b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/wan_causal_dmd_pipeline.py +++ b/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() - ) - - 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( + CausalDMDDenoisingStage( 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_decoding_stage() EntryClass = WanCausalDMDPipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines/wan_dmd_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/wan_dmd_pipeline.py index 3c973834c..73017f5b6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/wan_dmd_pipeline.py +++ b/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_stage( - stage_name="input_validation_stage", stage=InputValidationStage() + self.add_stages( + [ + 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_standard_text_encoding_stage() + + self.add_standard_timestep_preparation_stage() + self.add_standard_latent_preparation_stage() + + self.add_stages( + [ + DmdDenoisingStage( + transformer=self.get_module("transformer"), + scheduler=self.get_module("scheduler"), + ), + ] ) - 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=DmdDenoisingStage( - 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_decoding_stage() EntryClass = WanDMDPipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_dmd_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_dmd_pipeline.py index 23d560711..cb954d327 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_dmd_pipeline.py +++ b/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_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=DmdDenoisingStage( + self.add_standard_ti2v_stages( + image_vae_encoding_position="after_latent", + denoising_stage_factory=lambda: 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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_pipeline.py index 93a196870..984da085a 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_pipeline.py +++ b/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_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")) - ) + self.add_standard_ti2v_stages() EntryClass = WanImageToVideoPipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py index 8f1cbfc26..b52045754 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py +++ b/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_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), - ) + self.add_standard_t2i_stages() EntryClass = WanPipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py index f8fd441de..8f1f71474 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py +++ b/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_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")) - ) + self.add_standard_t2i_stages(prepare_extra_timestep_kwargs=[prepare_mu]) EntryClass = ZImagePipeline diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py b/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py index 62d69b0ce..be3ed50a4 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py +++ b/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 default_value - return self.modules[module_name] + return self.modules.get(module_name, default_value) 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"(? "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() diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_format_adapter.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_format_adapter.py index 656d795ab..889ceea7d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_format_adapter.py +++ b/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, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py index 34ff8658e..bc43cbf61 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py +++ b/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", diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/conditioning.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/conditioning.py deleted file mode 100644 index bdb403d20..000000000 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/conditioning.py +++ /dev/null @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py index ba3f5ea55..44cf8c119 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py +++ b/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: diff --git a/python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json b/python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json index 81a6f32b0..b4662e524 100644 --- a/python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json +++ b/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, diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 970089f62..41b7140a9 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/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, diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index c1920594b..f87caea7c 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/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