chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Mick
2026-02-21 16:35:47 +08:00
committed by GitHub
parent 33c33a7de9
commit b89ca65789
26 changed files with 444 additions and 1029 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
python/sglang/multimodal_gen/runtime/pipelines/comfyui_flux_pipeline.py
View File

@@ -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())}"
)

27
python/sglang/multimodal_gen/runtime/pipelines/comfyui_qwen_image_pipeline.py
View File

@@ -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())}"
)

27
python/sglang/multimodal_gen/runtime/pipelines/comfyui_zimage_pipeline.py
View File

@@ -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())}"
)

14
python/sglang/multimodal_gen/runtime/pipelines/diffusers_pipeline.py
View File

@@ -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]:

48
python/sglang/multimodal_gen/runtime/pipelines/flux.py
View File

@@ -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

80
python/sglang/multimodal_gen/runtime/pipelines/flux_2.py
View File

@@ -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],
)

16
python/sglang/multimodal_gen/runtime/pipelines/glm_image.py
View File

@@ -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]

47
python/sglang/multimodal_gen/runtime/pipelines/hunyuan_pipeline.py
View File

@@ -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

95
python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py
View File

@@ -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,
),
]
)

36
python/sglang/multimodal_gen/runtime/pipelines/mova_pipeline.py
View File

@@ -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",
)

152
python/sglang/multimodal_gen/runtime/pipelines/qwen_image.py
View File

@@ -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 = [

37
python/sglang/multimodal_gen/runtime/pipelines/wan_causal_dmd_pipeline.py
View File

@@ -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

58
python/sglang/multimodal_gen/runtime/pipelines/wan_dmd_pipeline.py
View File

@@ -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

77
python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_dmd_pipeline.py
View File

@@ -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

80
python/sglang/multimodal_gen/runtime/pipelines/wan_i2v_pipeline.py
View File

@@ -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

54
python/sglang/multimodal_gen/runtime/pipelines/wan_pipeline.py
View File

@@ -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

57
python/sglang/multimodal_gen/runtime/pipelines/zimage_pipeline.py
View File

@@ -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

277
python/sglang/multimodal_gen/runtime/pipelines_core/composed_pipeline_base.py
View File

@@ -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"(?<!^)(?=[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()

30
python/sglang/multimodal_gen/runtime/pipelines_core/lora_format_adapter.py
View File

@@ -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,

4
python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py
View File

@@ -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",

105
python/sglang/multimodal_gen/runtime/pipelines_core/stages/conditioning.py
View File

@@ -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

4
python/sglang/multimodal_gen/runtime/pipelines_core/stages/timestep_preparation.py
View File

@@ -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:

1
python/sglang/multimodal_gen/test/server/ascend/perf_baselines_npu.json
View File

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

101
python/sglang/multimodal_gen/test/server/perf_baselines.json
View File

@@ -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,

2
python/sglang/multimodal_gen/test/server/testcase_configs.py
View File

@@ -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