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

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

Browse Source
This commit is contained in:
Mick
2026-04-24 17:44:08 +08:00
committed by GitHub
parent b060a5ccfd
commit b3b03369a5
3 changed files with 64 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py
+13 -4
View File
@@ -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
# (which targets 1080p+ resolutions and was aligned against official
# LTX-2.3 HQ sigmas). Legacy one-stage and two-stage LTX-2.3 paths
# were baselined against the constant-anchor schedule.
# Gate on pipeline class to mirror the three official entry points:
# - HQ (`ti2vid_two_stages_hq.py:164`) calls
# `LTX2Scheduler.execute(latent=empty_latent, ...)` where
# `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
+16 -22
View File
@@ -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
# from the request seed and advanced by stage-1 packed shapes to match
# official LTX-2.3 HQ output. Legacy two-stage paths were baselined
# against `batch.generator`'s natural advance through stage-1, so keep
# them on the original `_randn_like_with_batch_generators` sampling.
is_hq_pipeline = server_args.pipeline_class_name == "LTX2TwoStageHQPipeline"
if is_hq_pipeline:
is_ltx23 = is_ltx23_native_variant(
server_args.pipeline_config.vae_config.arch_config
)
if is_ltx23:
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(
server_args.pipeline_config.vae_config.arch_config
):
if not is_ltx23:
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
# small non-zero sigma (0.0011) instead of 0.0, matching official
# LTX-2.3 HQ's last-step behavior. Legacy two-stage baselines used the
# un-extended schedule (final step goes to 0.0).
if (
server_args.pipeline_class_name == "LTX2TwoStageHQPipeline"
and self.distilled_sigmas[-1].item() == 0.0
):
# Inject `0.0011` before the terminal `0.0` to avoid the
# `sigma_next==0` singularity in res2s' `(sample - denoised) /
# (sigma - sigma_next)`. Official `res2s_denoising_loop` does this
# exact injection (samplers.py:262); official `euler_denoising_loop`
# does NOT — it uses `sigma_next` directly. So gate on the active
# sampler, not on the model variant.
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
+35 -34
View File
@@ -2632,44 +2632,45 @@
},
"ltx_2_3_hq_pipeline": {
"stages_ms": {
"InputValidationStage": 0.11,
"TextEncodingStage": 984.78,
"LTX2TextConnectorStage": 30.42,
"LTX2HalveResolutionStage": 0.1,
"LTX2LoRASwitchStage": 0.01,
"LTX2SigmaPreparationStage": 0.36,
"TimestepPreparationStage": 21.28,
"LTX2AVLatentPreparationStage": 0.13,
"LTX2ImageEncodingStage": 0.03,
"LTX2AVDenoisingStage": 20227.05,
"LTX2UpsampleStage": 157.73,
"LTX2RefinementStage": 1676.07,
"LTX2AVDecodingStage": 521.04,
"InputValidationStage": 0.09,
"TextEncodingStage": 987.02,
"LTX2TextConnectorStage": 31.22,
"LTX2HalveResolutionStage": 0.12,
"LTX2LoRASwitchStage": 0.03,
"LTX2SigmaPreparationStage": 0.56,
"TimestepPreparationStage": 24.77,
"LTX2AVLatentPreparationStage": 0.33,
"LTX2ImageEncodingStage": 0.04,
"LTX2AVDenoisingStage": 25873.2,
"LTX2UpsampleStage": 216.72,
"LTX2RefinementStage": 2553.3,
"LTX2AVDecodingStage": 585.0,
"per_frame_generation": null
},
"denoise_step_ms": {
"0": 1406.0,
"1": 1362.75,
"2": 1306.18,
"3": 1299.65,
"4": 1282.16,
"5": 1290.32,
"6": 1284.64,
"7": 1265.01,
"8": 1304.06,
"9": 1246.21,
"10": 1102.19,
"11": 1379.4,
"12": 1467.28,
"13": 1469.49,
"14": 734.96,
"15": 547.19,
"16": 543.27,
"17": 539.0
"0": 1512.13,
"1": 1506.38,
"2": 1523.98,
"3": 1533.71,
"4": 1569.51,
"5": 1546.87,
"6": 1545.47,
"7": 1528.27,
"8": 1550.64,
"9": 1548.24,
"10": 1531.65,
"11": 2211.84,
"12": 2306.24,
"13": 2270.41,
"14": 1143.7,
"15": 839.75,
"16": 808.41,
"17": 806.56
},
"expected_e2e_ms": 24150.97,
"expected_avg_denoise_ms": 1157.21,
"expected_median_denoise_ms": 1287.48
"expected_e2e_ms": 30848.37,
"expected_avg_denoise_ms": 1515.76,
"expected_median_denoise_ms": 1532.68,
"estimated_full_test_time_s": 363.2
}
}
}