import gradio as gr
from rembg import remove, new_session
from PIL import Image

# ১. মডেল সেশন গ্লোবাল করা (একবারই লোড হবে)
print("ইঞ্জিন চালু হচ্ছে... মডেল লোড করা হচ্ছে।")
SESSION = new_session("birefnet-general")

def process_image(img):
    if img is None:
        return None
    try:
        # Alpha Matting অপ্টিমাইজড প্রসেসিং
        output = remove(
            img,
            session=SESSION,
            alpha_matting=True,
            alpha_matting_foreground_threshold=240,
            alpha_matting_background_threshold=10,
            alpha_matting_erode_size=10,
            post_process_mask=True
        )
        return output
    except Exception as e:
        print(f"Error: {e}")
        return img

# ২. ইন্টারফেস ডিজাইন (Gradio 6.0 এর নিয়ম অনুযায়ী)
with gr.Blocks() as demo:
    gr.HTML("""
        <div style="text-align: center;">
            <h1 style="color: #FF4B4B; margin-bottom: 0;">🚀 প্রো-গ্রেড এআই ব্যাকগ্রাউন্ড রিমুভার</h1>
            <p style="color: #555;">BiRefNet মডেল ব্যবহার করে নিখুঁত স্বচ্ছ ছবি তৈরি করুন</p>
        </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            input_file = gr.Image(
                label="ছবি আপলোড করুন", 
                type="pil", 
                interactive=True
            )
            submit_btn = gr.Button("ব্যাকগ্রাউন্ড সরান ✨", variant="primary")
            
        with gr.Column(scale=1):
            # এখানে 'show_download_button' এরর দিচ্ছিল, তাই সেটি বাদ দেওয়া হয়েছে
            output_file = gr.Image(
                label="স্বচ্ছ ফলাফল", 
                type="pil"
            )

    submit_btn.click(
        fn=process_image, 
        inputs=input_file, 
        outputs=output_file
    )

    gr.Markdown("---")
    gr.Markdown("ℹ️ **তথ্য:** এটি ১৬জিবি র‍্যাম অপ্টিমাইজড। প্রথমবার ছবি প্রসেস করতে মডেল লোড হতে কয়েক সেকেন্ড সময় নিতে পারে।")

# ৩. লঞ্চ করার সময় থিম এবং সিএসএস পাস করা (নতুন নিয়ম)
if __name__ == "__main__":
    demo.queue(max_size=10).launch(
        theme=gr.themes.Soft(),
        ssr_mode=False
    )