import torch
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
from diffusers.utils import export_to_video
from IPython.display import HTML
from base64 import b64encode
import imageio
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from skimage.transform import resize
from IPython.display import HTML
import gradio as gr

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

def make_video(vid_prompt, duration):
  print(duration)
  video_duration_seconds = int(duration)
  num_frames = video_duration_seconds * 10
  video_frames = pipe(vid_prompt, negative_prompt="low quality", num_inference_steps=25, num_frames=num_frames).frames
  video_path = export_to_video(video_frames)
  return video_path

with gr.Blocks() as demo:
    gr.Markdown("Text to Video using Diffusion-based text-to-video generation model built by ModelScope (NVIDIA GPU driver must be already installed on your system.)")
    with gr.Tab("Text to video"):
        prompt = gr.Textbox(label= 'Text Prompt')
        video_duration = gr.Textbox(label= "Video Duration in Seconds (videos longer than 5 seconds can take a long time to create.)")
        video_output = gr.Video()
        create = gr.Button('Make Video')
        create.click(make_video,[prompt,video_duration], video_output)
demo.queue().launch(debug=True)