Update optimization.py

optimization.py

@@ -46,11 +46,11 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
     latent_height = height // VAE_SPATIAL_COMPRESSION_RATIO
     latent_width = width // VAE_SPATIAL_COMPRESSION_RATIO
 
-    with capture_component_call(pipeline, "transformer") as call:
-        pipeline(*args, **kwargs)
-
     @spaces.GPU(duration=1500)
     def compile_transformer():
+        with capture_component_call(pipeline, "transformer") as call:
+            pipeline(*args, **kwargs)
+        
         dynamic_shapes = tree_map_only((torch.Tensor, bool), lambda t: None, call.kwargs)
         dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
 
@@ -129,8 +129,8 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
     pipeline.transformer = combined_transformer
     # pipeline.transformer = cudagraph(combined_transformer)
 
-    with torch.no_grad():
-        pipeline.transformer(**call.kwargs)
+    # with torch.no_grad():
+    #     pipeline.transformer(**call.kwargs)
 
     pipeline.transformer.config = transformer_config  # pyright: ignore[reportAttributeAccessIssue]
     pipeline.transformer.dtype = transformer_dtype  # pyright: ignore[reportAttributeAccessIssue]