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

# Model ကို loading လုပ်ခြင်း (Qwen-Image-Edit-2509 အတွက် bfloat16 သို့မဟုတ် float16 သုံးနိုင်သည်)
model_id = "Qwen/Qwen-Image-Edit-2509"
pipe = QwenImageEditPlusPipeline.from_pretrained(model_id, torch_dtype=torch.bfloat16)
pipe.to("cuda")

def edit_image(input_image, prompt):
    if input_image is None or not prompt:
        return None
    
    # Model သို့ ပုံနှင့် Prompt ပေးပို့ခြင်း
    output = pipe(image=input_image, prompt=prompt).images[0]
    return output

# Gradio Interface တည်ဆောက်ခြင်း
with gr.Blocks() as demo:
    gr.Markdown("# 🎨 Qwen-Image-Edit-2509 AI Editor")
    gr.Markdown("Natural Language အသုံးပြုပြီး မိမိစိတ်ကြိုက် ပုံများကို ပြုပြင်ပြောင်းလဲလိုက်ပါ။")
    
    with gr.Row():
        with ui_col := gr.Column():
            image_input = gr.Image(type="pil", label="မူရင်းပုံတင်ရန်")
            prompt_input = gr.Textbox(lines=2, placeholder="ဥပမာ - Change the background to a sunny beach...", label="လုပ်ဆောင်စေလိုသည့် ညွှန်ကြားချက် (Prompt)")
            submit_btn = gr.Button("ပုံပြောင်းလဲမည်", variant="primary")
        
        with ui_col2 := gr.Column():
            image_output = gr.Image(type="pil", label="ရလဒ်ပုံစံ")
            
    submit_btn.click(
        fn=edit_image,
        inputs=[image_input, prompt_input],
        outputs=image_output
    )

demo.launch()