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

[diffusion] fix: improve LTX2.3 reference accuracy controls (#24022)

Browse Source
This commit is contained in:
Mick
2026-04-29 21:39:27 +08:00
committed by GitHub
parent 1279ae0787
commit e5c2a9b6cf
8 changed files with 98 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/sample/ltx_2.py
+3
View File
@@ -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
python/sglang/multimodal_gen/configs/utils.py
+1
View File
@@ -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() == "":
python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
+1
View File
@@ -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",
python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py
+15
View File
@@ -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:
python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py
+41 -29
View File
@@ -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
python/sglang/multimodal_gen/test/server/consistency_threshold.json
+19 -7
View File
@@ -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,
python/sglang/multimodal_gen/test/server/perf_baselines.json
+1 -1
View File
@@ -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
python/sglang/multimodal_gen/test/unit/test_server_args.py
+17
View File
@@ -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()