Update optimization.py

optimization.py

@@ -20,7 +20,13 @@ from optimization_utils import ZeroGPUCompiledModel
 P = ParamSpec('P')
 
 
-TRANSFORMER_DYNAMIC_SHAPES = {}
+TRANSFORMER_NUM_FRAMES_DIM = torch.export.Dim('num_frames', min=3, max=21)
+
+TRANSFORMER_DYNAMIC_SHAPES = {
+    'hidden_states': {
+        2: TRANSFORMER_NUM_FRAMES_DIM,
+    },
+}
 
 INDUCTOR_CONFIGS = {
     'conv_1x1_as_mm': True,
@@ -36,7 +42,7 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
 
     @spaces.GPU(duration=1500)
     def compile_transformer():
-
+        
         pipeline.load_lora_weights(
            "vrgamedevgirl84/Wan14BT2VFusioniX", 
            weight_name="FusionX_LoRa/Phantom_Wan_14B_FusionX_LoRA.safetensors", 
@@ -116,13 +122,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)