Update optimization.py

optimization.py

@@ -11,10 +11,6 @@ from torch.utils._pytree import tree_map_only
 from torchao.quantization import quantize_
 from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
 
-from optimization_utils import capture_component_call
-from optimization_utils import aoti_compile
-from optimization_utils import cudagraph
-
 
 P = ParamSpec('P')
 
@@ -41,13 +37,11 @@ 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:
+        with spaces.aoti_capture(pipeline.transformer) as call:
             pipeline(*args, **kwargs)
 
         dynamic_shapes = tree_map_only((torch.Tensor, bool, int), lambda t: None, call.kwargs)
         dynamic_shapes |= TRANSFORMER_DYNAMIC_SHAPES
-
-        quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
         
         exported = torch.export.export(
             mod=pipeline.transformer,
@@ -56,8 +50,6 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
             dynamic_shapes=dynamic_shapes,
         )
 
-        return aoti_compile(exported, INDUCTOR_CONFIGS)
+        return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
 
-    transformer_config = pipeline.transformer.config
-    pipeline.transformer = compile_transformer()
-    pipeline.transformer.config = transformer_config # pyright: ignore[reportAttributeAccessIssue]
+    spaces.aoti_apply(compile_transformer(), pipeline.transformer)