Update optimization.py

optimization.py

@@ -15,7 +15,7 @@ from torchao.quantization import Int8WeightOnlyConfig
 from optimization_utils import capture_component_call
 from optimization_utils import aoti_compile
 from optimization_utils import ZeroGPUCompiledModel
-
+from optimization_utils import drain_module_parameters
 
 P = ParamSpec('P')
 
@@ -43,6 +43,26 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
     @spaces.GPU(duration=1500)
     def compile_transformer():
         
+        pipeline.load_lora_weights(
+            "Kijai/WanVideo_comfy", 
+            weight_name="Lightx2v/lightx2v_T2V_14B_cfg_step_distill_v2_lora_rank128_bf16.safetensors", 
+            adapter_name="lightning"
+        )
+        kwargs_lora = {}
+        kwargs_lora["load_into_transformer_2"] = True
+        pipeline.load_lora_weights(
+            "Kijai/WanVideo_comfy", 
+             weight_name="Lightx2v/lightx2v_T2V_14B_cfg_step_distill_v2_lora_rank128_bf16.safetensors", 
+            #weight_name="Wan22-Lightning/Wan2.2-Lightning_T2V-A14B-4steps-lora_LOW_fp16.safetensors", 
+            adapter_name="lightning_2", **kwargs_lora
+        )
+        pipeline.set_adapters(["lightning", "lightning_2"], adapter_weights=[1., 1.])
+
+
+        pipeline.fuse_lora(adapter_names=["lightning"], lora_scale=3., components=["transformer"])
+        pipeline.fuse_lora(adapter_names=["lightning_2"], lora_scale=1., components=["transformer_2"])
+        pipeline.unload_lora_weights()
+        
         with capture_component_call(pipeline, 'transformer') as call:
             pipeline(*args, **kwargs)
         
@@ -105,13 +125,8 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
         else:
             return cp2(*args, **kwargs)
 
-    transformer_config = pipeline.transformer.config
-    transformer_dtype = pipeline.transformer.dtype
-
-    pipeline.transformer = combined_transformer_1
-    pipeline.transformer.config = transformer_config # pyright: ignore[reportAttributeAccessIssue]
-    pipeline.transformer.dtype = transformer_dtype # pyright: ignore[reportAttributeAccessIssue]
+    pipeline.transformer.forward = combined_transformer_1
+    drain_module_parameters(pipeline.transformer)
 
-    pipeline.transformer_2 = combined_transformer_2
-    pipeline.transformer_2.config = transformer_config # pyright: ignore[reportAttributeAccessIssue]
-    pipeline.transformer_2.dtype = transformer_dtype # pyright: ignore[reportAttributeAccessIssue]
+    pipeline.transformer_2.forward = combined_transformer_2
+    drain_module_parameters(pipeline.transformer_2)