Float8DynamicActivation quantization

optimization.py

@@ -8,6 +8,8 @@ from typing import ParamSpec
 import spaces
 import torch
 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
@@ -46,6 +48,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,

requirements.txt

@@ -1,3 +1,4 @@
+torchao
 transformers
 git+https://github.com/huggingface/diffusers.git
 accelerate