import gradio as gr
from diffusers import DiffusionPipeline
import torch
from PIL import Image

MODEL_ID = "Phr00t/Qwen-Image-Edit-Rapid-AIO"

pipe = DiffusionPipeline.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.float32,
    safety_checker=None,
    trust_remote_code=True
)

pipe.to("cpu")


def generate(prompt, image):
    if image is None:
        # text → image
        result = pipe(prompt)
        return result.images[0]
    else:
        # image editing mode
        img = Image.fromarray(image)
        result = pipe(prompt=prompt, image=img)
        return result.images[0]


demo = gr.Interface(
    fn=generate,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Image(label="Input Image (optional)")
    ],
    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()