import gradio as gr
from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
import torch
import accelerate
from PIL import Image
from diffusers.utils import export_to_video

model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning")
tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning")

pipe = DiffusionPipeline.from_pretrained("damo-vilab/text-to-video-ms-1.7b", variant="fp16")
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
pipe.to(device)

max_length = 16
num_beams = 4
gen_kwargs = {"max_length": max_length, "num_beams": num_beams}
def image_to_text(image_paths):
    images=[image_paths]
    
    pixel_values = feature_extractor(images=images, return_tensors="pt").pixel_values
    pixel_values = pixel_values.to(device)

    output_ids = model.generate(pixel_values, **gen_kwargs)

    preds = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
    preds = [pred.strip() for pred in preds]
    return preds[0]

def text_to_video(image_paths):
    prompt = image_to_text(image_paths)
    video_frames = pipe(prompt, num_inference_steps=25).frames
    video_path = export_to_video(video_frames)
    return video_path

article = """
## Descripción del proyecto
##### Queremos obtener un vídeo a partir de una imagen que tengamos. Por lo tanto, lo solucionaremos utilizando primero, un modelo de "image captioning" que nos genera un texto a partir de una imagen, y segundo utilizaremos un modelo que genera un video a partir de un texto. De esta manera el cliente en la interfaz deberá de subir como input un archivo de tipo imagen y recibirá como output un archivo de vídeo de unos pocos segundos.
### El primer modelo es el siguiente: https://huggingface.co/nlpconnect/vit-gpt2-image-captioning. Este modelo será el que reciba nuestro input, el cual deberá ser un archivo de tipo imagen. Como output nos da un texto, que será el input del segundo modelo.
### El segundo modelo es el siguiente: https://huggingface.co/damo-vilab/text-to-video-ms-1.7b. Este modelo debe recibir como input un texto, por ello, cuando este modelo reciba el texto del primer modelo generará el vídeo como output que será el output final que se mostrará en la interfaz al cliente que lo esté utilizando. Este modelo sólo puede recibir los inputs en inglés, pero no hay problema ya que el priemr modelo da el outpiut en inglés.
### Un problerma de la aplicación es que necesita bastantes recursos, y debido a los recursos tan limitados que nos ofrece Hugging Face gratuito deberemos esperar bastante tiempo desde que le damos el input hasta que nos de el output, unos 2000 segundos más o menos. Obviando ese detalle esta aplicación es muy útil para obtener un vídeo no muy largo de una imagen.
"""

title = "Proyecto 1: Aplicación de imagen a vídeo utilizando dos modelos de IA"
description = "En este proyecto, utilizando la siguiente interfaz, puedes obtener un vídeo corto a partir de una imagen"

interface = gr.Interface(
    fn=text_to_video, 
    inputs=gr.Image(type="pil"),
    outputs=gr.Video(), 
    title=title,
    description=description,
    article=article,

)

interface.launch(debug=True)