import torch
import re
import gradio as gr

from transformers import AutoTokenizer, ViTImageProcessor, VisionEncoderDecoderModel, AutoModelForSeq2SeqLM, pipeline
device='cpu'

encoder_checkpoint = "nlpconnect/vit-gpt2-image-captioning"
decoder_checkpoint = "nlpconnect/vit-gpt2-image-captioning"
model_checkpoint = "nlpconnect/vit-gpt2-image-captioning"

feature_extractor = ViTImageProcessor.from_pretrained(encoder_checkpoint)
tokenizer = AutoTokenizer.from_pretrained(decoder_checkpoint)
model = VisionEncoderDecoderModel.from_pretrained(model_checkpoint).to(device)
mod = AutoModelForSeq2SeqLM.from_pretrained("KarmaCST/nllb-200-distilled-600M-en-to-dz")
tok = AutoTokenizer.from_pretrained("KarmaCST/nllb-200-distilled-600M-en-to-dz")

src_lang = 'eng_Latn'
tgt_lang = "dzo_Tibt"

def predict(image,max_length=64, num_beams=4):
    image = image.convert('RGB')
    image = feature_extractor(image, return_tensors="pt").pixel_values.to(device)
    clean_text = lambda x: x.replace('<|endoftext|>','').split('\n')[0]
    caption_ids = model.generate(image, max_length = max_length)[0]
    caption_text = clean_text(tokenizer.decode(caption_ids))     
    translation_pipeline = pipeline("translation",model=mod,tokenizer=tok,src_lang=src_lang,tgt_lang=tgt_lang)
    result = translation_pipeline(caption_text)
    return result[0]['translation_text']

input = gr.Image(label="Upload any Image", type = 'pil')
output = gr.Textbox(type="text",label="Captions")  # Change type to "text"

examples = [f"example{i}.jpg" for i in range(1,7)]
title = "Image Captioning in Dzongkha "
interface = gr.Interface(        
    fn=predict,        
    inputs = input,        
    theme="grass",        
    outputs=output,        
    examples = examples,        
    title=title,        
    article="<h1>Created By:</h1>Mr. Karma Wangchuk<br>Lecturer<br>Information Technology Department<br>College of Science and Technology<br>Rinchending Phuentsholing<br>Chhukha Bhutan<br>",    
)
interface.launch(debug=True)