Regional AoT

optimization.py

@@ -4,8 +4,11 @@
 from typing import Any
 from typing import Callable
 from typing import ParamSpec
+
 import spaces
 import torch
+from spaces.zero.torch.aoti import ZeroGPUCompiledModel
+from spaces.zero.torch.aoti import ZeroGPUWeights
 
 from fa3 import FlashFusedFluxAttnProcessor3_0
 
@@ -25,10 +28,10 @@ INDUCTOR_CONFIGS = {
 
 def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kwargs):
 
-    @spaces.GPU(duration=1500)
-    def compile_transformer():
+    @spaces.GPU(duration=120)
+    def compile_transformer_block():
 
-        with spaces.aoti_capture(pipeline.transformer) as call:
+        with spaces.aoti_capture(pipeline.transformer.transformer_blocks[0]) as call:
             pipeline(*args, **kwargs)
 
         exported = torch.export.export(
@@ -37,8 +40,11 @@ def optimize_pipeline_(pipeline: Callable[P, Any], *args: P.args, **kwargs: P.kw
             kwargs=call.kwargs,
         )
 
-        return spaces.aoti_compile(exported, INDUCTOR_CONFIGS)
+        return spaces.aoti_compile(exported, INDUCTOR_CONFIGS).archive_file
 
     pipeline.transformer.fuse_qkv_projections()
     pipeline.transformer.set_attn_processor(FlashFusedFluxAttnProcessor3_0())
-    spaces.aoti_apply(compile_transformer(), pipeline.transformer)
+
+    archive_file = compile_transformer_block()
+    for block in pipeline.transformer.transformer_blocks:
+        block.forward = ZeroGPUCompiledModel(archive_file, ZeroGPUWeights(block.state_dict()))