Update app.py

app.py

@@ -119,14 +119,39 @@ GEMMA_REPO = "Lightricks/gemma-3-12b-it-qat-q4_0-unquantized"
 # =============================================================================
 
 class HQPipelineWithCachedLoRA(TI2VidTwoStagesHQPipeline):
-    """TI2VidTwoStagesHQPipeline with support for cached LoRA state."""
+    """
+    TI2VidTwoStagesHQPipeline modified to:
+    1. NOT accept or pass distilled_lora to parent init (enables preloading)
+    2. Handle ALL LoRAs via cached state (distilled + 12 custom)
+    """
     
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(
+        self,
+        checkpoint_path: str,
+        spatial_upsampler_path: str,
+        gemma_root: str,
+        loras: tuple = (),
+        quantization: QuantizationPolicy | None = None,
+    ):
+        # Call parent WITHOUT distilled_lora parameters
+        # We create minimal stage ledgers (no LoRAs) for preloading
+        super().__init__(
+            checkpoint_path=checkpoint_path,
+            distilled_lora=[],  # Empty - satisfies signature
+            distilled_lora_strength_stage_1=0.0,
+            distilled_lora_strength_stage_2=0.0,
+            spatial_upsampler_path=spatial_upsampler_path,
+            gemma_root=gemma_root,
+            loras=loras,
+            quantization=quantization,
+        )
+        
+        # Storage for cached LoRA states for each stage
         self._cached_state_stage1 = None
         self._cached_state_stage2 = None
     
     def apply_cached_lora_state(self, state_dict_stage1, state_dict_stage2=None):
+        """Apply pre-cached LoRA state to both stage transformers."""
         self._cached_state_stage1 = state_dict_stage1
         self._cached_state_stage2 = state_dict_stage2 if state_dict_stage2 else state_dict_stage1
     
@@ -147,6 +172,7 @@ class HQPipelineWithCachedLoRA(TI2VidTwoStagesHQPipeline):
         tiling_config: TilingConfig | None = None,
         enhance_prompt: bool = False,
     ):
+        # ... same as before ...
         assert_resolution(height=height, width=width, is_two_stage=True)
         
         device = self.device
@@ -154,7 +180,6 @@ class HQPipelineWithCachedLoRA(TI2VidTwoStagesHQPipeline):
         generator = torch.Generator(device=device).manual_seed(seed)
         noiser = GaussianNoiser(generator=generator)
         
-        # Apply cached LoRA state if available
         if self._cached_state_stage1 is not None:
             print("[LoRA] Applying cached state to stage 1 transformer...")
             t1 = self.stage_1_model_ledger.transformer()
@@ -354,12 +379,9 @@ print("Initializing HQ Pipeline...")
 
 pipeline = HQPipelineWithCachedLoRA(
     checkpoint_path=checkpoint_path,
-    distilled_lora=[],  # No distilled LoRA at init - applied via cached state
-    distilled_lora_strength_stage_1=0.0,  # Not used since loras is empty
-    distilled_lora_strength_stage_2=0.0,  # Not used since loras is empty
     spatial_upsampler_path=spatial_upsampler_path,
     gemma_root=gemma_root,
-    loras=(),  # No LoRAs at init - preloading will work
+    loras=[],  # No LoRAs at init - preloading works
     quantization=QuantizationPolicy.fp8_cast(),
 )