Pre-Compiled AoTI

app.py

@@ -6,10 +6,10 @@ import spaces
 
 from PIL import Image
 
-from optimization import optimize_pipeline_
+# from optimization import optimize_pipeline_
 from qwenimage.pipeline_qwen_image_edit import QwenImageEditPipeline
-from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
-from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
+# from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
+# from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3
 
 import os
 import base64
@@ -144,11 +144,14 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load the model pipeline
 pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit", torch_dtype=dtype).to(device)
-pipe.transformer.__class__ = QwenImageTransformer2DModel
-pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+
+spaces.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/Qwen-Image", variant="fa3")
+
+# pipe.transformer.__class__ = QwenImageTransformer2DModel
+# pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
 
 # --- Ahead-of-time compilation ---
-optimize_pipeline_(pipe, image=Image.new("RGB", (1024, 1024)), prompt="prompt")
+# optimize_pipeline_(pipe, image=Image.new("RGB", (1024, 1024)), prompt="prompt")
 
 # --- UI Constants and Helpers ---
 MAX_SEED = np.iinfo(np.int32).max

requirements.txt

@@ -7,4 +7,5 @@ safetensors
 sentencepiece
 dashscope
 torchvision
-peft
+peft
+torch==2.8