Update app.py

app.py

@@ -1,14 +1,21 @@
 import torch
 import spaces
 import gradio as gr
-from diffusers import DiffusionPipeline
+from diffusers import DiffusionPipeline, ZImagePipeline, ZImageTransformer2DModel
 
 from PIL import Image
 
 # Load the pipeline once at startup
 print("Loading Z-Image-Turbo pipeline...")
-pipe = DiffusionPipeline.from_pretrained(
+transformer = ZImageTransformer2DModel.from_pretrained(
+    "linoyts/beyond-reality-z-image-diffusers",
+    torch_dtype=torch.bfloat16,
+)
+
+# Load pipeline with custom transformer
+pipe = ZImagePipeline.from_pretrained(
     "Tongyi-MAI/Z-Image-Turbo",
+    transformer=transformer,
     torch_dtype=torch.bfloat16,
     low_cpu_mem_usage=False,
 )