import gradio as gr
import torch
from transformers import GPT2TokenizerFast, ViTImageProcessor, VisionEncoderDecoderModel 

# Setup device, model, tokenizer, and feature extractor
device = 'cpu'
model_checkpoint = "Stoneman/IG-caption-generator-nlpconnect-all"
feature_extractor = ViTImageProcessor.from_pretrained(model_checkpoint)
tokenizer = GPT2TokenizerFast.from_pretrained(model_checkpoint)
model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint).to(device)

# Prediction function
def predict(image, max_length=128):
    image = image.convert('RGB')
    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device)
    caption_ids = model.generate(pixel_values, max_length=max_length)[0]
    caption_text = tokenizer.decode(caption_ids, skip_special_tokens=True)
    return caption_text

# Define input and output components
input_component = gr.components.Image(label="Upload any Image", type="pil")
output_component = gr.components.Textbox(label="Captions")

# Example images
examples = [f"example{i}.JPG" for i in range(1, 10)]

# Interface
title = "IG-caption-generator"
description = "Made by: Jiayu Shi"
interface = gr.Interface(
    fn=predict,
    description=description,
    inputs=input_component,
    theme="huggingface",
    outputs=output_component,
    examples=examples,
    title=title,
)

# Launch interface
interface.launch(debug=True)