import gradio as gr
from transformers import pipeline
import io, base64
from PIL import Image
import numpy as np
import tensorflow as tf
import mediapy
import os
import sys
from huggingface_hub import snapshot_download

image_gen = gr.Interface.load("spaces/multimodalart/latentdiffusion")

os.system("git clone https://github.com/google-research/frame-interpolation")
sys.path.append("frame-interpolation")
from eval import interpolator, util

ffmpeg_path = util.get_ffmpeg_path()
mediapy.set_ffmpeg(ffmpeg_path)

model = snapshot_download(repo_id="akhaliq/frame-interpolation-film-style")
interpolator = interpolator.Interpolator(model, None)

def generate_story(choice, input_text):
    query = "<BOS> <{0}> {1}".format(choice, input_text)
    
    print(query)
    generated_text = story_gen(query)
    generated_text = generated_text[0]['generated_text']
    generated_text = generated_text.split('> ')[2]
    
    return generated_text
    
def generate_images(text):
    steps=50
    width=256
    height=256
    num_images=4
    diversity=4
    image_bytes = image_gen(text, steps, width, height, num_images, diversity)
    
    # Algo from spaces/Gradio-Blocks/latent_gpt2_story/blob/main/app.py
    generated_images = []
    for image in image_bytes[1]:
        image_str = image[0]
        image_str = image_str.replace("data:image/png;base64,","")
        decoded_bytes = base64.decodebytes(bytes(image_str, "utf-8"))
        img = Image.open(io.BytesIO(decoded_bytes))
        generated_images.append(img)
        
    return generated_images
    
def generate_interpolation(text):
    times_to_interpolate = 4
    
    generated_images = generate_images(text)
    
    generated_images[0].save('frame_0.png')
    generated_images[1].save('frame_1.png')
    generated_images[2].save('frame_2.png')
    generated_images[3].save('frame_3.png')
    
    input_frames = ["frame_0.png", "frame_1.png", "frame_2.png", "frame_3.png"]

    frames = list(util.interpolate_recursively_from_files(input_frames, times_to_interpolate, interpolator))

    mediapy.write_video("out.mp4", frames, fps=7)
    
    return "out.mp4"
    
    

demo = gr.Blocks()

with demo:
    input_start_text = gr.Textbox(placeholder='A teddy bear outer space', label="Starting Text")
    button_gen_video = gr.Button("Generate Video")
    output_interpolation = gr.Video(label="Generated Video")
            
    button_gen_video.click(fn=generate_interpolation, inputs=input_start_text, outputs=output_interpolation)

demo.launch(debug=True, enable_queue=True)