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

python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py

@@ -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:
+    def _can_short_circuit_lora_switch(self, phase: str) -> bool:
-        if phase == self._active_lora_phase:
+        return (
-            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):
+    def _build_lora_switch_spec(
-            self._active_lora_phase = phase
+        self, phase: str
-            return
+    ) -> 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_nicknames.append("ltx2_stage1_base")
-                    lora_nickname="ltx2_stage1_base",
+                lora_paths.append(self._stage1_lora_path)
-                    lora_path=self._stage1_lora_path,
+                lora_strengths.append(self._stage1_lora_scale)
-                    target="transformer",
+                lora_targets.append("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")
         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(
+                set_lora_kwargs["merge_weights"] = (
-                    self.server_args
+                    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

@@ -88,25 +88,13 @@ class LTX2AVDenoisingStage(LTX2DenoisingStage):
             if hasattr(batch, "extra")
             else ""
         )
-        if (
+        release_phase_state = (
-            pipeline is not None
+            getattr(pipeline, "release_ltx2_phase_state", None)
-            and getattr(pipeline, "_use_premerged_stage2_transformer", False)
+            if pipeline is not None
-            and server_args.dit_cpu_offload
+            else None
-            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):
+        if callable(release_phase_state):
-                release_to_snapshots()
+            release_phase_state(current_phase)
-            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()
 
         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

@@ -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/server_args.py

@@ -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 (
+        if not self._is_ltx23_two_stage_pipeline():
-            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:
             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

@@ -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,
+                "0": 40.09,
-                "1": 32.53,
+                "1": 38.17,
-                "2": 62.6,
+                "2": 54.71,
-                "3": 62.67,
+                "3": 64.29,
-                "4": 63.03,
+                "4": 65.0,
-                "5": 62.83,
+                "5": 65.56,
-                "6": 62.73,
+                "6": 64.61,
-                "7": 63.16,
+                "7": 64.76,
-                "8": 62.63
+                "8": 64.35
             },
-            "expected_e2e_ms": 957.78,
+            "expected_e2e_ms": 961.86,
-            "expected_avg_denoise_ms": 57.99,
+            "expected_avg_denoise_ms": 57.95,
-            "expected_median_denoise_ms": 64.64,
+            "expected_median_denoise_ms": 64.35
-            "estimated_full_test_time_s": 121.0
         },
         "qwen_image_edit_ti2i": {
             "stages_ms": {