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

device = 'cpu'

# Load the pretrained model, feature extractor, and tokenizer
model = VisionEncoderDecoderModel.from_pretrained("premanthcharan/Image_Captioning_Model").to(device)
feature_extractor = ViTImageProcessor.from_pretrained("premanthcharan/Image_Captioning_Model")
tokenizer = AutoTokenizer.from_pretrained("premanthcharan/Image_Captioning_Model")

def predict(image, max_length=64, num_beams=4):
    # Process the input image
    image = image.convert('RGB')
    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values.to(device)
    
    # Generate the caption
    caption_ids = model.generate(pixel_values, max_length=max_length, num_beams=num_beams)[0]
    
    # Decode and clean the generated caption
    caption = tokenizer.decode(caption_ids, skip_special_tokens=True)
    return caption

css = '''
h1#title {
  text-align: center;
}
h3#header {
  text-align: center;
}
img#overview {
  max-width: 800px;
  max-height: 600px;
}
img#style-image {
  max-width: 1000px;
  max-height: 600px;
}
'''

demo = gr.Blocks(css=css)

with demo:
    gr.Markdown('''<h1 id="title">Automated Image Captioning Using Generative AI: A Transformer based approach 🖼️</h1>''')
    gr.Markdown('Contributed by : Charan Gudivada, Premanth Alahari')
    
    with gr.Column():
        input_image = gr.Image(label="Upload your Image", type='pil')
        output_caption = gr.Textbox(label="Generated Caption")
        
    btn = gr.Button("Generate Caption")
    btn.click(fn=predict, inputs=input_image, outputs=output_caption)
    
with demo:
    gr.Markdown('''<h1 id="title">Features:</h1>''')
    gr.Markdown("1. Drop the Image Here or Click on Upload\n2. Click to Access Webcam\n3. Paste from Clipboard")
demo.launch()