Pre-Compiled AoTI

app.py

@@ -5,10 +5,10 @@ import torch
 import spaces
 from PIL import Image
 from diffusers import FlowMatchEulerDiscreteScheduler
-from optimization import optimize_pipeline_
+# from optimization import optimize_pipeline_
 from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
-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 math
 
 # --- Model Loading ---
@@ -51,9 +51,11 @@ pipe.load_lora_weights(
 )
 pipe.fuse_lora(lora_scale=1.0)
 
-pipe.transformer.__class__ = QwenImageTransformer2DModel
-pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
-optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")
+spaces.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/Qwen-Image", variant="fa3")
+
+#pipe.transformer.__class__ = QwenImageTransformer2DModel
+#pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+#optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")
 
 # --- Constants ---
 MAX_SEED = np.iinfo(np.int32).max

requirements.txt

@@ -8,4 +8,5 @@ dashscope
 kernels
 torchvision
 peft
-torchao==0.11.0
+torchao==0.11.0
+torch==2.8