import sys
from transformers.models.mt5 import tokenization_mt5
import transformers

# MT5Tokenizer ကို manually အတင်းထည့်ပေးခြင်း
transformers.MT5Tokenizer = tokenization_mt5.MT5Tokenizer
import gradio as gr
import torch
from diffusers import AutoPipelineForInpainting
from PIL import Image

# 1. Model ကို Load လုပ်ခြင်း (Mac GPU - MPS သုံးရန်)
device = "mps" if torch.backends.mps.is_available() else "cpu"
model_id = "diffusers/stable-diffusion-xl-1.0-inpainting-0.1"

pipe = AutoPipelineForInpainting.from_pretrained(
    model_id, 
    torch_dtype=torch.float16, 
    variant="fp16"
).to(device)

def extend_photo(image_data, prompt):
    # image_data မှာ image နဲ့ mask ပါဝင်ပါတယ်
    image = image_data["background"].convert("RGB")
    mask = image_data["layers"][0].convert("RGB")
    
    # Photo Extend လုပ်ငန်းစဉ်
    output = pipe(
        prompt=prompt,
        image=image,
        mask_image=mask,
        num_inference_steps=30
    ).images[0]
    
    return output

# 2. Gradio Interface တည်ဆောက်ခြင်း
with gr.Blocks() as demo:
    gr.Markdown("## SDXL Photo Extender (Mac Local)")
    with gr.Row():
        input_img = gr.ImageMask(label="ပုံထည့်ပြီး Extend လုပ်မည့်နေရာကို ခြစ်ပါ", type="pil")
        output_img = gr.Image(label="ရလဒ်")
    prompt_text = gr.Textbox(label="ဘာပုံ ဆက်ဆွဲရမလဲဆိုတာ ရေးပါ (ဥပမာ- beautiful forest background)")
    run_btn = gr.Button("Extend Photo")
    
    run_btn.click(fn=extend_photo, inputs=[input_img, prompt_text], outputs=output_img)

# 3. Launch လုပ်ခြင်း
# share=True ထားလျှင် Online ကနေပါ လှမ်းသုံးနိုင်တဲ့ link ရပါမယ်
demo.launch(share=True)