add tests and remove env variable

Signed-off-by: Yibin Li <109242046+yibinl-nvidia@users.noreply.github.com>

Author: yibinl-nvidia

tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py

@@ -4,7 +4,6 @@
 
 import json
 import math
-import os
 import time
 from concurrent.futures import ThreadPoolExecutor
 from typing import Any, Dict, List, Optional, Union
@@ -42,7 +41,6 @@
 STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0]
 _FP8_DTYPES = (torch.float8_e4m3fn, torch.float8_e5m2)
 _QKV_SUFFIXES = (".to_q", ".to_k", ".to_v")
-_DISABLE_OVERLAP_LORA_LOAD_ENV = "TRTLLM_LTX2_DISABLE_OVERLAP_LORA_LOAD"
 
 
 # ---------------------------------------------------------------------------
@@ -637,8 +635,7 @@ def load_standard_components(
 
         lora_executor = None
         lora_future = None
-        disable_overlap_lora_load = os.getenv(_DISABLE_OVERLAP_LORA_LOAD_ENV, "0") == "1"
-        if distilled_lora_path and not disable_overlap_lora_load:
+        if distilled_lora_path:
             logger.info(f"Starting distilled LoRA pre-compute from {distilled_lora_path}...")
             lora_executor = ThreadPoolExecutor(max_workers=1)
             lora_future = lora_executor.submit(
@@ -662,7 +659,6 @@ def load_standard_components(
                 spatial_upsampler_path,
                 distilled_lora_path,
                 lora_future,
-                disable_overlap_lora_load,
             )
         finally:
             if lora_executor is not None:
@@ -675,7 +671,6 @@ def _load_two_stage_components(
         spatial_upsampler_path: str,
         distilled_lora_path: str,
         lora_future,
-        disable_overlap_lora_load: bool,
     ) -> None:
         # --- Spatial upsampler ---
         if spatial_upsampler_path:
@@ -715,18 +710,10 @@ def _load_two_stage_components(
         # --- Distilled LoRA (pre-compute deltas) ---
         self._distilled_lora_deltas: Dict[str, torch.Tensor] = {}
         if distilled_lora_path:
-            if disable_overlap_lora_load:
-                logger.info(f"Loading distilled LoRA from {distilled_lora_path}...")
-                self._distilled_lora_deltas = _load_lora_deltas(
-                    distilled_lora_path,
-                    self.transformer,
-                    self._TRANSFORMER_PREFIX,
-                )
-            else:
-                logger.info("Waiting for distilled LoRA pre-compute...")
-                if lora_future is None:
-                    raise RuntimeError("Distilled LoRA pre-compute was not started.")
-                self._distilled_lora_deltas = lora_future.result()
+            logger.info("Waiting for distilled LoRA pre-compute...")
+            if lora_future is None:
+                raise RuntimeError("Distilled LoRA pre-compute was not started.")
+            self._distilled_lora_deltas = lora_future.result()
             logger.info(
                 f"Distilled LoRA ready: {len(self._distilled_lora_deltas)} parameter deltas"
             )

tests/unittest/_torch/visual_gen/test_ltx2_pipeline.py

@@ -757,6 +757,66 @@ def test_resolve_variant_requires_both_paths(self):
         assert result is LTX2Pipeline
 
 
+class TestTwoStageLoRALoadingOverlap:
+    """Test two-stage LoRA load orchestration without loading checkpoints."""
+
+    def test_load_standard_components_overlaps_lora_with_base_components(self, monkeypatch):
+        from types import SimpleNamespace
+
+        from tensorrt_llm._torch.visual_gen.models.ltx2 import pipeline_ltx2_two_stages
+        from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+        from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2_two_stages import (
+            LTX2TwoStagesPipeline,
+        )
+
+        events = []
+        lora_deltas = {"linear.weight": torch.ones(1)}
+
+        class FakeFuture:
+            def result(self):
+                events.append("lora_result")
+                return lora_deltas
+
+        class FakeExecutor:
+            def __init__(self, max_workers):
+                assert max_workers == 1
+
+            def submit(self, fn, *args):
+                events.append("lora_submit")
+                assert fn is pipeline_ltx2_two_stages._load_lora_deltas
+                assert args == (
+                    "/fake/lora.safetensors",
+                    pipeline.transformer,
+                    pipeline._TRANSFORMER_PREFIX,
+                )
+                return FakeFuture()
+
+            def shutdown(self, wait=True):
+                events.append(f"shutdown:{wait}")
+
+        def fake_base_load(self, checkpoint_dir, device, skip_components=None, **kwargs):
+            events.append("base_load")
+            assert checkpoint_dir == "/fake/checkpoint"
+            assert device == torch.device("cpu")
+
+        monkeypatch.setattr(pipeline_ltx2_two_stages, "ThreadPoolExecutor", FakeExecutor)
+        monkeypatch.setattr(LTX2Pipeline, "load_standard_components", fake_base_load)
+
+        pipeline = LTX2TwoStagesPipeline.__new__(LTX2TwoStagesPipeline)
+        pipeline.model_config = SimpleNamespace(
+            torch_dtype=torch.float32,
+            extra_attrs={
+                "distilled_lora_path": "/fake/lora.safetensors",
+            },
+        )
+        pipeline.transformer = object()
+
+        pipeline.load_standard_components("/fake/checkpoint", torch.device("cpu"))
+
+        assert events == ["lora_submit", "base_load", "lora_result", "shutdown:True"]
+        assert pipeline._distilled_lora_deltas is lora_deltas
+
+
 class TestTwoStageUpsamplerBuildingBlocks:
     """Test upsampler components without checkpoints (random weights)."""