import gradio as gr
import moviepy.video.io.ImageSequenceClip
from PIL import Image
from pydub import AudioSegment
from moviepy.editor import *
import numpy as np
import mutagen
from mutagen.mp3 import MP3
import cv2

HF_TOKEN = os.environ.get("HF_TOKEN")


def resize(img_list):
    resize_img_list = []
    for item in img_list:
        im = Image.open(item)
        imResize = im.resize((256,256), Image.ANTIALIAS)
        resize_img_list.append(np.array(imResize))
    return resize_img_list
                

def merge_audio_video(entities_num, resize_img_list, text_input):
    speech = text2speech(text_input)
    wav_audio = AudioSegment.from_file(speech, "flac")  #("/content/gdrive/My Drive/AI/audio1.flac", "flac")
    wav_audio.export("audio.mp3", format="mp3")  #("/content/gdrive/My Drive/AI/audio1.mp3", format="mp3")
    audio_length = int(MP3("audio.mp3").info.length)
    fps= entities_num / audio_length  #length of audio file   
    fps = float(format(fps, '.5f'))
    clip = moviepy.video.io.ImageSequenceClip.ImageSequenceClip(resize_img_list, fps=fps)
    clip.write_videofile('my_vid_tmp.mp4')
    videoclip = VideoFileClip('my_vid_tmp.mp4') #("/content/gdrive/My Drive/AI/my_video1.mp4")
    audioclip = AudioFileClip('audio.mp3') #.subclip(0, 15)
    mergedclip = videoclip.set_audio(audioclip)
    duration = mergedclip.duration
    frame_count = mergedclip.fps 
    return mergedclip
    
fastspeech = gr.Interface.load("huggingface/facebook/fastspeech2-en-ljspeech", api_key=HF_TOKEN)

def text2speech(text):
    speech = fastspeech(text)
    return speech
    
def engine(text_input):
    ner = gr.Interface.load("huggingface/flair/ner-english-ontonotes-large", api_key=HF_TOKEN)
    entities = ner(text_input)
    entities = [tupl for tupl in entities if None not in tupl]
    entities_num = len(entities)
    img_list = []
    for ent in entities:
        img = gr.Interface.load("spaces/multimodalart/latentdiffusion, api_key=HF_TOKEN")(ent[0],'50','256','256','1',10)[0] 
        img_list.append(img)
    resize_img_list = resize(img_list)
    mergedclip = merge_audio_video(entities_num, resize_img_list, text_input)
    mergedclip.to_videofile('mergedvideo.mp4')
    return 'mergedvideo.mp4' 
    
app = gr.Interface(engine, 
                   gr.inputs.Textbox(lines=5, label="Input Text"),
                   gr.outputs.Video(type=None, label='Final Merged video'), 
                   description="<div>🎭🎞️🍿 AI Movie Maker - Comedy 🎬 🧠 🎨</div>" ,
                   examples=
[
["Two space marines take up arms to save the planet from an alien invasion.  These two dashing strong men play a comedic role in the science fiction movie of the future where even barnaby bunny is willing to join their wacky gang of space marines to save the planet with good looks and comedy."],
["Two space marines take up arms to save the planet from an alien invasion. In each episode they play different roles with the main mission to save Wills girlfriend from alien attack. These two dashing strong men play a comedic role in the science fiction movie of the future, and make their second feature film appearance together in a comedy.  In the short-lived TV series The Twilight Space Marines, Two space marines take up arms to save the planet from an alien invasion."],
["These two dashing strong men play a comedic role in the science fiction movie of the future but also serve as a bit of sci-fi history lesson the first time we saw them one of them held a gun to the head of another in anger due to a wrestling disagreement about a famous clearwater bar karaoke and wrestling champion. Barnaby bunny is willing to join their wacky gang of space marines for their adventure in the universe. But, its not just any adventure. They have got to survive a dangerous world, find their destiny, and defeat a cosmic villain before These two dashing strong men play a comedic role in the science fiction movie of the future."],
["To save the planet with good looks and comedy is not easy and they are both pro expert mobsters and wrestling federation superstars. It is a wonder I have been able to continue working as I have when so many others have given up  And here he is, a mere nine years after the film release, interviewed. Two space marines who take up arms and wrestling moves and jokes to save the planet from invading bad guys.  it is hilarious. That is why they are so perfect as a comedic double act, so we asked if they will be up for a sequel to charm us once again.  Dudes we are ready!"]
],
                   title="AI Pipeline Multi Model 🎭🎞️🍿 Movie Maker 🎬 🧠 🎨",
                   article="<br><div></div>"

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