Update optimization.py

optimization.py

@@ -28,14 +28,16 @@ INDUCTOR_CONFIGS = {
 
 def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
 
+    blocks = pipeline.transformer.transformer_blocks
+
     @spaces.GPU(duration=120)
     def compile_transformer_block():
-
-        with spaces.aoti_capture(pipeline.transformer.transformer_blocks[0]) as call:
+        
+        with spaces.aoti_capture(blocks[0]) as call:
             pipeline(*args, **kwargs)
 
         exported = torch.export.export(
-            mod=pipeline.transformer,
+            mod=blocks[0],
             args=call.args,
             kwargs=call.kwargs,
         )
@@ -46,5 +48,5 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
     pipeline.transformer.set_attn_processor(FlashFusedFluxAttnProcessor3_0())
 
     archive_file = compile_transformer_block()
-    for block in pipeline.transformer.transformer_blocks:
+    for block in blocks:
         block.forward = ZeroGPUCompiledModel(archive_file, ZeroGPUWeights(block.state_dict()))