Update optimization.py

optimization.py

@@ -1,6 +1,25 @@
 import torch
 import torchao
 from torchao.quantization import DEFAULT_INT4_AUTOQUANT_CLASS_LIST
+P = ParamSpec('P')
+
+
+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,
+    'epilogue_fusion': False,
+    'coordinate_descent_tuning': True,
+    'coordinate_descent_check_all_directions': True,
+    'max_autotune': True,
+    'triton.cudagraphs': True,
+}
 
 def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
     print("[optimize_pipeline_] Starting pipeline optimization")