Pre-Compile AoTI

app.py

@@ -5,10 +5,10 @@ import torch
 import spaces
 
 from PIL import Image
-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 os
 import base64
@@ -164,12 +164,14 @@ device = "cuda" if torch.cuda.is_available() else "cpu"
 # Load the model pipeline
 pipe = QwenImageEditPlusPipeline.from_pretrained("Qwen/Qwen-Image-Edit-2509", torch_dtype=dtype).to(device)
 
+spaces.aoti_blocks_load(pipe.transformer, "zerogpu-aoti/Qwen-Image", variant="fa3")
+
 # Apply the same optimizations from the first version
-pipe.transformer.__class__ = QwenImageTransformer2DModel
-pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
+# pipe.transformer.__class__ = QwenImageTransformer2DModel
+# pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
 
 # --- Ahead-of-time compilation ---
-optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")
+# optimize_pipeline_(pipe, image=[Image.new("RGB", (1024, 1024)), Image.new("RGB", (1024, 1024))], prompt="prompt")
 
 # --- UI Constants and Helpers ---
 MAX_SEED = np.iinfo(np.int32).max

requirements.txt

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