Dynamic shapes

app.py

@@ -18,12 +18,7 @@ from optimization import optimize_pipeline_
 MAX_SEED = np.iinfo(np.int32).max
 
 pipe = FluxKontextPipeline.from_pretrained("black-forest-labs/FLUX.1-Kontext-dev", torch_dtype=torch.bfloat16).to("cuda")
-
-optimize_pipeline_(pipe,
-    image=Image.new('RGB', (512, 512)),
-    prompt='prompt',
-    guidance_scale=2.5,
-)
+optimize_pipeline_(pipe, image=None, prompt='prompt')
 
 @spaces.GPU
 def infer(input_image, prompt, seed=42, randomize_seed=False, guidance_scale=2.5, steps=28, progress=gr.Progress(track_tqdm=True)):

optimization.py

@@ -29,10 +29,25 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
 
     @spaces.GPU(duration=1500)
     def compile_transformer():
+
         with capture_component_call(pipeline, 'transformer') as call:
             pipeline(*args, **kwargs)
+
+        hidden_dim = torch.export.Dim('hidden', min=4096, max=8212)
+        dynamic_shapes = {
+            'hidden_states': {1: hidden_dim},
+            'img_ids': {0: hidden_dim},
+        }
+
         pipeline.transformer.fuse_qkv_projections()
-        exported = torch.export.export(pipeline.transformer, args=call.args, kwargs=call.kwargs)
+
+        exported = torch.export.export(
+            mod=pipeline.transformer,
+            args=call.args,
+            kwargs=call.kwargs,
+            dynamic_shapes=dynamic_shapes,
+        )
+
         return aoti_compile(exported, INDUCTOR_CONFIGS)
 
     transformer_config = pipeline.transformer.config