Update app.py

app.py

@@ -221,28 +221,20 @@ label, .gr-box label {
     border-bottom-style: solid;
 }
 """
-base_model_id = "Tongyi-MAI/Z-Image-Turbo"
-
 # Load the pipeline once at startup
 print("Loading Z-Image-Turbo pipeline...")
-#import torch
-from diffusers.models import ZImageTransformer2DModel
-from huggingface_hub import hf_hub_download
-
-model = ZImageTransformer2DModel.from_single_file(
-    hf_hub_download(
-        "xiaopeng/Z-IMAGE-Update", "zImageTurbo_40BF16AIO.safetensors"
-    ),
+pipe = DiffusionPipeline.from_pretrained(
+    "Tongyi-MAI/Z-Image-Turbo",
     torch_dtype=torch.bfloat16,
+    low_cpu_mem_usage=False,
 )
+pipe.to("cuda")
 
-# Build the Turbo pipeline once with our custom transformer weights.
-pipe = ZImagePipeline.from_pretrained(
-    base_model_id,
-    torch_dtype=torch.bfloat16,
-    transformer=model,
-)
-pipe = pipe.to("cuda")
+# ======== AoTI compilation + FA3 ========
+# pipe.transformer.layers._repeated_blocks = ["ZImageTransformerBlock"]
+# spaces.aoti_blocks_load(pipe.transformer.layers, "zerogpu-aoti/Z-Image", variant="fa3")
+
+print("Pipeline loaded!")
 
 @spaces.GPU
 def generate_image(