Update optimization.py

optimization.py

@@ -8,9 +8,8 @@ from typing import ParamSpec
 import spaces
 import torch
 from torch.utils._pytree import tree_map_only
-
-from optimization_utils import capture_component_call
-from optimization_utils import aoti_compile
+from torchao.quantization import quantize_
+from torchao.quantization import Float8DynamicActivationFloat8WeightConfig
 
 
 P = ParamSpec('P')
@@ -38,7 +37,7 @@ 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), lambda t: None, call.kwargs)
@@ -46,6 +45,8 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
 
         pipeline.transformer.fuse_qkv_projections()
 
+        quantize_(pipeline.transformer, Float8DynamicActivationFloat8WeightConfig())
+        
         exported = torch.export.export(
             mod=pipeline.transformer,
             args=call.args,
@@ -53,8 +54,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)