minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[diffusion] refactor: refactor realtime and model-specific stage modules (#27697)

Browse Source
This commit is contained in:
Mick
2026-06-10 09:39:06 +08:00
committed by GitHub
parent bcd9c5a903
commit af55025644
26 changed files with 251 additions and 200 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/pipelines/hunyuan3d_pipeline.py
+1 -1
View File
@@ -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,
python/sglang/multimodal_gen/runtime/pipelines/lingbot_world_causal_dmd_pipeline.py
+6 -4
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py
+7 -3
View File
@@ -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,
python/sglang/multimodal_gen/runtime/pipelines/sana_wm_realtime_pipeline.py
+4 -4
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/__init__.py
-79
View File
@@ -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",
]
python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation.py
-56
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/__init__.py
+25
View File
@@ -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",
]
python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_paint.py → python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/paint.py
View File
python/sglang/multimodal_gen/runtime/pipelines_core/stages/hunyuan3d_shape.py → python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/hunyuan3d/shape.py
View File
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/__init__.py
+37
View File
@@ -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",
]
python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py → python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/decoding_av.py
View File
python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py → python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising.py
View File
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py → python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/denoising_av.py
+1 -1
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py → python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/latent_preparation_av.py
View File
python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py → python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/text_connector.py
View File
python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py → python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/ltx_2/upsampling.py
View File
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/sana_wm/realtime_stage.py
+1 -1
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/__init__.py
+31
View File
@@ -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",
]
python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_diffusion.py → python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/base.py
View File
python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_input_validation.py → python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/input_validation.py
View File
python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/latent_preparation.py
+74
View File
@@ -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
python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_text_encoding.py → python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/text_encoding.py
+37 -34
View File
@@ -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(
python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime_vae.py → python/sglang/multimodal_gen/runtime/pipelines_core/stages/realtime/vae.py
View File
python/sglang/multimodal_gen/test/unit/realtime/test_realtime_runtime.py
+11 -5
View File
@@ -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,
)
python/sglang/multimodal_gen/test/unit/realtime/test_realtime_vae.py
+7 -3
View File
@@ -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):
python/sglang/multimodal_gen/test/unit/test_disagg_roles.py
+9 -9
View File
@@ -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))