import gradio as gr
import torch
from diffusers import StableDiffusionInpaintPipeline, DPMSolverMultistepScheduler
from PIL import Image

# ၁။ Model Card ရွေးချယ်မှု
model_id = "Sanster/Realistic_Vision_V1.4-inpainting"

# ၂။ Pipeline ကို Load လုပ်ပြီး Optimization များထည့်သွင်းခြင်း
pipe = StableDiffusionInpaintPipeline.from_pretrained(
    model_id, 
    torch_dtype=torch.float32 # CPU အတွက် ပိုငြိမ်သည်
)

# Speed အတွက် Scheduler ပြောင်းခြင်း
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)

# Safety Checker (NSFW Filter) ကို လုံးဝ ပိတ်ခြင်း
pipe.safety_checker = lambda images, **kwargs: (images, [False] * len(images))

# CPU Optimization များ
pipe.enable_attention_slicing()
# pipe.enable_sequential_cpu_offload() # RAM သိပ်မရှိလျှင် ဖွင့်ရန်

def predict(image_dict, prompt):
    # ပုံအရွယ်အစားကို 512 အောက်ထားခြင်းက အမြန်ဆုံးဖြစ်သည်
    init_image = image_dict["background"].convert("RGB").resize((512, 512))
    mask_image = image_dict["layers"][0].convert("RGB").resize((512, 512))
    
    # ၃။ အမြန်ဆုံး အနေအထားဖြင့် ပုံထုတ်ခြင်း (Steps ကို ၂၀ ထားပါ)
    output = pipe(
        prompt=prompt, 
        image=init_image, 
        mask_image=mask_image,
        num_inference_steps=20, # အမြန်နှုန်းအတွက် အဓိက အချက်
        guidance_scale=7.5
    ).images[0]
    
    return output

# ၄။ Gradio UI တည်ဆောက်ခြင်း
with gr.Blocks() as demo:
    gr.Markdown("### ⚡ Fast AI Inpainting (CPU Optimized)")
    with gr.Row():
        img = gr.Image(label="Upload Image & Paint Mask", tool="sketch", type="pil")
        prompt = gr.Textbox(label="Prompt (e.g., 'white shirt')")
    btn = gr.Button("Generate Fast")
    result = gr.Image(label="Result")
    
    btn.click(predict, inputs=[img, prompt], outputs=result)

demo.launch()