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

[diffusion] refactor: LTX2.3 code cleanup (#23207)

Browse Source
This commit is contained in:
Mick
2026-04-20 19:02:05 +08:00
committed by GitHub
parent da62e90904
commit 0be6ab04dd
6 changed files with 493 additions and 540 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py
+54 -28
View File
@@ -690,7 +690,7 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
upsampler_path = server_args.component_paths.get("spatial_upsampler")
if not upsampler_path:
raise ValueError(
"LTX2TwoStagePipeline requires --spatial-upsampler-path "
f"{self.pipeline_name} requires --spatial-upsampler-path "
"(component_paths['spatial_upsampler'])."
)
module, memory_usage = PipelineComponentLoader.load_component(
@@ -705,7 +705,7 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
distilled_lora_path = server_args.component_paths.get("distilled_lora")
if not distilled_lora_path:
raise ValueError(
"LTX2TwoStagePipeline requires --distilled-lora-path "
f"{self.pipeline_name} requires --distilled-lora-path "
"(component_paths['distilled_lora'])."
)
self._distilled_lora_path = distilled_lora_path
@@ -741,6 +741,10 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
"""Release inactive premerged DiTs according to the selected device mode."""
self._device_manager.release_premerged_transformers()
def release_ltx2_phase_state(self, phase: str | None) -> None:
if phase == "stage2":
self.release_premerged_transformers_to_cpu_snapshots()
def ensure_ltx2_phase_ready(self, phase: str | None) -> None:
self._device_manager.ensure_phase_ready(phase)
@@ -753,32 +757,28 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
"resident",
)
def switch_lora_phase(self, phase: str) -> None:
if phase == self._active_lora_phase:
return
def _can_short_circuit_lora_switch(self, phase: str) -> bool:
return (
phase in ("stage1", "stage2")
and self._use_premerged_stage2_transformer
and self._stage1_lora_path is None
)
if self._device_manager.switch_phase(phase):
self._active_lora_phase = phase
return
def _build_lora_switch_spec(
self, phase: str
) -> tuple[list[str], list[str], list[float], list[str]]:
lora_nicknames: list[str] = []
lora_paths: list[str] = []
lora_strengths: list[float] = []
lora_targets: list[str] = []
if phase == "stage1":
if self._stage1_lora_path:
self.set_lora(
lora_nickname="ltx2_stage1_base",
lora_path=self._stage1_lora_path,
target="transformer",
strength=self._stage1_lora_scale,
)
else:
# Stage 1 must run on the base transformer weights. If stage 2 left the
# distilled adapter active, stage 1 quality drifts away from the official
# two-stage pipeline immediately.
self.deactivate_lora_weights(target="transformer")
lora_nicknames.append("ltx2_stage1_base")
lora_paths.append(self._stage1_lora_path)
lora_strengths.append(self._stage1_lora_scale)
lora_targets.append("transformer")
elif phase == "stage2":
lora_nicknames = []
lora_paths = []
lora_strengths = []
lora_targets = []
if self._stage1_lora_path:
lora_nicknames.append("ltx2_stage1_base")
lora_paths.append(self._stage1_lora_path)
@@ -788,20 +788,46 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
lora_paths.append(self._distilled_lora_path)
lora_strengths.append(1.0)
lora_targets.append("transformer")
self.set_lora(
else:
raise ValueError(f"Unknown LTX2 two-stage LoRA phase: {phase}")
return lora_nicknames, lora_paths, lora_strengths, lora_targets
def switch_lora_phase(self, phase: str) -> None:
if phase == self._active_lora_phase:
return
if self._device_manager.switch_phase(
phase
) and self._can_short_circuit_lora_switch(phase):
self._active_lora_phase = phase
return
lora_nicknames, lora_paths, lora_strengths, lora_targets = (
self._build_lora_switch_spec(phase)
)
if lora_nicknames:
set_lora_kwargs = dict(
lora_nickname=lora_nicknames,
lora_path=lora_paths,
target=lora_targets,
strength=lora_strengths,
)
if phase == "stage2":
# Official LTX-2.3 two-stage builds stage 2 with distilled LoRA fused
# into the transformer weights. Legacy LTX-2 should keep the
# preexisting unmerged behavior to avoid regressing stage 2 quality.
merge_weights=self._should_merge_stage2_distilled_lora(
self.server_args
),
set_lora_kwargs["merge_weights"] = (
self._should_merge_stage2_distilled_lora(self.server_args)
)
self.set_lora(
**set_lora_kwargs,
)
else:
raise ValueError(f"Unknown LTX2 two-stage LoRA phase: {phase}")
# Stage 1 must run on the base transformer weights. If stage 2 left the
# distilled adapter active, stage 1 quality drifts away from the official
# two-stage pipeline immediately.
self.deactivate_lora_weights(target="transformer")
self._active_lora_phase = phase
python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py
+7 -19
View File
@@ -88,25 +88,13 @@ class LTX2AVDenoisingStage(LTX2DenoisingStage):
if hasattr(batch, "extra")
else ""
)
if (
pipeline is not None
and getattr(pipeline, "_use_premerged_stage2_transformer", False)
and server_args.dit_cpu_offload
and not server_args.use_fsdp_inference
and current_phase == "stage2"
):
release_to_snapshots = getattr(
pipeline, "release_premerged_transformers_to_cpu_snapshots", None
)
if callable(release_to_snapshots):
release_to_snapshots()
else:
for dit in filter(None, [self.transformer]):
param = next(dit.parameters(), None)
if param is not None and param.device.type == "cuda":
dit.to("cpu")
if torch.get_device_module().is_available():
torch.get_device_module().empty_cache()
release_phase_state = (
getattr(pipeline, "release_ltx2_phase_state", None)
if pipeline is not None
else None
)
if callable(release_phase_state):
release_phase_state(current_phase)
if isinstance(self.transformer, OffloadableDiTMixin):
for manager in self.transformer.layerwise_offload_managers:
python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py
+5 -2
View File
@@ -15,7 +15,10 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.validators import (
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.server_args import (
ServerArgs,
is_ltx2_two_stage_pipeline_name,
)
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
logger = init_logger(__name__)
@@ -62,7 +65,7 @@ class LTX2AVLatentPreparationStage(LatentPreparationStage):
server_args: ServerArgs,
):
if is_ltx23_native_variant(server_args.pipeline_config.vae_config.arch_config):
if server_args.pipeline_class_name == "LTX2TwoStagePipeline":
if is_ltx2_two_stage_pipeline_name(server_args.pipeline_class_name):
return server_args.pipeline_config.get_latent_dtype(
batch.prompt_embeds[0].dtype
)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py
+398 -468
View File
File diff suppressed because it is too large Load Diff
python/sglang/multimodal_gen/runtime/server_args.py
+12 -5
View File
@@ -73,6 +73,7 @@ logger = init_logger(__name__)
# GPUs on the faster no-offload default while preserving some headroom.
WAN_LAYERWISE_OFFLOAD_AUTO_DISABLE_MEM_GB = 130
LTX2_TWO_STAGE_DEVICE_MODES = ("original", "snapshot", "resident")
LTX2_TWO_STAGE_PIPELINE_NAMES = ("LTX2TwoStagePipeline",)
# H200-class GPUs (>=130 GiB total) can usually keep both LTX2 DiTs resident.
LTX2_RESIDENT_AUTO_ENABLE_MEM_GB = 130
@@ -84,6 +85,10 @@ def _normalize_ltx2_two_stage_device_mode(mode: str | None) -> str | None:
return mode
def is_ltx2_two_stage_pipeline_name(pipeline_class_name: str | None) -> bool:
return pipeline_class_name in LTX2_TWO_STAGE_PIPELINE_NAMES
class Backend(str, Enum):
"""
Enumeration for different model backends.
@@ -394,11 +399,7 @@ class ServerArgs(DisaggArgsMixin):
self.vae_cpu_offload = True
def _adjust_ltx2_two_stage_device_mode(self):
is_ltx23_two_stage = self.pipeline_class_name == "LTX2TwoStagePipeline" and (
self._is_ltx23_model_path(self.model_path)
or is_ltx23_native_variant(self.pipeline_config.vae_config.arch_config)
)
if not is_ltx23_two_stage:
if not self._is_ltx23_two_stage_pipeline():
return
mode = self.ltx2_two_stage_device_mode
@@ -449,6 +450,12 @@ class ServerArgs(DisaggArgsMixin):
)
return "snapshot"
def _is_ltx23_two_stage_pipeline(self) -> bool:
return is_ltx2_two_stage_pipeline_name(self.pipeline_class_name) and (
self._is_ltx23_model_path(self.model_path)
or is_ltx23_native_variant(self.pipeline_config.vae_config.arch_config)
)
def _adjust_attention_backend(self):
if self.attention_backend in ["fa3", "fa4"]:
self.attention_backend = "fa"
python/sglang/multimodal_gen/test/server/perf_baselines.json
+17 -18
View File
@@ -664,28 +664,27 @@
},
"zimage_image_t2i_2_gpus": {
"stages_ms": {
"TimestepPreparationStage": 35.17,
"DecodingStage": 9.73,
"TextEncodingStage": 307.62,
"LatentPreparationStage": 0.13,
"InputValidationStage": 0.05,
"DenoisingStage": 525.42
"TextEncodingStage": 309.68,
"LatentPreparationStage": 0.14,
"TimestepPreparationStage": 37.19,
"DenoisingStage": 525.39,
"DecodingStage": 10.18
},
"denoise_step_ms": {
"0": 19.97,
"1": 32.53,
"2": 62.6,
"3": 62.67,
"4": 63.03,
"5": 62.83,
"6": 62.73,
"7": 63.16,
"8": 62.63
"0": 40.09,
"1": 38.17,
"2": 54.71,
"3": 64.29,
"4": 65.0,
"5": 65.56,
"6": 64.61,
"7": 64.76,
"8": 64.35
},
"expected_e2e_ms": 957.78,
"expected_avg_denoise_ms": 57.99,
"expected_median_denoise_ms": 64.64,
"estimated_full_test_time_s": 121.0
"expected_e2e_ms": 961.86,
"expected_avg_denoise_ms": 57.95,
"expected_median_denoise_ms": 64.35
},
"qwen_image_edit_ti2i": {
"stages_ms": {