From e5c2a9b6cf35e163c606ce2349b4c132d3c2e58c Mon Sep 17 00:00:00 2001 From: Mick Date: Wed, 29 Apr 2026 21:39:27 +0800 Subject: [PATCH] [diffusion] fix: improve LTX2.3 reference accuracy controls (#24022) --- .../multimodal_gen/configs/sample/ltx_2.py | 3 + python/sglang/multimodal_gen/configs/utils.py | 1 + .../runtime/models/dits/ltx_2.py | 1 + .../pipelines_core/stages/ltx_2_denoising.py | 15 ++++ .../pipelines_core/stages/text_connector.py | 70 +++++++++++-------- .../test/server/consistency_threshold.json | 26 +++++-- .../test/server/perf_baselines.json | 2 +- .../test/unit/test_server_args.py | 17 +++++ 8 files changed, 98 insertions(+), 37 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/sample/ltx_2.py b/python/sglang/multimodal_gen/configs/sample/ltx_2.py index 86837abb7..e074ace6e 100644 --- a/python/sglang/multimodal_gen/configs/sample/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/sample/ltx_2.py @@ -64,6 +64,7 @@ class LTX23SamplingParams(LTX2SamplingParams): audio_modality_scale: float = 3.0 audio_skip_step: int = 0 audio_stg_blocks: list[int] = field(default_factory=lambda: [28]) + skip_v2a_cross_attn_for_video_gt: bool = False def build_request_extra(self) -> dict[str, Any]: extra = super().build_request_extra() @@ -81,6 +82,8 @@ class LTX23SamplingParams(LTX2SamplingParams): "audio_skip_step": self.audio_skip_step, "audio_stg_blocks": self.audio_stg_blocks, } + if self.skip_v2a_cross_attn_for_video_gt: + extra["ltx2_skip_v2a_cross_attn_for_video_gt"] = True return extra diff --git a/python/sglang/multimodal_gen/configs/utils.py b/python/sglang/multimodal_gen/configs/utils.py index d2cc69adb..11565db01 100644 --- a/python/sglang/multimodal_gen/configs/utils.py +++ b/python/sglang/multimodal_gen/configs/utils.py @@ -19,6 +19,7 @@ def update_config_from_args( # Handle top-level attributes (no prefix) args_not_to_remove = [ "model_path", + "disable_autocast", ] args_to_remove = [] if prefix.strip() == "": diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index 10af35ea9..31918203e 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -1450,6 +1450,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): base_num_frames=cross_attn_pos_embed_max_pos, sampling_rate=16000, hop_length=160, + scale_factors=self.audio_scale_factors, theta=float(arch.positional_embedding_theta), causal_offset=causal_offset, modality="audio", diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py index f95a1aac3..91c0b2274 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py @@ -1446,6 +1446,9 @@ class LTX2DenoisingStage(DenoisingStage): and int(getattr(batch, "ltx2_num_image_tokens", 0)) > 0 ) ) + skip_v2a_cross_attn_for_video_gt = bool( + batch.extra.get("ltx2_skip_v2a_cross_attn_for_video_gt", False) + ) def evaluate_stage1_guided_x0( *, @@ -1476,6 +1479,9 @@ class LTX2DenoisingStage(DenoisingStage): encoder_hidden_states=encoder_hidden_states, audio_encoder_hidden_states=audio_encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, + disable_v2a_cross_attn=( + skip_v2a_cross_attn_for_video_gt + ), ) ) v_neg, a_v_neg = step.current_model( @@ -1485,6 +1491,9 @@ class LTX2DenoisingStage(DenoisingStage): encoder_hidden_states=negative_encoder_hidden_states, audio_encoder_hidden_states=negative_audio_encoder_hidden_states, encoder_attention_mask=negative_encoder_attention_mask, + disable_v2a_cross_attn=( + skip_v2a_cross_attn_for_video_gt + ), ) ) @@ -1510,6 +1519,9 @@ class LTX2DenoisingStage(DenoisingStage): skip_audio_self_attn_blocks=tuple( stage1_guider_params["audio_stg_blocks"] ), + disable_v2a_cross_attn=( + skip_v2a_cross_attn_for_video_gt + ), ) ) v_ptb = v_ptb.float() @@ -1539,12 +1551,14 @@ class LTX2DenoisingStage(DenoisingStage): encoder_hidden_states=encoder_hidden_states, audio_encoder_hidden_states=audio_encoder_hidden_states, encoder_attention_mask=encoder_attention_mask, + disable_v2a_cross_attn=skip_v2a_cross_attn_for_video_gt, ), LTX2GuidancePassSpec( name="neg", encoder_hidden_states=negative_encoder_hidden_states, audio_encoder_hidden_states=negative_audio_encoder_hidden_states, encoder_attention_mask=negative_encoder_attention_mask, + disable_v2a_cross_attn=skip_v2a_cross_attn_for_video_gt, ), ] if need_perturbed: @@ -1560,6 +1574,7 @@ class LTX2DenoisingStage(DenoisingStage): skip_audio_self_attn_blocks=tuple( stage1_guider_params["audio_stg_blocks"] ), + disable_v2a_cross_attn=skip_v2a_cross_attn_for_video_gt, ) ) if need_modality: diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py index c2969e99c..d82796aee 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py @@ -45,47 +45,59 @@ class LTX2TextConnectorStage(PipelineStage): else None ) - # Handle CFG: Concatenate negative and positive inputs - if batch.do_classifier_free_guidance: - # Concatenate: [Negative, Positive] - prompt_embeds = torch.cat([neg_prompt_embeds, prompt_embeds], dim=0) - prompt_attention_mask = torch.cat( - [neg_prompt_attention_mask, prompt_attention_mask], dim=0 + if prompt_embeds is None or prompt_attention_mask is None: + raise ValueError( + "LTX2TextConnectorStage requires prompt embeddings and " + "attention mask." ) - # Prepare additive mask for connectors (as per Diffusers implementation) - dtype = prompt_embeds.dtype - - additive_attention_mask = (prompt_attention_mask.to(torch.int64) - 1).to( - dtype - ) * torch.finfo(dtype).max - - # Call connectors - # Expects: prompt_embeds, attention_mask, additive_mask=True - with set_forward_context(current_timestep=None, attn_metadata=None): - connector_prompt_embeds, connector_audio_prompt_embeds, connector_mask = ( - self.connectors( - prompt_embeds, additive_attention_mask, additive_mask=True + if batch.do_classifier_free_guidance: + if neg_prompt_embeds is None or neg_prompt_attention_mask is None: + raise ValueError( + "LTX2TextConnectorStage requires negative prompt embeddings " + "and attention mask when classifier-free guidance is enabled." ) - ) - # Split results if CFG was enabled - if batch.do_classifier_free_guidance: - neg_embeds, pos_embeds = connector_prompt_embeds.chunk(2, dim=0) - neg_audio_embeds, pos_audio_embeds = connector_audio_prompt_embeds.chunk( - 2, dim=0 - ) - neg_mask, pos_mask = connector_mask.chunk(2, dim=0) + # Official LTX-2.3 processes positive and negative prompts through + # the connector independently; batching shifts output numerics. + dtype = prompt_embeds.dtype + pos_additive_mask = (prompt_attention_mask.to(torch.int64) - 1).to( + dtype + ) * torch.finfo(dtype).max + neg_additive_mask = (neg_prompt_attention_mask.to(torch.int64) - 1).to( + dtype + ) * torch.finfo(dtype).max + + with set_forward_context(current_timestep=None, attn_metadata=None): + pos_embeds, pos_audio_embeds, pos_mask = self.connectors( + prompt_embeds, pos_additive_mask, additive_mask=True + ) + neg_embeds, neg_audio_embeds, neg_mask = self.connectors( + neg_prompt_embeds, neg_additive_mask, additive_mask=True + ) batch.prompt_embeds = [pos_embeds] batch.audio_prompt_embeds = [pos_audio_embeds] batch.prompt_attention_mask = pos_mask - batch.negative_prompt_embeds = [neg_embeds] batch.negative_audio_prompt_embeds = [neg_audio_embeds] batch.negative_attention_mask = neg_mask else: - # Update positive fields + # Prepare additive mask for connectors (as per diffusers implementation) + dtype = prompt_embeds.dtype + additive_attention_mask = (prompt_attention_mask.to(torch.int64) - 1).to( + dtype + ) * torch.finfo(dtype).max + + with set_forward_context(current_timestep=None, attn_metadata=None): + ( + connector_prompt_embeds, + connector_audio_prompt_embeds, + connector_mask, + ) = self.connectors( + prompt_embeds, additive_attention_mask, additive_mask=True + ) + batch.prompt_embeds = [connector_prompt_embeds] batch.audio_prompt_embeds = [connector_audio_prompt_embeds] batch.prompt_attention_mask = connector_mask diff --git a/python/sglang/multimodal_gen/test/server/consistency_threshold.json b/python/sglang/multimodal_gen/test/server/consistency_threshold.json index d41e7aa66..a0adacefa 100644 --- a/python/sglang/multimodal_gen/test/server/consistency_threshold.json +++ b/python/sglang/multimodal_gen/test/server/consistency_threshold.json @@ -104,16 +104,16 @@ "mean_abs_diff_threshold": 8.0 }, "ltx_2.3_one_stage_ti2v": { - "clip_threshold": 0.57, + "clip_threshold": 0.64, "ssim_threshold": 0.42, - "psnr_threshold": 9.0, - "mean_abs_diff_threshold": 57.0 + "psnr_threshold": 8.8, + "mean_abs_diff_threshold": 59.0 }, "ltx_2.3_two_stage_t2v_2gpus": { - "clip_threshold": 0.78, - "ssim_threshold": 0.20, - "psnr_threshold": 12.7, - "mean_abs_diff_threshold": 49.0 + "clip_threshold": 0.80, + "ssim_threshold": 0.12, + "psnr_threshold": 12.2, + "mean_abs_diff_threshold": 50.5 }, "wan2_1_t2v_1.3b_teacache_enabled": { "clip_threshold": 0.93, @@ -246,6 +246,18 @@ "ssim_threshold": 0.89, "psnr_threshold": 24.0, "mean_abs_diff_threshold": 10.0 + }, + "ltx_2_3_hq_pipeline": { + "clip_threshold": 0.78, + "ssim_threshold": 0.48, + "psnr_threshold": 12.0, + "mean_abs_diff_threshold": 48.0 + }, + "ltx_2_3_two_stage_ti2v_2gpus": { + "clip_threshold": 0.92, + "ssim_threshold": 0.58, + "psnr_threshold": 17.5, + "mean_abs_diff_threshold": 20.0 } }, "default_clip_threshold_image": 0.92, diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index b3f7dc5ef..9772146ae 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -2576,7 +2576,7 @@ "LTX2AVLatentPreparationStage": 0.33, "LTX2ImageEncodingStage": 0.04, "LTX2AVDenoisingStage": 25873.2, - "LTX2UpsampleStage": 216.72, + "LTX2UpsampleStage": 752.07, "LTX2RefinementStage": 2553.3, "LTX2AVDecodingStage": 585.0, "per_frame_generation": null diff --git a/python/sglang/multimodal_gen/test/unit/test_server_args.py b/python/sglang/multimodal_gen/test/unit/test_server_args.py index 3151287f6..17b6c84ee 100644 --- a/python/sglang/multimodal_gen/test/unit/test_server_args.py +++ b/python/sglang/multimodal_gen/test/unit/test_server_args.py @@ -210,6 +210,23 @@ class TestPipelineResolutionCliOverride(unittest.TestCase): self.assertEqual(server_args.pipeline_config.resolution, 768) + def test_disable_autocast_is_preserved_after_pipeline_config_resolution(self): + parser = FlexibleArgumentParser() + ServerArgs.add_cli_args(parser) + argv = [ + "--model-path", + "Qwen/Qwen-Image-Layered", + "--disable-autocast", + "true", + ] + + with patch.object(sys, "argv", ["sglang"] + argv): + args, unknown_args = parser.parse_known_args(argv) + server_args = ServerArgs.from_cli_args(args, unknown_args) + + self.assertTrue(server_args.pipeline_config.disable_autocast) + self.assertTrue(server_args.disable_autocast) + if __name__ == "__main__": unittest.main()