import gradio as gr
import torch
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration, VisionEncoderDecoderModel, ViTFeatureExtractor, AutoTokenizer

# Load BLIP model
blip_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")
blip_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")

# Load ViT-GPT2 model
gpt2_model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
gpt2_feature_extractor = ViTFeatureExtractor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
gpt2_tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")

# Set device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
gpt2_model.to(device)

# Generation parameters
max_length = 16
num_beams = 4
gen_kwargs = {"max_length": max_length, "num_beams": num_beams}

def blip_caption(img_path, min_len, max_len):
    raw_image = Image.open(img_path).convert('RGB')
    inputs = blip_processor(raw_image, return_tensors="pt")
    out = blip_model.generate(**inputs, min_length=min_len, max_length=max_len)
    return blip_processor.decode(out[0], skip_special_tokens=True)

def gpt2_caption(img_path):
    raw_image = Image.open(img_path).convert("RGB")
    pixel_values = gpt2_feature_extractor(images=raw_image, return_tensors="pt").pixel_values
    pixel_values = pixel_values.to(device)
    output_ids = gpt2_model.generate(pixel_values, **gen_kwargs)
    preds = gpt2_tokenizer.batch_decode(output_ids, skip_special_tokens=True)
    return preds[0].strip()

def generate_captions(img, min_len, max_len):
    blip_result = blip_caption(img, min_len, max_len)
    gpt2_result = gpt2_caption(img)
    return blip_result, gpt2_result

iface = gr.Interface(
    fn=generate_captions,
    inputs=[
        gr.Image(type='filepath', label='Image'),
        gr.Slider(label='Minimum Length', minimum=1, maximum=100, value=30),
        gr.Slider(label='Maximum Length', minimum=1, maximum=1000, value=100)
    ],
    outputs=[
        gr.Textbox(label='BLIP Caption'),
        gr.Textbox(label='GPT-2 Caption')
    ],
    title='Image Captioning',
    description="This application generates descriptive captions for images using two advanced models: BLIP and ViT-GPT-2. Simply upload an image and receive two unique captions, showcasing different perspectives from each model. Customize the caption length with easy-to-use sliders and enjoy a seamless, interactive experience. Perfect for content creation, accessibility, and research.",
    theme=gr.themes.Base(primary_hue="teal", secondary_hue="teal", neutral_hue="slate"),
)

iface.launch()