import gradio as gr
import cv2
import numpy as np
import torch
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler

pipe = DiffusionPipeline.from_pretrained("cerspense/zeroscope_v2_576w", torch_dtype=torch.float16)
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe.enable_model_cpu_offload()

def export_to_video(video_frames):
    # Convert the nested list to a NumPy array and scale values to 0-255 range
    video_frames = np.array(video_frames)
    video_frames = (video_frames * 255).astype(np.uint8)
    
    # Get the dimensions of the frames
    height, width, channels = video_frames.shape[2:]
    
    # Define the video writer object
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for .mp4 files
    fps = 24  # Frames per second
    video_writer = cv2.VideoWriter('output_video.mp4', fourcc, fps, (width, height))
    
    # Write each frame to the video
    for i in range(video_frames.shape[0]):
        frame = video_frames[i]
        video_writer.write(frame)
    
    # Release the video writer object
    video_writer.release()
    
    print("Video has been created successfully.")
    return 'output_video.mp4'

def infer(prompt):
    negative_prompt = "text, watermark, copyright, blurry, nsfw"
    video_frames = pipe(prompt, negative_prompt=negative_prompt, num_inference_steps=40, height=320, width=576, num_frames=24).frames
    video_path = export_to_video(video_frames)
    print(video_path)
    return video_path

css = """
#col-container {max-width: 510px; margin-left: auto; margin-right: auto;}
a {text-decoration-line: underline; font-weight: 600;}
.animate-spin {
  animation: spin 1s linear infinite;
}

@keyframes spin {
  from {
      transform: rotate(0deg);
  }
  to {
      transform: rotate(360deg);
  }
}

#share-btn-container {
  display: flex; 
  padding-left: 0.5rem !important; 
  padding-right: 0.5rem !important; 
  background-color: #000000; 
  justify-content: center; 
  align-items: center; 
  border-radius: 9999px !important; 
  max-width: 15rem;
  height: 36px;
}

div#share-btn-container > div {
    flex-direction: row;
    background: black;
    align-items: center;
}

#share-btn-container:hover {
  background-color: #060606;
}

#share-btn {
  all: initial; 
  color: #ffffff;
  font-weight: 600; 
  cursor:pointer; 
  font-family: 'IBM Plex Sans', sans-serif; 
  margin-left: 0.5rem !important; 
  padding-top: 0.5rem !important; 
  padding-bottom: 0.5rem !important;
  right:0;
}

#share-btn * {
  all: unset;
}

#share-btn-container div:nth-child(-n+2){
  width: auto !important;
  min-height: 0px !important;
}

#share-btn-container .wrap {
  display: none !important;
}

#share-btn-container.hidden {
  display: none!important;
}
img[src*='#center'] { 
    display: inline-block;
    margin: unset;
}

.footer {
        margin-bottom: 45px;
        margin-top: 10px;
        text-align: center;
        border-bottom: 1px solid #e5e5e5;
    }
    .footer>p {
        font-size: .8rem;
        display: inline-block;
        padding: 0 10px;
        transform: translateY(10px);
        background: white;
    }
    .dark .footer {
        border-color: #303030;
    }
    .dark .footer>p {
        background: #0b0f19;
    }
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            <h1 style="text-align: center;">Zeroscope Text-to-Video</h1>
            <p style="text-align: center;">
            A watermark-free Modelscope-based video model optimized for producing high-quality 16:9 compositions and a smooth video output. <br />
            </p>
                        
            """
        )

        prompt_in = gr.Textbox(label="Prompt", placeholder="Darth Vader is surfing on waves", elem_id="prompt-in")
        #neg_prompt = gr.Textbox(label="Negative prompt", value="text, watermark, copyright, blurry, nsfw", elem_id="neg-prompt-in")
        #inference_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, step=1, value=40, interactive=False)
        submit_btn = gr.Button("Submit")
        video_result = gr.Video(label="Video Output", elem_id="video-output")

        with gr.Row():
            gr.Markdown("""
                [![Duplicate this Space](https://huggingface.co/datasets/huggingface/badges/raw/main/duplicate-this-space-lg.svg#center)](https://huggingface.co/spaces/fffiloni/zeroscope-cloning?duplicate=true)
            """)

        gr.HTML("""
            <div class="footer">
                <p>
                <a href="https://huggingface.co/cerspense/zeroscope_v2_576w" target="_blank">Zeroscope v2 576w model</a> by @cerspense - 
                Demo by 🤗 <a href="https://twitter.com/fffiloni" target="_blank">Sylvain Filoni</a>
                </p>
            </div>
            <div id="may-like-container" style="display: flex;justify-content: center;flex-direction: column;align-items: center;">
                <p style="font-size: 0.8em;margin-bottom: 4px;">You may also like: </p>
                <div id="may-like" style="display:flex; align-items:center; justify-content: center;height:20px;">
                    <svg height="20" width="148" style="margin-left:4px">       
                        <a href="https://huggingface.co/spaces/fffiloni/zeroscope-XL" target="_blank">
                            <image href="https://img.shields.io/badge/🤗 Spaces-Zeroscope XL-blue" src="https://img.shields.io/badge/🤗 Spaces-Image to Music-blue.png" height="20"/>
                        </a>
                    </svg>
                </div>
            </div>
        """)
        
    submit_btn.click(fn=infer,
                    inputs=[prompt_in],
                    outputs=[video_result],
                    api_name="zrscp")

demo.queue(max_size=12).launch(show_api=False)