Update app.py

app.py

@@ -24,6 +24,10 @@ 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):
@@ -105,98 +109,15 @@ def normalize_and_save_video(input_video_path, output_video_path):
     print(f"NORMALIZE DONE!")
     return output_video_path
 
-def chunkify(video_path, fps, nb_frames):
-    chunks_array = []
-
-    video_capture = cv2.VideoCapture(video_path)
-    chunk_start_frame = 0
-    frames_per_chunk = 12
-
-    while chunk_start_frame < nb_frames:
-        chunk_end_frame = min(chunk_start_frame + frames_per_chunk, nb_frames)
-
-        video_capture.set(cv2.CAP_PROP_POS_FRAMES, chunk_start_frame)
-        success, frame = video_capture.read()
-        if not success:
-            break
-
-        chunk_name = f"chunk_{chunk_start_frame}-{chunk_end_frame}.mp4"
-        chunk_video = cv2.VideoWriter(chunk_name, cv2.VideoWriter_fourcc(*"mp4v"), fps, (frame.shape[1], frame.shape[0]))
-
-        for frame_number in range(chunk_start_frame, chunk_end_frame):
-            video_capture.set(cv2.CAP_PROP_POS_FRAMES, frame_number)
-            success, frame = video_capture.read()
-            if not success:
-                break
-
-            chunk_video.write(frame)
-
-        chunk_video.release()
-        chunks_array.append(chunk_name)
-
-        chunk_start_frame += frames_per_chunk
-
-    video_capture.release()
-    print(f"CHUNKS: {chunks_array}")
-    return chunks_array
-    
-        
-def run_inference_by_chunkify(prompt, video_path, condition, video_length):
-
-    # DOESN'T WORK
-
-    # Get FPS of original video input
-    target_fps = get_video_dimension(video_path)[2] 
-    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)
-
-    # Chunkify the video into 12 frames chunks
-    chunks = chunkify(resized, target_fps, total_frames)
-
-    output_path = 'output/'
-    os.makedirs(output_path, exist_ok=True)
-
-    processed_chunks = []
     
-    for index, chunk_path in enumerate(chunks):
-        if index == 0 :
-            print(f"Chunk #{index}: {chunk_path}")
-            
-            # Check if the file already exists
-            if os.path.exists(os.path.join(output_path, f"{index}.mp4")):
-                # Delete the existing file
-                os.remove(os.path.join(output_path, f"{index}.mp4"))
-    
-        #if video_length > 12:
-        #    command = f"python inference.py --prompt '{prompt}' --condition '{condition}' --video_path '{video_path}' --output_path '{output_path}' --width 512 --height 512 --fps 8 --video_length {video_length} --is_long_video"
-        #else:
-            command = f"python inference.py --prompt '{prompt}' --condition '{condition}' --video_path '{chunk_path}' --output_path '{output_path}' --temp_chunk_path '{index}'  --width 512 --height 512 --fps 8 --video_length {video_length} --is_long_video"
-            subprocess.run(command, shell=True)
-    
-            # Construct the video path
-            video_path_output = os.path.join(output_path, f"{index}.mp4")
-    
-            # Append processed chunk to final array
-            processed_chunks.append(video_path_output)
-        else:
-            print("finished")
-
-    print(f"PROCESSED CHUNKS: {processed_chunks}")
-
-    return "done", processed_chunks[0] 
 
-
-def run_inference(prompt, video_path, condition, video_length):
-
-    
+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
@@ -218,16 +139,21 @@ def run_inference(prompt, video_path, condition, video_length):
 
     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} --video_length {video_length} --is_long_video"
+        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} --video_length {video_length}"
+        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", video_path_output 
+    return "done", 'resized_final.mp4'
 
  
 
@@ -244,10 +170,10 @@ with gr.Blocks(css=css) as demo:
                 #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.Row():
+                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)
+                    seed = gr.Number(label="seed", value=42)
                 submit_btn = gr.Button("Submit")
             with gr.Column():
                 video_res = gr.Video(label="result")
@@ -260,7 +186,8 @@ with gr.Blocks(css=css) as demo:
                      inputs=[prompt,
                              video_path,
                              condition,
-                             video_length
+                             video_length,
+                             seed
                             ],
                     outputs=[status, video_res])