Update optimization.py

optimization.py

@@ -14,8 +14,8 @@ from torch.utils._pytree import tree_map
 P = ParamSpec('P')
 
 
-TRANSFORMER_IMAGE_SEQ_LENGTH_DIM = torch.export.Dim('image_seq_length')
-TRANSFORMER_TEXT_SEQ_LENGTH_DIM = torch.export.Dim('text_seq_length')
+TRANSFORMER_IMAGE_SEQ_LENGTH_DIM = torch.export.Dim('image_seq_length', min=1024, max=16384)
+TRANSFORMER_TEXT_SEQ_LENGTH_DIM = torch.export.Dim('text_seq_length', min=64, max=1024)
 
 TRANSFORMER_DYNAMIC_SHAPES = {
     'hidden_states': {
@@ -27,11 +27,10 @@ TRANSFORMER_DYNAMIC_SHAPES = {
     'encoder_hidden_states_mask': {
         1: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
     },
-    #'image_rotary_emb': ({
-    #    0: TRANSFORMER_IMAGE_SEQ_LENGTH_DIM,
-    #}, {
-    #    0: TRANSFORMER_TEXT_SEQ_LENGTH_DIM,
-    #}),
+    'image_rotary_emb': (
+        {0: TRANSFORMER_IMAGE_SEQ_LENGTH_DIM},  # vid_freqs: [img_seq_len, rope_dim]
+        {0: TRANSFORMER_TEXT_SEQ_LENGTH_DIM},   # txt_freqs: [txt_seq_len, rope_dim]
+    ),
 }