[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
)
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):
super().initialize_pipeline(server_args)
server_args.component_paths = _resolve_ltx2_two_stage_component_paths(
@@ -967,16 +974,15 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline):
strength=lora_strengths,
)
if phase == "stage2":
# Official LTX-2.3 two-stage builds stage 2 with distilled LoRA fused
# into the transformer weights. Legacy LTX-2 should keep the
# preexisting unmerged behavior to avoid regressing stage 2 quality.
set_lora_kwargs["merge_weights"] = (
self._should_merge_stage2_distilled_lora(self.server_args)
# premerged modes keep official LTX-2.3 fused stage-2 LoRA; original
# avoids single-DiT request-time merge/unmerge with dynamic LoRA
set_lora_kwargs["merge_weights"] = self._should_merge_lora_for_phase(
phase
)
elif phase == "stage1" and self.pipeline_name == "LTX2TwoStageHQPipeline":
# Official HQ also builds stage 1 with distilled LoRA fused.
set_lora_kwargs["merge_weights"] = (
self._should_merge_stage2_distilled_lora(self.server_args)
set_lora_kwargs["merge_weights"] = self._should_merge_lora_for_phase(
phase
)
self.set_lora(
**set_lora_kwargs,
@@ -1362,7 +1362,7 @@
"1": 115.0,
"2": 130.0
},
"expected_e2e_ms": 2684.91,
"expected_e2e_ms": 3150.0,
"expected_avg_denoise_ms": 95.0,
"expected_median_denoise_ms": 120.0,
"estimated_full_test_time_s": 125.2
@@ -2477,7 +2477,7 @@
"32": 184.4
},
"expected_e2e_ms": 18039.38,
"expected_avg_denoise_ms": 183.75,
"expected_avg_denoise_ms": 271.95,
"expected_median_denoise_ms": 179.26,
"estimated_full_test_time_s": 160.0
},
@@ -2487,7 +2487,7 @@
"TextEncodingStage": 1554.35,
"LTX2TextConnectorStage": 27.64,
"LTX2HalveResolutionStage": 0.05,
"LTX2LoRASwitchStage": 114.37,
"LTX2LoRASwitchStage": 70.0,
"LTX2SigmaPreparationStage": 0.4,
"TimestepPreparationStage": 23.51,
"LTX2AVLatentPreparationStage": 0.15,
@@ -2531,9 +2531,9 @@
"29": 237.2,
"30": 127.08,
"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_median_denoise_ms": 246.85,
"estimated_full_test_time_s": 170.0