Update optimization.py

optimization.py

@@ -8,16 +8,19 @@ from torchao.quantization import quantize_
 from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
 import spaces
 import torch
-from torch.utils._pytree import tree_map, tree_map_only
+from torch.utils._pytree import tree_map
 
 
 P = ParamSpec('P')
 
 
-TRANSFORMER_HIDDEN_DIM = torch.export.Dim.AUTO(min=3584, max=9727)
+TRANSFORMER_IMAGE_SEQ_LENGTH_DIM = torch.export.Dim.AUTO(min=3584, max=9727)
+TRANSFORMER_IMAGE_HEIGHT_DIM = torch.export.Dim.DYNAMIC
+TRANSFORMER_IMAGE_WIDTH_DIM = torch.export.Dim.DYNAMIC
 
 TRANSFORMER_DYNAMIC_SHAPES = {
-    'hidden_states': {1: TRANSFORMER_HIDDEN_DIM},
+    'hidden_states': {1: TRANSFORMER_IMAGE_SEQ_LENGTH_DIM},
+    'img_shapes': [(None, TRANSFORMER_IMAGE_HEIGHT_DIM, TRANSFORMER_IMAGE_WIDTH_DIM)]
 }
 
 
@@ -40,7 +43,6 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
             pipeline(*args, **kwargs)
 
         dynamic_shapes = tree_map(lambda t: None, call.kwargs)
-        # dynamic_shapes = tree_map_only((torch.Tensor, bool), lambda x: None, call.kwargs)
         dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
 
         quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
@@ -49,7 +51,7 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
             mod=pipeline.transformer,
             args=call.args,
             kwargs=call.kwargs,
-            # dynamic_shapes=dynamic_shapes,
+            dynamic_shapes=dynamic_shapes,
         )
 
         return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)