From af550256441980d7d0d4ac50641b90a481e2c01f Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 10 Jun 2026 09:39:06 +0800 Subject: [PATCH] [diffusion] refactor: refactor realtime and model-specific stage modules (#27697) --- .../runtime/pipelines/hunyuan3d_pipeline.py | 2 +- .../lingbot_world_causal_dmd_pipeline.py | 10 ++- .../runtime/pipelines/ltx_2_pipeline.py | 10 ++- .../pipelines/sana_wm_realtime_pipeline.py | 8 +- .../runtime/pipelines_core/stages/__init__.py | 79 ------------------- .../stages/latent_preparation.py | 56 ------------- .../hunyuan3d/__init__.py | 25 ++++++ .../hunyuan3d/paint.py} | 0 .../hunyuan3d/shape.py} | 0 .../model_specific_stages/ltx_2/__init__.py | 37 +++++++++ .../ltx_2}/decoding_av.py | 0 .../ltx_2/denoising.py} | 0 .../ltx_2}/denoising_av.py | 2 +- .../ltx_2}/latent_preparation_av.py | 0 .../ltx_2}/text_connector.py | 0 .../ltx_2}/upsampling.py | 0 .../sana_wm/realtime_stage.py | 2 +- .../stages/realtime/__init__.py | 31 ++++++++ .../base.py} | 0 .../input_validation.py} | 0 .../stages/realtime/latent_preparation.py | 74 +++++++++++++++++ .../text_encoding.py} | 71 +++++++++-------- .../{realtime_vae.py => realtime/vae.py} | 0 .../unit/realtime/test_realtime_runtime.py | 16 ++-- .../test/unit/realtime/test_realtime_vae.py | 10 ++- .../test/unit/test_disagg_roles.py | 18 ++--- 26 files changed, 251 insertions(+), 200 deletions(-) create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/__init__.py rename python/sglang/multimodal_gen/runtime/pipelines_core/stages/{hunyuan3d_paint.py => model_specific_stages/hunyuan3d/paint.py} (100%) rename python/sglang/multimodal_gen/runtime/pipelines_core/stages/{hunyuan3d_shape.py => model_specific_stages/hunyuan3d/shape.py} (100%) create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/__init__.py rename python/sglang/multimodal_gen/runtime/pipelines_core/stages/{ => model_specific_stages/ltx_2}/decoding_av.py (100%) rename python/sglang/multimodal_gen/runtime/pipelines_core/stages/{ltx_2_denoising.py => model_specific_stages/ltx_2/denoising.py} (100%) rename python/sglang/multimodal_gen/runtime/pipelines_core/stages/{ => model_specific_stages/ltx_2}/denoising_av.py (99%) rename python/sglang/multimodal_gen/runtime/pipelines_core/stages/{ => model_specific_stages/ltx_2}/latent_preparation_av.py (100%) rename python/sglang/multimodal_gen/runtime/pipelines_core/stages/{ => model_specific_stages/ltx_2}/text_connector.py (100%) rename python/sglang/multimodal_gen/runtime/pipelines_core/stages/{ => model_specific_stages/ltx_2}/upsampling.py (100%) create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/__init__.py rename python/sglang/multimodal_gen/runtime/pipelines_core/stages/{realtime_diffusion.py => realtime/base.py} (100%) rename python/sglang/multimodal_gen/runtime/pipelines_core/stages/{realtime_input_validation.py => realtime/input_validation.py} (100%) create mode 100644 python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/latent_preparation.py rename python/sglang/multimodal_gen/runtime/pipelines_core/stages/{realtime_text_encoding.py => realtime/text_encoding.py} (64%) rename python/sglang/multimodal_gen/runtime/pipelines_core/stages/{realtime_vae.py => realtime/vae.py} (100%) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py index 5842f65c6..c36ef347f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py @@ -28,7 +28,7 @@ from sglang.multimodal_gen.runtime.loader.utils import get_param_names_mapping from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) -from sglang.multimodal_gen.runtime.pipelines_core.stages import ( +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.hunyuan3d import ( Hunyuan3DPaintPostprocessStage, Hunyuan3DPaintPreprocessStage, Hunyuan3DPaintTexGenStage, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/lingbot_world_causal_dmd_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/lingbot_world_causal_dmd_pipeline.py index e3f1300f5..4ae196c8e 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/lingbot_world_causal_dmd_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/lingbot_world_causal_dmd_pipeline.py @@ -15,17 +15,19 @@ from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipeline from sglang.multimodal_gen.runtime.pipelines_core.stages import ( AuxiliaryConditionEncodingStage, - CausalVaeDecodingStage, DMDTimestepPreparationStage, ImageEncodingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world import ( + LingBotWorldCausalDMDDenoisingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime import ( + CausalVaeDecodingStage, RealtimeChunkLatentPreparationStage, RealtimeImageVAEEncodingStage, RealtimeInputValidationStage, RealtimeTextEncodingStage, ) -from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world import ( - LingBotWorldCausalDMDDenoisingStage, -) from sglang.multimodal_gen.runtime.server_args import ServerArgs 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 e8b27be0a..463acad85 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -35,18 +35,22 @@ from sglang.multimodal_gen.runtime.pipelines_core.lora_pipeline import LoRAPipel from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages import ( InputValidationStage, + TextEncodingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage +from sglang.multimodal_gen.runtime.pipelines_core.stages.image_encoding import ( + LTX2ImageEncodingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2 import ( LTX2AVDecodingStage, LTX2AVDenoisingStage, LTX2AVLatentPreparationStage, LTX2HalveResolutionStage, - LTX2ImageEncodingStage, LTX2LoRASwitchStage, LTX2RefinementStage, LTX2TextConnectorStage, LTX2UpsampleStage, - TextEncodingStage, ) -from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ( LTX2_RESIDENT_AUTO_ENABLE_MEM_GB, diff --git a/python/sglang/multimodal_gen/runtime/pipelines/sana_wm_realtime_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/sana_wm_realtime_pipeline.py index 31f2ea6b2..2c1f8e3a1 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/sana_wm_realtime_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/sana_wm_realtime_pipeline.py @@ -7,10 +7,6 @@ from sglang.multimodal_gen.configs.pipeline_configs.sana_wm import ( from sglang.multimodal_gen.runtime.pipelines.sana_wm_pipeline import ( SanaWMTwoStagePipeline, ) -from sglang.multimodal_gen.runtime.pipelines_core.stages import ( - RealtimeInputValidationStage, - RealtimeTextEncodingStage, -) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm import ( SanaWMTextEncodingStage, ) @@ -31,6 +27,10 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.s from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.sana_wm.streaming_refiner import ( SanaWMStreamingRefinerStage, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime import ( + RealtimeInputValidationStage, + RealtimeTextEncodingStage, +) from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.hf_diffusers_utils import maybe_download_model 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 1936eaddf..7c124b21f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py @@ -20,67 +20,20 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.condition_encoding impo ConditionEncodingStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding import DecodingStage -from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding_av import ( - LTX2AVDecodingStage, -) from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising import DenoisingStage -from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_av import ( - LTX2AVDenoisingStage, - LTX2RefinementStage, -) from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_dmd import ( DmdDenoisingStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.encoding import EncodingStage - -# Hunyuan3D paint stages -from sglang.multimodal_gen.runtime.pipelines_core.stages.hunyuan3d_paint import ( - Hunyuan3DPaintPostprocessStage, - Hunyuan3DPaintPreprocessStage, - Hunyuan3DPaintTexGenStage, -) - -# Hunyuan3D shape stages -from sglang.multimodal_gen.runtime.pipelines_core.stages.hunyuan3d_shape import ( - Hunyuan3DShapeBeforeDenoisingStage, - Hunyuan3DShapeDenoisingStage, - Hunyuan3DShapeExportStage, - Hunyuan3DShapeSaveStage, -) from sglang.multimodal_gen.runtime.pipelines_core.stages.image_encoding import ( ImageEncodingStage, ImageVAEEncodingStage, - LTX2ImageEncodingStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.input_validation import ( InputValidationStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import ( LatentPreparationStage, - RealtimeChunkLatentPreparationStage, -) -from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation_av import ( - LTX2AVLatentPreparationStage, -) -from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import ( - LTX2DenoisingStage, -) -from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_diffusion import ( - RealtimeDiffusionStage, - RealtimeStageComponent, -) -from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_input_validation import ( - RealtimeInputValidationStage, -) -from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_text_encoding import ( - RealtimeTextEncodingStage, -) -from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_vae import ( - CausalVaeDecodingStage, - RealtimeImageVAEEncodingStage, -) -from sglang.multimodal_gen.runtime.pipelines_core.stages.text_connector import ( - LTX2TextConnectorStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.text_encoding import ( TextEncodingStage, @@ -89,54 +42,22 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.timestep_preparation im DMDTimestepPreparationStage, TimestepPreparationStage, ) -from sglang.multimodal_gen.runtime.pipelines_core.stages.upsampling import ( - LTX2HalveResolutionStage, - LTX2LoRASwitchStage, - LTX2UpsampleStage, -) __all__ = [ "PipelineStage", "InputValidationStage", - "RealtimeInputValidationStage", - "RealtimeDiffusionStage", - "RealtimeStageComponent", "TimestepPreparationStage", "DMDTimestepPreparationStage", "LatentPreparationStage", - "RealtimeChunkLatentPreparationStage", "ComfyUILatentPreparationStage", - "LTX2AVLatentPreparationStage", "DenoisingStage", "DmdDenoisingStage", - "LTX2DenoisingStage", - "LTX2AVDenoisingStage", "CausalDMDDenoisingStage", "EncodingStage", "ConditionEncodingStage", "AuxiliaryConditionEncodingStage", "DecodingStage", - "CausalVaeDecodingStage", - "LTX2AVDecodingStage", "ImageEncodingStage", "ImageVAEEncodingStage", - "RealtimeImageVAEEncodingStage", - "LTX2ImageEncodingStage", "TextEncodingStage", - "RealtimeTextEncodingStage", - "LTX2TextConnectorStage", - # Hunyuan3D shape stages - "Hunyuan3DShapeBeforeDenoisingStage", - "Hunyuan3DShapeDenoisingStage", - "Hunyuan3DShapeExportStage", - "Hunyuan3DShapeSaveStage", - # Hunyuan3D paint stages - "Hunyuan3DPaintPreprocessStage", - "Hunyuan3DPaintTexGenStage", - "Hunyuan3DPaintPostprocessStage", - # LTX-2 two-stage - "LTX2RefinementStage", - "LTX2HalveResolutionStage", - "LTX2LoRASwitchStage", - "LTX2UpsampleStage", ] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py index 3d8ef095e..f9a28a21d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py @@ -377,59 +377,3 @@ class LatentPreparationStage(PipelineStage): # result.add_check("latents", batch.latents, [V.is_tensor, V.with_dims(5)]) result.add_check("raw_latent_shape", batch.raw_latent_shape, V.is_tuple) return result - - -class RealtimeChunkLatentPreparationStage(LatentPreparationStage): - """Prepare one realtime causal DiT chunk from the encoded condition shape.""" - - def get_forward_latent_num_frames( - self, - batch: Req, - server_args: ServerArgs, - ) -> int: - return int( - batch.realtime_chunk_size - or self.transformer.config.arch_config.num_frames_per_block - ) - - def get_latent_preparation_spec( - self, - batch: Req, - server_args: ServerArgs, - batch_size: int, - num_frames: int, - device: torch.device | str, - ) -> LatentPreparationSpec: - condition_latent = batch.image_latent - assert condition_latent is not None, ( - "Realtime chunk latent preparation requires image_latent. " - "Ensure the condition VAE encoding stage runs before this stage." - ) - return LatentPreparationSpec( - shape=( - condition_latent.shape[0], - self.transformer.config.arch_config.out_channels, - num_frames, - condition_latent.shape[3], - condition_latent.shape[4], - ), - dtype=condition_latent.dtype, - device=device, - prepare_latent_ids=False, - pack_latents=False, - ) - - def should_scale_initial_noise(self, batch: Req, server_args: ServerArgs) -> bool: - return False - - def requires_batch_height_width(self, batch: Req, server_args: ServerArgs) -> bool: - return False - - def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: - result = VerificationResult() - result.add_check( - "image_latent", batch.image_latent, [V.is_tensor, V.with_dims(5)] - ) - result.add_check("generator", batch.generator, V.generator_or_list_generators) - result.add_check("latents", batch.latents, V.none_or_tensor) - return result diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/__init__.py new file mode 100644 index 000000000..79480ddaf --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/__init__.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""Hunyuan3D-specific pipeline stages""" + +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.hunyuan3d.paint import ( + Hunyuan3DPaintPostprocessStage, + Hunyuan3DPaintPreprocessStage, + Hunyuan3DPaintTexGenStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.hunyuan3d.shape import ( + Hunyuan3DShapeBeforeDenoisingStage, + Hunyuan3DShapeDenoisingStage, + Hunyuan3DShapeExportStage, + Hunyuan3DShapeSaveStage, +) + +__all__ = [ + "Hunyuan3DPaintPostprocessStage", + "Hunyuan3DPaintPreprocessStage", + "Hunyuan3DPaintTexGenStage", + "Hunyuan3DShapeBeforeDenoisingStage", + "Hunyuan3DShapeDenoisingStage", + "Hunyuan3DShapeExportStage", + "Hunyuan3DShapeSaveStage", +] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_paint.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_paint.py rename to python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_shape.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/shape.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_shape.py rename to python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/shape.py diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/__init__.py new file mode 100644 index 000000000..8853051a8 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/__init__.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""LTX-2-specific pipeline stages""" + +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.decoding_av import ( + LTX2AVDecodingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.denoising import ( + LTX2DenoisingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.denoising_av import ( + LTX2AVDenoisingStage, + LTX2RefinementStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.latent_preparation_av import ( + LTX2AVLatentPreparationStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.text_connector import ( + LTX2TextConnectorStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.upsampling import ( + LTX2HalveResolutionStage, + LTX2LoRASwitchStage, + LTX2UpsampleStage, +) + +__all__ = [ + "LTX2AVDecodingStage", + "LTX2AVDenoisingStage", + "LTX2AVLatentPreparationStage", + "LTX2DenoisingStage", + "LTX2HalveResolutionStage", + "LTX2LoRASwitchStage", + "LTX2RefinementStage", + "LTX2TextConnectorStage", + "LTX2UpsampleStage", +] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py rename to python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py rename to python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising_av.py similarity index 99% rename from python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py rename to python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising_av.py index 7b37a0d0d..dcae54d27 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising_av.py @@ -12,7 +12,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( StageParallelismType, ) -from sglang.multimodal_gen.runtime.pipelines_core.stages.ltx_2_denoising import ( +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.denoising import ( LTX2DenoisingStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/latent_preparation_av.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py rename to python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/latent_preparation_av.py diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/text_connector.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py rename to python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/text_connector.py diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/upsampling.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py rename to python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/upsampling.py diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_stage.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_stage.py index ea97ee27a..5f34d4bfd 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_stage.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_stage.py @@ -11,7 +11,7 @@ import torch from PIL import Image from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req -from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_diffusion import ( +from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.base import ( RealtimeDiffusionStage, ) from sglang.multimodal_gen.runtime.server_args import ServerArgs diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/__init__.py new file mode 100644 index 000000000..1e9e27ad9 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/__init__.py @@ -0,0 +1,31 @@ +# SPDX-License-Identifier: Apache-2.0 + +"""realtime pipeline stages shared by interactive diffusion models""" + +from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.base import ( + RealtimeDiffusionStage, + RealtimeStageComponent, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.input_validation import ( + RealtimeInputValidationStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.latent_preparation import ( + RealtimeChunkLatentPreparationStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.text_encoding import ( + RealtimeTextEncodingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.vae import ( + CausalVaeDecodingStage, + RealtimeImageVAEEncodingStage, +) + +__all__ = [ + "CausalVaeDecodingStage", + "RealtimeChunkLatentPreparationStage", + "RealtimeDiffusionStage", + "RealtimeImageVAEEncodingStage", + "RealtimeInputValidationStage", + "RealtimeStageComponent", + "RealtimeTextEncodingStage", +] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_diffusion.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/base.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_diffusion.py rename to python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/base.py diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_input_validation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/input_validation.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_input_validation.py rename to python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/input_validation.py diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/latent_preparation.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/latent_preparation.py new file mode 100644 index 000000000..c4d32f6eb --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/latent_preparation.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch + +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import ( + LatentPreparationSpec, + LatentPreparationStage, +) +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 + + +class RealtimeChunkLatentPreparationStage(LatentPreparationStage): + """Prepare one realtime causal DiT chunk from the encoded condition shape.""" + + def get_forward_latent_num_frames( + self, + batch: Req, + server_args: ServerArgs, + ) -> int: + return int( + batch.realtime_chunk_size + or self.transformer.config.arch_config.num_frames_per_block + ) + + def get_latent_preparation_spec( + self, + batch: Req, + server_args: ServerArgs, + batch_size: int, + num_frames: int, + device: torch.device | str, + ) -> LatentPreparationSpec: + condition_latent = batch.image_latent + assert condition_latent is not None, ( + "Realtime chunk latent preparation requires image_latent. " + "Ensure the condition VAE encoding stage runs before this stage." + ) + return LatentPreparationSpec( + shape=( + condition_latent.shape[0], + self.transformer.config.arch_config.out_channels, + num_frames, + condition_latent.shape[3], + condition_latent.shape[4], + ), + dtype=condition_latent.dtype, + device=device, + prepare_latent_ids=False, + pack_latents=False, + ) + + def should_scale_initial_noise(self, batch: Req, server_args: ServerArgs) -> bool: + return False + + def requires_batch_height_width(self, batch: Req, server_args: ServerArgs) -> bool: + return False + + def verify_input(self, batch: Req, server_args: ServerArgs) -> VerificationResult: + result = VerificationResult() + result.add_check( + "image_latent", batch.image_latent, [V.is_tensor, V.with_dims(5)] + ) + result.add_check("generator", batch.generator, V.generator_or_list_generators) + result.add_check("latents", batch.latents, V.none_or_tensor) + return result diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_text_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/text_encoding.py similarity index 64% rename from python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_text_encoding.py rename to python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/text_encoding.py index ddd1636c3..bab58c4fb 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_text_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/text_encoding.py @@ -49,6 +49,34 @@ def _copy_seq_lens( return [list(seq_lens) for seq_lens in value] +_TEXT_CACHE_TENSOR_LIST_FIELDS = ( + "prompt_embeds", + "pooled_embeds", + "prompt_attention_mask", + "prompt_embeds_mask", + "negative_prompt_embeds", + "neg_pooled_embeds", + "negative_attention_mask", + "negative_prompt_embeds_mask", +) +_TEXT_CACHE_SEQ_LENS_FIELDS = ( + "prompt_seq_lens", + "negative_prompt_seq_lens", +) +_TEXT_CACHE_FIELDS = _TEXT_CACHE_TENSOR_LIST_FIELDS + _TEXT_CACHE_SEQ_LENS_FIELDS +_TEXT_CACHE_DEFAULT_EMPTY_LIST_FIELDS = { + "prompt_embeds", + "pooled_embeds", + "neg_pooled_embeds", +} + + +def _copy_text_cache_field(name: str, value): + if name in _TEXT_CACHE_SEQ_LENS_FIELDS: + return _copy_seq_lens(value) + return _copy_tensor_list(value) + + class RealtimeTextState(BaseRealtimeState): def __init__(self): super().__init__() @@ -66,16 +94,8 @@ class RealtimeTextState(BaseRealtimeState): def clear_text_cache(self): self.cache_key = None - self.prompt_embeds = None - self.pooled_embeds = None - self.prompt_attention_mask = None - self.prompt_embeds_mask = None - self.prompt_seq_lens = None - self.negative_prompt_embeds = None - self.neg_pooled_embeds = None - self.negative_attention_mask = None - self.negative_prompt_embeds_mask = None - self.negative_prompt_seq_lens = None + for field in _TEXT_CACHE_FIELDS: + setattr(self, field, None) def dispose(self): super().dispose() @@ -97,33 +117,16 @@ class RealtimeTextEncodingStage(TextEncodingStage): ) def _restore_cached_outputs(self, batch: Req, state: RealtimeTextState) -> Req: - batch.prompt_embeds = _copy_tensor_list(state.prompt_embeds) or [] - batch.pooled_embeds = _copy_tensor_list(state.pooled_embeds) or [] - batch.prompt_attention_mask = _copy_tensor_list(state.prompt_attention_mask) - batch.prompt_embeds_mask = _copy_tensor_list(state.prompt_embeds_mask) - batch.prompt_seq_lens = _copy_seq_lens(state.prompt_seq_lens) - batch.negative_prompt_embeds = _copy_tensor_list(state.negative_prompt_embeds) - batch.neg_pooled_embeds = _copy_tensor_list(state.neg_pooled_embeds) or [] - batch.negative_attention_mask = _copy_tensor_list(state.negative_attention_mask) - batch.negative_prompt_embeds_mask = _copy_tensor_list( - state.negative_prompt_embeds_mask - ) - batch.negative_prompt_seq_lens = _copy_seq_lens(state.negative_prompt_seq_lens) + for field in _TEXT_CACHE_FIELDS: + value = _copy_text_cache_field(field, getattr(state, field)) + if value is None and field in _TEXT_CACHE_DEFAULT_EMPTY_LIST_FIELDS: + value = [] + setattr(batch, field, value) return batch def _store_outputs(self, batch: Req, state: RealtimeTextState) -> None: - state.prompt_embeds = _copy_tensor_list(batch.prompt_embeds) - state.pooled_embeds = _copy_tensor_list(batch.pooled_embeds) - state.prompt_attention_mask = _copy_tensor_list(batch.prompt_attention_mask) - state.prompt_embeds_mask = _copy_tensor_list(batch.prompt_embeds_mask) - state.prompt_seq_lens = _copy_seq_lens(batch.prompt_seq_lens) - state.negative_prompt_embeds = _copy_tensor_list(batch.negative_prompt_embeds) - state.neg_pooled_embeds = _copy_tensor_list(batch.neg_pooled_embeds) - state.negative_attention_mask = _copy_tensor_list(batch.negative_attention_mask) - state.negative_prompt_embeds_mask = _copy_tensor_list( - batch.negative_prompt_embeds_mask - ) - state.negative_prompt_seq_lens = _copy_seq_lens(batch.negative_prompt_seq_lens) + for field in _TEXT_CACHE_FIELDS: + setattr(state, field, _copy_text_cache_field(field, getattr(batch, field))) @torch.no_grad() def forward( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_vae.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/vae.py similarity index 100% rename from python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_vae.py rename to python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/vae.py diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py index 2ed9507a2..6cb8cc8e2 100644 --- a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py +++ b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py @@ -38,10 +38,10 @@ from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBa from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.lingbot_world import ( LingBotWorldCausalDMDDenoisingStage, ) -from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_diffusion import ( +from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.base import ( RealtimeDiffusionStage, ) -from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_input_validation import ( +from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.input_validation import ( RealtimeInputValidationStage, RealtimeInputValidationState, ) @@ -74,8 +74,14 @@ class _State(BaseRealtimeState): self.disposed = True +class _TestRealtimeDiffusionStage(RealtimeDiffusionStage): + def forward(self, batch, component_manager=None): + del batch, component_manager + raise NotImplementedError + + def test_realtime_diffusion_stage_declares_long_lived_components(): - stage = RealtimeDiffusionStage() + stage = _TestRealtimeDiffusionStage() server_args = SimpleNamespace( pipeline_config=SimpleNamespace(dit_precision="bf16", vae_precision="fp32") ) @@ -93,7 +99,7 @@ def test_realtime_diffusion_stage_declares_long_lived_components(): def test_realtime_diffusion_stage_requires_session(): - stage = RealtimeDiffusionStage(default_height=480, default_width=832) + stage = _TestRealtimeDiffusionStage(default_height=480, default_width=832) req = _Req(session=None) try: @@ -884,7 +890,7 @@ def test_realtime_chunk_latent_preparation_uses_chunk_spec(): import torch - from sglang.multimodal_gen.runtime.pipelines_core.stages import ( + from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime import ( RealtimeChunkLatentPreparationStage, ) diff --git a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_vae.py b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_vae.py index d33aaff25..0cf3f8306 100644 --- a/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_vae.py +++ b/python/sglang/multimodal_gen/test/unit/realtime/test_realtime_vae.py @@ -4,7 +4,7 @@ from types import SimpleNamespace import torch -from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime_vae import ( +from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime.vae import ( CausalVaeDecodingStage, RealtimeVAEDecodeState, ) @@ -22,7 +22,9 @@ def test_realtime_vae_decode_state_clears_model_cache_on_dispose(): def test_causal_vae_decoding_stage_keeps_wan_decoder_cache(monkeypatch): - from sglang.multimodal_gen.runtime.pipelines_core.stages import realtime_vae + from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime import ( + vae as realtime_vae, + ) class _WanVAE: def __init__(self): @@ -104,7 +106,9 @@ def test_causal_vae_decoding_stage_keeps_wan_decoder_cache(monkeypatch): def test_causal_vae_decoding_stage_prefers_native_causal_decode(monkeypatch): - from sglang.multimodal_gen.runtime.pipelines_core.stages import realtime_vae + from sglang.multimodal_gen.runtime.pipelines_core.stages.realtime import ( + vae as realtime_vae, + ) class _NativeCausalVAE: def __init__(self): diff --git a/python/sglang/multimodal_gen/test/unit/test_disagg_roles.py b/python/sglang/multimodal_gen/test/unit/test_disagg_roles.py index 6bc441937..05a0cf917 100644 --- a/python/sglang/multimodal_gen/test/unit/test_disagg_roles.py +++ b/python/sglang/multimodal_gen/test/unit/test_disagg_roles.py @@ -39,20 +39,20 @@ from sglang.multimodal_gen.runtime.pipelines.wan_i2v_pipeline import ( from sglang.multimodal_gen.runtime.pipelines_core.composed_pipeline_base import ( ComposedPipelineBase, ) -from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_av import ( - LTX2RefinementStage, -) -from sglang.multimodal_gen.runtime.pipelines_core.stages.hunyuan3d_shape import ( - Hunyuan3DShapeBeforeDenoisingStage, - Hunyuan3DShapeExportStage, - Hunyuan3DShapeSaveStage, -) from sglang.multimodal_gen.runtime.pipelines_core.stages.image_encoding import ( ImageVAEEncodingStage, ) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.helios_denoising import ( HeliosChunkedDenoisingStage, ) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.hunyuan3d.shape import ( + Hunyuan3DShapeBeforeDenoisingStage, + Hunyuan3DShapeExportStage, + Hunyuan3DShapeSaveStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.denoising_av import ( + LTX2RefinementStage, +) from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.mova import ( MOVADecodingStage, MOVADenoisingStage, @@ -628,7 +628,7 @@ class TestHunyuan3DShapeStageRuntimeDtype(_GlobalStageArgsMixin, unittest.TestCa ) with patch( - "sglang.multimodal_gen.runtime.pipelines_core.stages.hunyuan3d_shape.logger.warning" + "sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.hunyuan3d.shape.logger.warning" ) as mock_warning: dtype = stage._resolve_runtime_dtype(torch.zeros(1, dtype=torch.float16))