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

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):
    video = cv2.VideoCapture(filepath) 
    frame_count = int(video.get(cv2.CAP_PROP_FRAME_COUNT))
    video.release()

    #LIMITS
    if frame_count > 12 :
        frame_count = 24 # limit to 24 frames to avoid cuDNN errors
    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))
    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 normalize_and_save_video(input_video_path, output_video_path):
    print(f"NORMALIZING ...")
    cap = cv2.VideoCapture(input_video_path)

    # Get video properties
    frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps = cap.get(cv2.CAP_PROP_FPS)

    # Create VideoWriter object to save the normalized video
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # Specify the codec (e.g., 'mp4v', 'XVID', 'MPEG')
    out = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))

    # Iterate through each frame in the video
    for _ in range(frame_count):
        ret, frame = cap.read()
        if not ret:
            break

        # Convert frame to floating point
        frame = frame.astype(np.float32)

        # Normalize pixel values to the range [0, 1]
        frame /= 255.0

        # Convert normalized frame back to 8-bit unsigned integer
        frame = (frame * 255.0).astype(np.uint8)

        # Write the normalized frame to the output video file
        out.write(frame)

    # Release the VideoCapture and VideoWriter objects
    cap.release()
    out.release()

    print(f"NORMALIZE DONE!")
    return output_video_path

    

def run_inference(prompt, video_path, condition, video_length, seed):

    # Get FPS of original video input
    target_fps = get_video_dimension(video_path)[2] 
    if target_fps > 12 :
        print(f"FPS is too high")
        target_fps = 12
    print(f"INPUT FPS: {target_fps}")
    
    # Count total frames according to fps
    total_frames = get_video_dimension(video_path)[3]

    # Resize the video
    resized = resize_video(video_path, 'resized.mp4', 512, 512, target_fps)

    # normalize pixels
    normalized = normalize_and_save_video(resized, 'normalized.mp4')

    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}' --condition '{condition}' --video_path '{normalized}' --output_path '{output_path}' --temp_chunk_path 'result' --width 512 --height 512 --fps {target_fps} --seed {seed} --video_length {video_length} --is_long_video"
    else:
        command = f"python inference.py --prompt '{prompt}' --condition '{condition}' --video_path '{normalized}' --output_path '{output_path}' --temp_chunk_path 'result'  --width 512 --height 512 --fps {target_fps} --seed {seed} --video_length {video_length}"
    subprocess.run(command, shell=True)
    
    # 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)

    print(f"FINISHED !")
    return "done", 'resized_final.mp4'

 

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():
                    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,
                      inputs=[video_path],
                      outputs=[video_length]
                     )
    submit_btn.click(fn=run_inference, 
                     inputs=[prompt,
                             video_path,
                             condition,
                             video_length,
                             seed
                            ],
                    outputs=[status, video_res])

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