app.py

import gradio as gr
import os
import subprocess
import cv2
import numpy as np
from moviepy.editor import VideoFileClip, concatenate_videoclips
import math

from huggingface_hub import snapshot_download

os.environ['CUDA_LAUNCH_BLOCKING'] = '1'


model_ids = [
    'runwayml/stable-diffusion-v1-5',
    'lllyasviel/sd-controlnet-depth', 
    'lllyasviel/sd-controlnet-canny', 
    'lllyasviel/sd-controlnet-openpose',
]
for model_id in model_ids:
    model_name = model_id.split('/')[-1]
    snapshot_download(model_id, local_dir=f'checkpoints/{model_name}')



def get_frame_count(filepath):
    if filepath is not None:
        video = cv2.VideoCapture(filepath) 
        frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    
        video.release()

        LIMITS
        if frame_count > 36 :
            frame_count = 36 # limit to 236 frames to avoid cuDNN errors

        return gr.update(maximum=frame_count)

    else:
        return gr.update(value=1, maximum=12 )

def get_video_dimension(filepath):
    video = cv2.VideoCapture(filepath)
    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = int(video.get(cv2.CAP_PROP_FPS))
    frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    video.release()
    return width, height, fps, frame_count

def resize_video(input_vid, output_vid, width, height, fps):
    print(f"RESIZING ...")
    # Open the input video file
    video = cv2.VideoCapture(input_vid)

    # Get the original video's width and height
    original_width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
    original_height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))

    # Create a VideoWriter object to write the resized video
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Codec for the output video
    output_video = cv2.VideoWriter(output_vid, fourcc, fps, (width, height))

    while True:
        # Read a frame from the input video
        ret, frame = video.read()
        if not ret:
            break

        # Resize the frame to the desired dimensions
        resized_frame = cv2.resize(frame, (width, height))

        # Write the resized frame to the output video file
        output_video.write(resized_frame)

    # Release the video objects
    video.release()
    output_video.release()
    print(f"RESIZE VIDEO DONE!")
    return output_vid

def make_nearest_multiple_of_32(number):
    remainder = number % 32
    if remainder <= 16:
        number -= remainder
    else:
        number += 32 - remainder
    return number 

def change_video_fps(input_path):
    print(f"CHANGING FIANL OUTPUT FPS")
    cap = cv2.VideoCapture(input_path)
    # Check if the final file already exists
    if os.path.exists('output_video.mp4'):
        # Delete the existing file
        os.remove('output_video.mp4')
    output_path = 'output_video.mp4'
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    output_fps = 12
    output_size = (int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)), int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)))
    out = cv2.VideoWriter(output_path, fourcc, output_fps, output_size)
    
    frame_count = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if not ret:
            break
        
        # Write the current frame to the output video multiple times to increase the frame rate
        for _ in range(output_fps // 8):
            out.write(frame)
        
        frame_count += 1
        print(f'Processed frame {frame_count}')

    cap.release()
    out.release()
    cv2.destroyAllWindows()

    return 'output_video.mp4'

def run_inference(prompt, video_path, condition, video_length, seed, steps):
    
    seed = math.floor(seed)
    o_width = get_video_dimension(video_path)[0]
    o_height = get_video_dimension(video_path)[1]

    # Prepare dimensions
    if o_width > 512 :
        # Calculate the new height while maintaining the aspect ratio
        n_height = int(o_height / o_width * 512)
        n_width = 512

    # Make sure new dimensions are multipe of 32
    r_width = make_nearest_multiple_of_32(n_width)
    r_height = make_nearest_multiple_of_32(n_height)
    print(f"multiple of 32 sizes : {r_width}x{r_height}")

    # Get FPS of original video input
    original_fps = get_video_dimension(video_path)[2] 
    if original_fps > 12 :
        print(f"FPS is too high: {original_fps}")          
        target_fps = 12
    print(f"NEW INPUT FPS: {target_fps}, NEW LENGTH: {video_length}")
    
    # Check if the resized file already exists
    if os.path.exists('resized.mp4'):
        # Delete the existing file
        os.remove('resized.mp4')
    
    resized = resize_video(video_path, 'resized.mp4', r_width, r_height, target_fps)
    resized_video_fcount = get_video_dimension(resized)[3]
    print(f"RESIZED VIDEO FRAME COUNT: {resized_video_fcount}")

    # Make sure new total frame count is enough to handle chosen video length
    # if video_length > resized_video_fcount :
    # video_length = int((target_fps * video_length) / original_fps)
    
    output_path = 'output/'
    os.makedirs(output_path, exist_ok=True)
            
    # Check if the file already exists
    if os.path.exists(os.path.join(output_path, f"result.mp4")):
        # Delete the existing file
        os.remove(os.path.join(output_path, f"result.mp4"))

    print(f"RUNNING INFERENCE ...")
    if video_length > 12:
        command = f"python inference.py --prompt '{prompt}' --inference_steps {steps} --condition '{condition}' --video_path '{resized}' --output_path '{output_path}' --temp_chunk_path 'result' --width {r_width} --height {r_height} --fps {target_fps} --seed {seed} --video_length {video_length} --smoother_steps 19 20 --is_long_video"
    else:
        command = f"python inference.py --prompt '{prompt}' --inference_steps {steps} --condition '{condition}' --video_path '{resized}' --output_path '{output_path}' --temp_chunk_path 'result'  --width {r_width} --height {r_height} --fps {target_fps} --seed {seed} --video_length {video_length} --smoother_steps 19 20"
    
    try:
        subprocess.run(command, shell=True)
    except cuda.Error as e:
        return f"CUDA Error: {e}", None
    except RuntimeError as e:
        return f"Runtime Error: {e}", None

    # Construct the video path
    video_path_output = os.path.join(output_path, f"result.mp4")

    # Resize to original video input size
    #o_width = get_video_dimension(video_path)[0]
    #o_height = get_video_dimension(video_path)[1]
    #resize_video(video_path_output, 'resized_final.mp4', o_width, o_height, target_fps)

    # Check generated video FPS
    gen_fps = get_video_dimension(video_path_output)[2] 
    print(f"GEN VIDEO FPS: {gen_fps}")
    final = change_video_fps(video_path_output)
    print(f"FINISHED !")
    return "done", final
 

css="""
#col-container {max-width: 810px; margin-left: auto; margin-right: auto;}
"""
with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("""
            <h1 style="text-align: center;">ControlVideo</h1>
        """)
        with gr.Row():
            with gr.Column():
                #video_in = gr.Video(source="upload", type="filepath", visible=True)
                video_path = gr.Video(source="upload", type="filepath", visible=True)
                prompt = gr.Textbox(label="prompt")
                with gr.Column():
                    video_length = gr.Slider(label="Video length", info="How many frames do you want to process ? For demo purpose, max is set to 24", minimum=1, maximum=12, step=1, value=2)
                    with gr.Row():
                        condition = gr.Dropdown(label="Condition", choices=["depth", "canny", "pose"], value="depth")
                        seed = gr.Number(label="seed", value=42)
                    inference_steps = gr.Slider(label="Inference steps", minimum=25, maximum=50, step=1, value=25)
                submit_btn = gr.Button("Submit")
            with gr.Column():
                video_res = gr.Video(label="result")
                status = gr.Textbox(label="result")
    video_path.change(fn=get_frame_count,
                      inputs=[video_path],
                      outputs=[video_length],
                      queue=False
                     )
    submit_btn.click(fn=run_inference, 
                     inputs=[prompt,
                             video_path,
                             condition,
                             video_length,
                             seed,
                             inference_steps
                            ],
                    outputs=[status, video_res])

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