[diffusion] fix: use dynamic LoRA for LTX2 original stage-two (#25591)

This commit is contained in:
Mick
2026-05-18 23:05:29 +08:00
committed by GitHub
parent d1acd62d29
commit 110bbdcad7
2 changed files with 18 additions and 12 deletions
@@ -806,6 +806,13 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
server_args.pipeline_config.vae_config.arch_config server_args.pipeline_config.vae_config.arch_config
) )
def _should_merge_lora_for_phase(self, phase: str) -> bool:
if phase == "stage2" and self._ltx2_residency.mode == "original":
# original mode reuses one DiT for both phases; dynamic LoRA avoids
# request-time merge/unmerge without keeping another DiT resident
return False
return self._should_merge_stage2_distilled_lora(self.server_args)
def initialize_pipeline(self, server_args: ServerArgs): def initialize_pipeline(self, server_args: ServerArgs):
super().initialize_pipeline(server_args) super().initialize_pipeline(server_args)
server_args.component_paths = _resolve_ltx2_two_stage_component_paths( server_args.component_paths = _resolve_ltx2_two_stage_component_paths(
@@ -967,16 +974,15 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
strength=lora_strengths, strength=lora_strengths,
) )
if phase == "stage2": if phase == "stage2":
# Official LTX-2.3 two-stage builds stage 2 with distilled LoRA fused # premerged modes keep official LTX-2.3 fused stage-2 LoRA; original
# into the transformer weights. Legacy LTX-2 should keep the # avoids single-DiT request-time merge/unmerge with dynamic LoRA
# preexisting unmerged behavior to avoid regressing stage 2 quality. set_lora_kwargs["merge_weights"] = self._should_merge_lora_for_phase(
set_lora_kwargs["merge_weights"] = ( phase
self._should_merge_stage2_distilled_lora(self.server_args)
) )
elif phase == "stage1" and self.pipeline_name == "LTX2TwoStageHQPipeline": elif phase == "stage1" and self.pipeline_name == "LTX2TwoStageHQPipeline":
# Official HQ also builds stage 1 with distilled LoRA fused. # Official HQ also builds stage 1 with distilled LoRA fused.
set_lora_kwargs["merge_weights"] = ( set_lora_kwargs["merge_weights"] = self._should_merge_lora_for_phase(
self._should_merge_stage2_distilled_lora(self.server_args) phase
) )
self.set_lora( self.set_lora(
**set_lora_kwargs, **set_lora_kwargs,
@@ -1362,7 +1362,7 @@
"1": 115.0, "1": 115.0,
"2": 130.0 "2": 130.0
}, },
"expected_e2e_ms": 2684.91, "expected_e2e_ms": 3150.0,
"expected_avg_denoise_ms": 95.0, "expected_avg_denoise_ms": 95.0,
"expected_median_denoise_ms": 120.0, "expected_median_denoise_ms": 120.0,
"estimated_full_test_time_s": 125.2 "estimated_full_test_time_s": 125.2
@@ -2477,7 +2477,7 @@
"32": 184.4 "32": 184.4
}, },
"expected_e2e_ms": 18039.38, "expected_e2e_ms": 18039.38,
"expected_avg_denoise_ms": 183.75, "expected_avg_denoise_ms": 271.95,
"expected_median_denoise_ms": 179.26, "expected_median_denoise_ms": 179.26,
"estimated_full_test_time_s": 160.0 "estimated_full_test_time_s": 160.0
}, },
@@ -2487,7 +2487,7 @@
"TextEncodingStage": 1554.35, "TextEncodingStage": 1554.35,
"LTX2TextConnectorStage": 27.64, "LTX2TextConnectorStage": 27.64,
"LTX2HalveResolutionStage": 0.05, "LTX2HalveResolutionStage": 0.05,
"LTX2LoRASwitchStage": 114.37, "LTX2LoRASwitchStage": 70.0,
"LTX2SigmaPreparationStage": 0.4, "LTX2SigmaPreparationStage": 0.4,
"TimestepPreparationStage": 23.51, "TimestepPreparationStage": 23.51,
"LTX2AVLatentPreparationStage": 0.15, "LTX2AVLatentPreparationStage": 0.15,
@@ -2531,9 +2531,9 @@
"29": 237.2, "29": 237.2,
"30": 127.08, "30": 127.08,
"31": 123.1, "31": 123.1,
"32": 128.06 "32": 209.80
}, },
"expected_e2e_ms": 17660.77, "expected_e2e_ms": 12000.0,
"expected_avg_denoise_ms": 238.8, "expected_avg_denoise_ms": 238.8,
"expected_median_denoise_ms": 246.85, "expected_median_denoise_ms": 246.85,
"estimated_full_test_time_s": 170.0 "estimated_full_test_time_s": 170.0