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

# মডেল লোড করা (ফ্রি সিপইউ এর জন্য অপ্টিমাইজড)
model_id = "timbrooks/instruct-pix2pix"
pipe = StableDiffusionInstructPix2PixPipeline.from_pretrained(
    model_id, torch_dtype=torch.float32, safety_checker=None
)
pipe.to("cpu") # আপনার ফ্রি স্পেসের জন্য সিপইউ ব্যবহার করা হয়েছে

def edit_image(input_image, instruction):
    if input_image is None or instruction == "":
        return None
    
    # এআই প্রসেসিং
    edited_image = pipe(
        prompt=instruction, 
        image=input_image, 
        num_inference_steps=5, # সিপইউ তে দ্রুত করার জন্য স্টেপ কমানো হয়েছে
        image_guidance_scale=1.5,
        guidance_scale=7.5
    ).images[0]
    
    return edited_image

# ইন্টারফেস ডিজাইন
with gr.Blocks() as demo:
    gr.Markdown("## 🎨 AI Photo Editor (Prompt Based)")
    with gr.Row():
        with gr.Column():
            img_in = gr.Image(type="pil", label="আপনার ছবি দিন")
            prompt_in = gr.Textbox(label="কি এডিট করতে চান?", placeholder="যেমন: Make the dress red or change background to moon")
            btn = gr.Button("এডিট করুন")
        with gr.Column():
            img_out = gr.Image(type="pil", label="এডিট করা ছবি")

    btn.click(edit_image, inputs=[img_in, prompt_in], outputs=img_out)

demo.launch()