[diffusion] fix: unify LTX-2.3 HQ codepath gates for all LTX-2.3 variants (#23624)

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

@@ -157,11 +157,20 @@ class LTX2SigmaPreparationStage(PipelineStage):
     def forward(self, batch: Req, server_args: ServerArgs) -> Req:
         batch.extra["ltx2_phase"] = "stage1"
         if is_ltx23_native_variant(server_args.pipeline_config.vae_config.arch_config):
-            # Resolution-aware sigma shift is only required for the HQ pipeline
+            # Gate on pipeline class to mirror the three official entry points:
-            # (which targets 1080p+ resolutions and was aligned against official
+            # - HQ (`ti2vid_two_stages_hq.py:164`) calls
-            # LTX-2.3 HQ sigmas). Legacy one-stage and two-stage LTX-2.3 paths
+            #   `LTX2Scheduler.execute(latent=empty_latent, ...)` where
-            # were baselined against the constant-anchor schedule.
+            #   `empty_latent` is built from the **half-resolution** stage-1
+            #   shape → resolution-aware sigma shift.
+            # - Non-HQ two-stage (`ti2vid_two_stages.py:145`) and
+            #   one-stage (`ti2vid_one_stage.py:138`) call
+            #   `LTX2Scheduler.execute(steps=...)` with no `latent` →
+            #   falls back to `default_number_of_tokens = MAX_SHIFT_ANCHOR
+            #   = 4096` → constant-anchor sigma shift.
             if server_args.pipeline_class_name == "LTX2TwoStageHQPipeline":
+                # batch.height/width have already been halved by
+                # LTX2HalveResolutionStage, so these latents are the
+                # half-resolution stage-1 shape (matches `empty_latent`).
                 latent_num_frames = (int(batch.num_frames) - 1) // int(
                     server_args.pipeline_config.vae_temporal_compression
                 ) + 1

python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py

@@ -234,13 +234,10 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
         if self._should_reset_stage2_generators(server_args):
             self._reset_stage2_generators(batch)
         noise_scale = float(self.distilled_sigmas[0].item())
-        # HQ pipeline uses a dedicated, deterministic renoise generator seeded
+        is_ltx23 = is_ltx23_native_variant(
-        # from the request seed and advanced by stage-1 packed shapes to match
+            server_args.pipeline_config.vae_config.arch_config
-        # official LTX-2.3 HQ output. Legacy two-stage paths were baselined
+        )
-        # against `batch.generator`'s natural advance through stage-1, so keep
+        if is_ltx23:
-        # them on the original `_randn_like_with_batch_generators` sampling.
-        is_hq_pipeline = server_args.pipeline_class_name == "LTX2TwoStageHQPipeline"
-        if is_hq_pipeline:
             video_reference_for_gen = (
                 batch.latents if isinstance(batch.latents, torch.Tensor) else None
             )
@@ -259,7 +256,7 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
                 zero_clean_latent=True,
                 clean_latent_background=batch.ltx2_ti2v_clean_latent_background,
             )
-            if is_hq_pipeline:
+            if is_ltx23:
                 video_noise = self._ltx2_renoise_like(
                     prepared_latents, renoise_generator
                 )
@@ -271,7 +268,7 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
                 denoise_mask.to(device=prepared_latents.device, dtype=torch.float32)
                 * noise_scale
             )
-            if is_hq_pipeline:
+            if is_ltx23:
                 batch.latents = (
                     video_noise.float() * scaled_mask
                     + prepared_latents.float() * (1.0 - scaled_mask)
@@ -281,7 +278,7 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
                     video_noise * scaled_mask + prepared_latents * (1 - scaled_mask)
                 ).to(prepared_latents.dtype)
         else:
-            if is_hq_pipeline:
+            if is_ltx23:
                 video_noise = self._ltx2_renoise_like(batch.latents, renoise_generator)
                 batch.latents = (
                     video_noise.float() * noise_scale
@@ -296,7 +293,7 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
                 ).to(batch.latents.dtype)
 
         if isinstance(batch.audio_latents, torch.Tensor):
-            if is_hq_pipeline:
+            if is_ltx23:
                 audio_noise = self._ltx2_renoise_like(
                     batch.audio_latents, renoise_generator
                 )
@@ -316,9 +313,7 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
                     audio_noise * audio_scaled_mask
                     + batch.audio_latents * (1 - audio_scaled_mask)
                 ).to(batch.audio_latents.dtype)
-        if not is_ltx23_native_variant(
+        if not is_ltx23:
-            server_args.pipeline_config.vae_config.arch_config
-        ):
             batch.latents = batch.latents.to(
                 device=batch.latents.device, dtype=torch.float32
             )
@@ -334,14 +329,13 @@ class LTX2RefinementStage(LTX2AVDenoisingStage):
         self.scheduler = copy.deepcopy(original_scheduler)
         distilled_device = self.scheduler.sigmas.device
         num_steps = len(self.distilled_sigmas) - 1
-        # HQ pipeline extends the sigma schedule so the final step targets a
+        # Inject `0.0011` before the terminal `0.0` to avoid the
-        # small non-zero sigma (0.0011) instead of 0.0, matching official
+        # `sigma_next==0` singularity in res2s' `(sample - denoised) /
-        # LTX-2.3 HQ's last-step behavior. Legacy two-stage baselines used the
+        # (sigma - sigma_next)`. Official `res2s_denoising_loop` does this
-        # un-extended schedule (final step goes to 0.0).
+        # exact injection (samplers.py:262); official `euler_denoising_loop`
-        if (
+        # does NOT — it uses `sigma_next` directly. So gate on the active
-            server_args.pipeline_class_name == "LTX2TwoStageHQPipeline"
+        # sampler, not on the model variant.
-            and self.distilled_sigmas[-1].item() == 0.0
+        if self.sampler_name == "res2s" and self.distilled_sigmas[-1].item() == 0.0:
-        ):
             scheduler_sigmas = torch.cat(
                 [
                     self.distilled_sigmas[:-1],

python/sglang/multimodal_gen/test/server/perf_baselines.json

@@ -2632,44 +2632,45 @@
         },
         "ltx_2_3_hq_pipeline": {
             "stages_ms": {
-                "InputValidationStage": 0.11,
+                "InputValidationStage": 0.09,
-                "TextEncodingStage": 984.78,
+                "TextEncodingStage": 987.02,
-                "LTX2TextConnectorStage": 30.42,
+                "LTX2TextConnectorStage": 31.22,
-                "LTX2HalveResolutionStage": 0.1,
+                "LTX2HalveResolutionStage": 0.12,
-                "LTX2LoRASwitchStage": 0.01,
+                "LTX2LoRASwitchStage": 0.03,
-                "LTX2SigmaPreparationStage": 0.36,
+                "LTX2SigmaPreparationStage": 0.56,
-                "TimestepPreparationStage": 21.28,
+                "TimestepPreparationStage": 24.77,
-                "LTX2AVLatentPreparationStage": 0.13,
+                "LTX2AVLatentPreparationStage": 0.33,
-                "LTX2ImageEncodingStage": 0.03,
+                "LTX2ImageEncodingStage": 0.04,
-                "LTX2AVDenoisingStage": 20227.05,
+                "LTX2AVDenoisingStage": 25873.2,
-                "LTX2UpsampleStage": 157.73,
+                "LTX2UpsampleStage": 216.72,
-                "LTX2RefinementStage": 1676.07,
+                "LTX2RefinementStage": 2553.3,
-                "LTX2AVDecodingStage": 521.04,
+                "LTX2AVDecodingStage": 585.0,
                 "per_frame_generation": null
             },
             "denoise_step_ms": {
-                "0": 1406.0,
+                "0": 1512.13,
-                "1": 1362.75,
+                "1": 1506.38,
-                "2": 1306.18,
+                "2": 1523.98,
-                "3": 1299.65,
+                "3": 1533.71,
-                "4": 1282.16,
+                "4": 1569.51,
-                "5": 1290.32,
+                "5": 1546.87,
-                "6": 1284.64,
+                "6": 1545.47,
-                "7": 1265.01,
+                "7": 1528.27,
-                "8": 1304.06,
+                "8": 1550.64,
-                "9": 1246.21,
+                "9": 1548.24,
-                "10": 1102.19,
+                "10": 1531.65,
-                "11": 1379.4,
+                "11": 2211.84,
-                "12": 1467.28,
+                "12": 2306.24,
-                "13": 1469.49,
+                "13": 2270.41,
-                "14": 734.96,
+                "14": 1143.7,
-                "15": 547.19,
+                "15": 839.75,
-                "16": 543.27,
+                "16": 808.41,
-                "17": 539.0
+                "17": 806.56
             },
-            "expected_e2e_ms": 24150.97,
+            "expected_e2e_ms": 30848.37,
-            "expected_avg_denoise_ms": 1157.21,
+            "expected_avg_denoise_ms": 1515.76,
-            "expected_median_denoise_ms": 1287.48
+            "expected_median_denoise_ms": 1532.68,
+            "estimated_full_test_time_s": 363.2
         }
     }
 }