AOT: Load Pre-compiled FLUX.2 blocks from hub (#5)

- AOT: Load Pre-compiled FLUX.2 blocks from hub (e7db332e79b3f00f7252605838483b3d6ec4bcfd)


Co-authored-by: Charles Bensimon <[email protected]>

app.py

@@ -9,7 +9,6 @@ import spaces
 import torch
 from diffusers import Flux2Pipeline, Flux2Transformer2DModel
 from diffusers import BitsAndBytesConfig as DiffBitsAndBytesConfig
-from optimization import optimize_pipeline_
 import requests
 from PIL import Image
 import json
@@ -81,18 +80,8 @@ pipe = Flux2Pipeline.from_pretrained(
 )
 pipe.to(device)
 
-pipe.transformer.set_attention_backend("_flash_3_hub")
-
-# Optimization runs once at startup
-optimize_pipeline_(
-    pipe,
-    image=[Image.new("RGB", (1024, 1024))],
-    prompt_embeds = remote_text_encoder("prompt").to(device),
-    guidance_scale=2.5,
-    width=1024,
-    height=1024,
-    num_inference_steps=1
-)
+# Pull pre-compiled Flux2 Transformer blocks from HF hub
+spaces.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/FLUX.2", variant="fa3")
 
 def image_to_data_uri(img):
     buffered = io.BytesIO()