import gradio as gr
import torch
from transformers import BlipProcessor, BlipForConditionalGeneration

MODEL_NAME = "Salesforce/blip-image-captioning-base"

# --- مدل را بارگیری کن
try:
    device = "cuda" if torch.cuda.is_available() else "cpu"
    caption_processor = BlipProcessor.from_pretrained(MODEL_NAME)
    caption_model = BlipForConditionalGeneration.from_pretrained(MODEL_NAME)
    caption_model.to(device)
    _load_error = None
except Exception as e:
    caption_processor = None
    caption_model = None
    _load_error = str(e)

# --- تابع captioning
def caption_image(image):
    if _load_error:
        return f"❌ Model load error: {_load_error}"
    if image is None:
        return "⚠️ لطفاً یک تصویر آپلود کنید."

    inputs = caption_processor(image, return_tensors="pt").to(device)
    out = caption_model.generate(**inputs, max_new_tokens=30)
    caption = caption_processor.decode(out[0], skip_special_tokens=True)
    return caption

# --- رابط کاربری Gradio
status_text = (
    "✅ Model loaded successfully"
    if caption_model is not None and not _load_error
    else f"❌ Error: {_load_error}"
    if _load_error
    else "⏳ Loading model..."
)

with gr.Blocks(title="Image Captioning App") as demo:
    gr.Markdown("## 🖼️ Image Captioning with BLIP\nUpload an image and get an automatic caption.")
    gr.Markdown(f"**Status:** {status_text}")

    with gr.Row():
        image_input = gr.Image(type="pil", label="Upload Image")
        caption_output = gr.Textbox(label="Generated Caption", interactive=False)

    generate_btn = gr.Button("Generate Caption")

    generate_btn.click(fn=caption_image, inputs=image_input, outputs=caption_output)

demo.launch()