import gradio as gr
from diffusers import StableDiffusionInpaintPipeline
import torch
from PIL import Image
import base64
import io

# Load the Qwen Inpainting model
pipe = StableDiffusionInpaintPipeline.from_pretrained(
    "InstantX/Qwen-Image-ControlNet-Inpainting",
    torch_dtype=torch.float16
).to("cuda")

def edit_image(image, mask, prompt):
    image = image.convert("RGB")
    mask = mask.convert("L")  # Black = edit, White = keep
    result = pipe(
        prompt=prompt,
        image=image,
        mask_image=mask,
        num_inference_steps=30
    ).images[0]
    return result

# Gradio UI
demo = gr.Interface(
    fn=edit_image,
    inputs=[
        gr.Image(type="pil", label="Input Image"),
        gr.Image(type="pil", label="Mask Image"),
        gr.Textbox(label="Prompt")
    ],
    outputs=gr.Image(type="pil", label="Edited Image"),
    title="Qwen Image Inpainting"
)

demo.launch()