Update app.py

app.py

@@ -1,4 +1,9 @@
 from PIL import Image
+import os
+import cv2
+import numpy as np
+from PIL import Image
+from moviepy.editor import *
 import gradio as gr
 from diffusers import StableDiffusionControlNetPipeline, ControlNetModel, UniPCMultistepScheduler
 import torch
@@ -20,7 +25,51 @@ pipe.enable_xformers_memory_efficient_attention()
 pipe.enable_model_cpu_offload()
 pipe.enable_attention_slicing()
 
-def infer(
+def get_frames(video_in):
+    frames = []
+    #resize the video
+    clip = VideoFileClip(video_in)
+    
+    #check fps
+    if clip.fps > 30:
+        print("vide rate is over 30, resetting to 30")
+        clip_resized = clip.resize(height=512)
+        clip_resized.write_videofile("video_resized.mp4", fps=30)
+    else:
+        print("video rate is OK")
+        clip_resized = clip.resize(height=512)
+        clip_resized.write_videofile("video_resized.mp4", fps=clip.fps)
+    
+    print("video resized to 512 height")
+    
+    # Opens the Video file with CV2
+    cap= cv2.VideoCapture("video_resized.mp4")
+    
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    print("video fps: " + str(fps))
+    i=0
+    while(cap.isOpened()):
+        ret, frame = cap.read()
+        if ret == False:
+            break
+        cv2.imwrite('kang'+str(i)+'.jpg',frame)
+        frames.append('kang'+str(i)+'.jpg')
+        i+=1
+    
+    cap.release()
+    cv2.destroyAllWindows()
+    print("broke the video into frames")
+    
+    return frames, fps
+
+def create_video(frames, fps):
+    print("building video result")
+    clip = ImageSequenceClip(frames, fps=fps)
+    clip.write_videofile("_result.mp4", fps=fps)
+    
+    return "_result.mp4"
+
+def process_brightness(
         prompt,
         negative_prompt,
         conditioning_image,
@@ -57,12 +106,63 @@ def infer(
 
     return output_image
 
+def infer(video_in, prompt,
+                    negative_prompt,
+                    conditioning_image,
+                    num_inference_steps=30,
+                    size=768,
+                    guidance_scale=7.0,
+                    seed=1234
+         ):
+    
+    # 1. break video into frames and get FPS
+    break_vid = get_frames(video_in)
+    frames_list= break_vid[0]
+    fps = break_vid[1]
+    #n_frame = int(trim_value*fps)
+    n_frame = len(frames_list)
+    
+    if n_frame >= len(frames_list):
+        print("video is shorter than the cut value")
+        n_frame = len(frames_list)
+    
+    # 2. prepare frames result arrays
+    result_frames = []
+    print("set stop frames to: " + str(n_frame))
+    
+    for i, image in enumerate(frames_list[0:int(n_frame)]):
+
+        image = Image.open(image).convert("RGB")
+        image = np.array(image)
+        output_frame = process_brightness(image,
+                            prompt,
+                            negative_prompt,
+                            conditioning_image,
+                            num_inference_steps=30,
+                            size=768,
+                            guidance_scale=7.0,
+                            seed=1234
+                           )
+        print(output_frame)
+        
+        image = Image.open(output_frame)
+        #image = Image.fromarray(output_frame[0])
+        image.save("_frame_" + str(i) + ".jpeg")
+        result_frames.append("_frame_" + str(i) + ".jpeg")
+        print("frame " + str(i) + "/" + str(n_frame) + ": done;")
+
+    
+    final_vid = create_video(result_frames, fps)
+
+    
+    return final_vid
+
 with gr.Blocks() as demo:
     gr.Markdown(
         """
-    # ControlNet on Brightness
+    # ControlNet on Brightness • Video
 
-    This is a demo on ControlNet based on brightness.
+    This is a demo on ControlNet based on brightness for video.
     """)
 
     with gr.Row():
@@ -73,8 +173,10 @@ with gr.Blocks() as demo:
             negative_prompt = gr.Textbox(
                 label="Negative Prompt",
             )
-            conditioning_image = gr.Image(
-                label="Conditioning Image",
+            video_in = gr.Video(
+                label="Conditioning Video",
+                source="upload",
+                type="filepath"     
             )
             with gr.Accordion('Advanced options', open=False):
                 with gr.Row():
@@ -116,7 +218,7 @@ with gr.Blocks() as demo:
     submit_btn.click(
         fn=infer,
         inputs=[
-            prompt, negative_prompt, conditioning_image, num_inference_steps, size, guidance_scale, seed
+            video_in, prompt, negative_prompt, conditioning_image, num_inference_steps, size, guidance_scale, seed
         ],
         outputs=output
     )
@@ -134,7 +236,7 @@ with gr.Blocks() as demo:
         ],
         outputs=output,
         fn=infer,
-        cache_examples=True,
+        cache_examples=False,
     )
     gr.Markdown(
         """