app.py

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

from huggingface_hub import snapshot_download

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_in_duration(filepath):
    video = cv2.VideoCapture(filepath)
    fps = video.get(cv2.CAP_PROP_FPS)
    frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = frame_count / fps
    width = int(video.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(video.get(cv2.CAP_PROP_FRAME_HEIGHT))
    video.release()
    return gr.update(maximum=frame_count)

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))
    video.release()
    return width, height, fps

def adjust_to_multiple_of_12(number):
    remainder = number % 12
    if remainder != 0:
        adjustment = 12 - remainder
        number += adjustment
    return number
    
def resize_video(input_file):
    # Load the video clip
    clip = VideoFileClip(input_file)
    print(f"WIDTH TARGET: 512")
    # Calculate the aspect ratio
    ratio = 512 / clip.size[0]
    new_height = int(clip.size[1] * ratio)
    new_height_adjusted = adjust_to_multiple_of_12(new_height)
    new_width_adjusted = adjust_to_multiple_of_12(512)
    print(f"OLD H: {new_height} | NEW H: {new_height_adjusted}")
    print(f"OLD W: 512 | NEW W: {new_width_adjusted}")

    # Resize the video clip
    resized_clip = clip.resize(width=new_width_adjusted, height=new_height_adjusted)
    
    # Check if the file already exists
    if os.path.exists('video_resized.mp4'):
        # Delete the existing file
        os.remove('video_resized.mp4')
        
    # Write the resized video to a file
    resized_clip.write_videofile('video_resized.mp4', codec="libx264")

    # Close the video clip
    clip.close()

    
    #final_video_resized = os.path.join(temp_output_path, 'video_resized.mp4')
    test_w, test_h, fps = get_video_dimension('video_resized.mp4')
    print(f"resized clip dims : {test_w}, {test_h}, {fps}")
    return 'video_resized.mp4'
        
def run_inference(prompt, video_path, condition, video_length):

    # Specify the input and output paths
    input_vid = video_path

    # Call the function to resize the video
    resized_video_path = resize_video(input_vid)
    print(f"PATH TO RESIZED: {resized_video_path}")
    width, height, fps = get_video_dimension(resized_video_path)

    print(f"{width} x {height} | {fps}")
    
    output_path = 'output/'
    os.makedirs(output_path, exist_ok=True)

    # Construct the final video path
    video_path_output = os.path.join(output_path, f"{prompt}.mp4")

    # Check if the file already exists
    if os.path.exists(video_path_output):
        # Delete the existing file
        os.remove(video_path_output)

    if video_length > 12:
        command = f"python inference.py --prompt '{prompt}' --condition '{condition}' --video_path '{resized_video_path}' --output_path '{output_path}' --width {width} --height {height} --fps {fps} --video_length {video_length} --is_long_video"
    else:
        command = f"python inference.py --prompt '{prompt}' --condition '{condition}' --video_path '{resized_video_path}' --output_path '{output_path}' --width {width} --height {height} --fps {fps} --video_length {video_length}"
    subprocess.run(command, shell=True)

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

        

    return "done", video_path_output 

def run_inference_chunks(prompt, video_path, condition, video_length):

    # Specify the input and output paths
    input_vid = video_path
    resized_vid = 'resized.mp4'

    # Call the function to resize the video
    video_path = resize_video(input_vid, resized_vid, width=512)
    width, height, fps = get_video_dimension(video_path)

    print(f"{width} x {height} | {fps}")

    # Split the video into chunks mp4 of 12 frames at video fps
    # Store chunks as mp4 paths in an array

    # For each mp4 chunks in chunks arrays, run command
    # store video result in processed chunks array
    
    output_path = 'output/'
    os.makedirs(output_path, exist_ok=True)

    # Construct the final video path
    video_path_output = os.path.join(output_path, f"{prompt}.mp4")

    # Check if the file already exists
    if os.path.exists(video_path_output):
        # Delete the existing file
        os.remove(video_path_output)

    if video_length > 12:
        command = f"python inference.py --prompt '{prompt}' --condition '{condition}' --video_path '{video_path}' --output_path '{output_path}' --width {width} --height {height} --fps {fps} --video_length {video_length} --is_long_video"
    else:
        command = f"python inference.py --prompt '{prompt}' --condition '{condition}' --video_path '{video_path}' --output_path '{output_path}' --width {width} --height {height} --fps {fps} --video_length {video_length}"
    subprocess.run(command, shell=True)

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

        

    return "done", video_path_output 

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_path = gr.Video(source="upload", type="filepath")
                prompt = gr.Textbox(label="prompt")
                with gr.Row():
                    condition = gr.Dropdown(label="Condition", choices=["depth", "canny", "pose"], value="depth")
                    video_length = gr.Slider(label="Video length", info="How many frames do you want to process ?", minimum=1, maximum=12, step=1, value=2)
                #seed = gr.Number(label="seed", value=42)
                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_in_duration,
                      inputs=[video_path],
                      outputs=[video_length]
                     ).then(fn=resize_video, inputs=[video_path], outputs=[video_path])
    submit_btn.click(fn=run_inference, 
                     inputs=[prompt,
                             video_path,
                             condition,
                             video_length
                            ],
                    outputs=[status, video_res])

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