Update app.py

app.py

@@ -1,28 +1,31 @@
 import gradio as gr
+from diffusers import DiffusionPipeline
 import torch
 from PIL import Image
-from transformers import AutoProcessor, AutoModelForCausalLM
 
-model_id = "Phr00t/Qwen-Image-Edit-Rapid-AIO"
+MODEL_ID = "Phr00t/Qwen-Image-Edit-Rapid-AIO"
 
-processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    torch_dtype=torch.bfloat16,
-    device_map="auto",
+pipe = DiffusionPipeline.from_pretrained(
+    MODEL_ID,
+    torch_dtype=torch.float32,
+    safety_checker=None,
     trust_remote_code=True
 )
 
-def generate(prompt, input_img):
-    if input_img is not None:
-        img = Image.fromarray(input_img)
-        inputs = processor(prompt, image=img, return_tensors="pt").to(model.device)
+pipe.to("cpu")
+
+
+def generate(prompt, image):
+    if image is None:
+        # text → image
+        result = pipe(prompt)
+        return result.images[0]
     else:
-        inputs = processor(prompt, return_tensors="pt").to(model.device)
+        # image editing mode
+        img = Image.fromarray(image)
+        result = pipe(prompt=prompt, image=img)
+        return result.images[0]
 
-    output = model.generate(**inputs)
-    image = processor.batch_decode(output, output_type="pil")[0]
-    return image
 
 demo = gr.Interface(
     fn=generate,
@@ -30,9 +33,9 @@ demo = gr.Interface(
         gr.Textbox(label="Prompt"),
         gr.Image(label="Input Image (optional)")
     ],
-    outputs=gr.Image(label="Generated Image"),
-    title="Qwen Image Edit Rapid AIO",
-    description="Uses Qwen2-VL model for text-to-image and image editing."
+    outputs=gr.Image(label="Generated / Edited Image"),
+    title="Qwen Image Edit Rapid AIO (CPU Friendly)",
+    description="Works on CPU with Diffusers. Supports text-to-image AND image editing."
 )
 
 demo.launch()