import gradio as gr
from transformers import BlipProcessor, BlipForConditionalGeneration, MarianMTModel, MarianTokenizer
from PIL import Image
import torch

# Load BLIP model for image captioning
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")

# Load MarianMT model for English to Arabic translation
translator_model_ar = MarianMTModel.from_pretrained("Helsinki-NLP/opus-mt-en-ar")
translator_tokenizer_ar = MarianTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-ar")

# Prediction function
def predict(image):
    inputs = processor(image, return_tensors="pt")
    out = model.generate(**inputs)
    english_caption = processor.decode(out[0], skip_special_tokens=True)

    # Translate to Arabic
    tokens = translator_tokenizer_ar.prepare_seq2seq_batch([english_caption], return_tensors="pt")
    translated = translator_model_ar.generate(**tokens)
    arabic_caption = translator_tokenizer_ar.decode(translated[0], skip_special_tokens=True)

    return arabic_caption

# Gradio interface
demo = gr.Interface(fn=predict, inputs=gr.Image(type="pil"), outputs="text", title="Image Captioning (English to Arabic)")
demo.launch()