Update optimization.py

optimization.py

@@ -36,10 +36,12 @@ INDUCTOR_CONFIGS = {
     'triton.cudagraphs': True,
 }
 
+from torchao.quantization import DEFAULT_INT4_AUTOQUANT_CLASS_LIST
+
 def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
     print("[optimize_pipeline_] Starting pipeline optimization")
 
-    # Text encoder: move to CPU first, then quantize+compile to avoid early CUDA init
+    # --- TEXT ENCODER ---
     pipeline.text_encoder = pipeline.text_encoder.cpu()
     pipeline.text_encoder = torchao.autoquant(
         torch.compile(pipeline.text_encoder, mode='max-autotune'),
@@ -49,6 +51,7 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
 
     @spaces.GPU(duration=1500)
     def compile_transformer():
+        # --- LOAD LORAS ---
         print("[compile_transformer] Loading LoRA weights")
         pipeline.load_lora_weights(
             "DeepBeepMeep/Wan2.2",
@@ -63,23 +66,30 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
         )
         pipeline.set_adapters(["lightning", "lightning_2"], adapter_weights=[1.0, 1.0])
 
+        # --- FUSE & UNLOAD ---
         print("[compile_transformer] Fusing LoRA weights")
         pipeline.fuse_lora(adapter_names=["lightning"], lora_scale=3.0, components=["transformer"])
         pipeline.fuse_lora(adapter_names=["lightning_2"], lora_scale=1.0, components=["transformer_2"])
         pipeline.unload_lora_weights()
 
-        print("[compile_transformer] Running dummy forward pass to capture shapes")
+        # --- DUMMY FORWARD ---
+        print("[compile_transformer] Capturing shapes")
         with torch.inference_mode():
             with capture_component_call(pipeline, 'transformer') as call:
                 pipeline(*args, **kwargs)
 
-        dynamic_shapes = tree_map_only((torch.Tensor, bool), lambda t: None, call.kwargs)
-        dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
+        hidden_states: torch.Tensor = call.kwargs['hidden_states']
+        hidden_states_transposed = hidden_states.transpose(-1, -2).contiguous()
 
-        # Now that we're inside GPU context, move and compile transformers
+        # --- MOVE TO CUDA BEFORE AUTOQUANT ---
         pipeline.transformer = pipeline.transformer.to("cuda")
         pipeline.transformer_2 = pipeline.transformer_2.to("cuda")
 
+        # Sanity: Ensure parameters exist before quantization
+        assert any(p.numel() > 0 for p in pipeline.transformer.parameters()), "Transformer has no params!"
+        assert any(p.numel() > 0 for p in pipeline.transformer_2.parameters()), "Transformer_2 has no params!"
+
+        # --- AUTOQUANT + COMPILE ---
         compiled_transformer = torchao.autoquant(
             torch.compile(pipeline.transformer, mode='max-autotune'),
             qtensor_class_list=DEFAULT_INT4_AUTOQUANT_CLASS_LIST
@@ -89,28 +99,26 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
             qtensor_class_list=DEFAULT_INT4_AUTOQUANT_CLASS_LIST
         )
 
-        # Patch forward with landscape/portrait logic
-        hidden_states: torch.Tensor = call.kwargs['hidden_states']
-        hidden_states_transposed = hidden_states.transpose(-1, -2).contiguous()
-
+        # --- PATCH FOR LANDSCAPE/PORTRAIT ---
         def combined_transformer_1(*a, **k):
             if k['hidden_states'].shape[-1] > k['hidden_states'].shape[-2]:
                 return compiled_transformer(*a, **k)
-            k_mod = k.copy()
+            k_mod = dict(k)
             k_mod['hidden_states'] = hidden_states_transposed
             return compiled_transformer(*a, **k_mod)
 
         def combined_transformer_2(*a, **k):
             if k['hidden_states'].shape[-1] > k['hidden_states'].shape[-2]:
                 return compiled_transformer_2(*a, **k)
-            k_mod = k.copy()
+            k_mod = dict(k)
             k_mod['hidden_states'] = hidden_states_transposed
             return compiled_transformer_2(*a, **k_mod)
 
         pipeline.transformer.forward = combined_transformer_1
-        drain_module_parameters(pipeline.transformer)
-
         pipeline.transformer_2.forward = combined_transformer_2
+
+        # --- NOW drain parameters to save VRAM ---
+        drain_module_parameters(pipeline.transformer)
         drain_module_parameters(pipeline.transformer_2)
 
         print("[compile_transformer] Transformers autoquantized, compiled, and patched")
@@ -118,3 +126,4 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
 
     cl1, cl2 = compile_transformer()
     print("[optimize_pipeline_] Optimization complete")
+