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

# Load the pipeline
pipe = QwenImageEditPipeline.from_pretrained("Qwen/Qwen-Image-Edit").to("cuda")

def edit_image(image, prompt):
    result = pipe(
        image=image,
        prompt=prompt,
        num_inference_steps=30,
        generator=torch.manual_seed(0),
    ).images[0]
    return result

with gr.Blocks() as demo:
    gr.Markdown("## 🖼️ Qwen Image Edit Demo")
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type="pil")
            prompt = gr.Textbox(label="Edit instruction")
            btn = gr.Button("Generate")
        with gr.Column():
            output_image = gr.Image(type="pil")

    btn.click(fn=edit_image, inputs=[input_image, prompt], outputs=output_image)

demo.launch()