Fix dynamic shapes

optimization.py

@@ -7,6 +7,7 @@ from typing import ParamSpec
 
 import spaces
 import torch
+from torch.utils._pytree import tree_map_only
 
 from pipeline_utils import capture_component_call
 from zerogpu import aoti_compile
@@ -34,7 +35,8 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
             pipeline(*args, **kwargs)
 
         hidden_dim = torch.export.Dim('hidden', min=4096, max=8212)
-        dynamic_shapes = {
+        dynamic_shapes = tree_map_only((torch.Tensor, bool), lambda t: None, call.kwargs)
+        dynamic_shapes |= {
             'hidden_states': {1: hidden_dim},
             'img_ids': {0: hidden_dim},
         }