nivere/ControlNet-Video
1
1from __future__ import annotations2import gradio as gr3import os4import cv25import numpy as np6from PIL import Image7from moviepy.editor import *8from share_btn import community_icon_html, loading_icon_html, share_js9 10import pathlib11import shlex12import subprocess13 14if os.getenv('SYSTEM') == 'spaces':15 with open('patch') as f:16 subprocess.run(shlex.split('patch -p1'), stdin=f, cwd='ControlNet')17 18base_url = 'https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/'19 20names = [21 'body_pose_model.pth',22 'dpt_hybrid-midas-501f0c75.pt',23 'hand_pose_model.pth',24 'mlsd_large_512_fp32.pth',25 'mlsd_tiny_512_fp32.pth',26 'network-bsds500.pth',27 'upernet_global_small.pth',28]29 30for name in names:31 command = f'wget https://huggingface.co/lllyasviel/ControlNet/resolve/main/annotator/ckpts/{name} -O {name}'32 out_path = pathlib.Path(f'ControlNet/annotator/ckpts/{name}')33 if out_path.exists():34 continue35 subprocess.run(shlex.split(command), cwd='ControlNet/annotator/ckpts/')36 37from model import Model38model = Model()39 40 41def controlnet(i, prompt, control_task, seed_in, ddim_steps, scale):42 img= Image.open(i)43 np_img = np.array(img)44 45 a_prompt = "best quality, extremely detailed"46 n_prompt = "longbody, lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality"47 num_samples = 148 image_resolution = 51249 detect_resolution = 51250 eta = 0.051 low_threshold = 10052 high_threshold = 20053 value_threshold = 0.154 distance_threshold = 0.155 bg_threshold = 0.456 57 if control_task == 'Canny':58 result = model.process_canny(np_img, prompt, a_prompt, n_prompt, num_samples,59 image_resolution, ddim_steps, scale, seed_in, eta, low_threshold, high_threshold)60 elif control_task == 'Depth':61 result = model.process_depth(np_img, prompt, a_prompt, n_prompt, num_samples,62 image_resolution, detect_resolution, ddim_steps, scale, seed_in, eta)63 elif control_task == 'Hed':64 result = model.process_hed(np_img, prompt, a_prompt, n_prompt, num_samples,65 image_resolution, detect_resolution, ddim_steps, scale, seed_in, eta)66 elif control_task == 'Hough':67 result = model.process_hough(np_img, prompt, a_prompt, n_prompt, num_samples,68 image_resolution, detect_resolution, ddim_steps, scale, seed_in, eta, value_threshold,69 distance_threshold)70 elif control_task == 'Normal':71 result = model.process_normal(np_img, prompt, a_prompt, n_prompt, num_samples,72 image_resolution, detect_resolution, ddim_steps, scale, seed_in, eta, bg_threshold)73 elif control_task == 'Pose':74 result = model.process_pose(np_img, prompt, a_prompt, n_prompt, num_samples,75 image_resolution, detect_resolution, ddim_steps, scale, seed_in, eta)76 elif control_task == 'Scribble':77 result = model.process_scribble(np_img, prompt, a_prompt, n_prompt, num_samples,78 image_resolution, ddim_steps, scale, seed_in, eta)79 elif control_task == 'Seg':80 result = model.process_seg(np_img, prompt, a_prompt, n_prompt, num_samples,81 image_resolution, detect_resolution, ddim_steps, scale, seed_in, eta)82 83 #print(result[0])84 processor_im = Image.fromarray(result[0])85 processor_im.save("process_" + control_task + "_" + str(i) + ".jpeg")86 im = Image.fromarray(result[1])87 im.save("your_file" + str(i) + ".jpeg")88 return "your_file" + str(i) + ".jpeg", "process_" + control_task + "_" + str(i) + ".jpeg"89 90 91def get_frames(video_in):92 frames = []93 #resize the video94 clip = VideoFileClip(video_in)95 96 #check fps97 if clip.fps > 30:98 print("vide rate is over 30, resetting to 30")99 clip_resized = clip.resize(height=512)100 clip_resized.write_videofile("video_resized.mp4", fps=30)101 else:102 print("video rate is OK")103 clip_resized = clip.resize(height=512)104 clip_resized.write_videofile("video_resized.mp4", fps=clip.fps)105 106 print("video resized to 512 height")107 108 # Opens the Video file with CV2109 cap= cv2.VideoCapture("video_resized.mp4")110 111 fps = cap.get(cv2.CAP_PROP_FPS)112 print("video fps: " + str(fps))113 i=0114 while(cap.isOpened()):115 ret, frame = cap.read()116 if ret == False:117 break118 cv2.imwrite('kang'+str(i)+'.jpg',frame)119 frames.append('kang'+str(i)+'.jpg')120 i+=1121 122 cap.release()123 cv2.destroyAllWindows()124 print("broke the video into frames")125 126 return frames, fps127 128 129def convert(gif):130 if gif != None:131 clip = VideoFileClip(gif.name)132 clip.write_videofile("my_gif_video.mp4")133 return "my_gif_video.mp4"134 else:135 pass136 137 138def create_video(frames, fps, type):139 print("building video result")140 clip = ImageSequenceClip(frames, fps=fps)141 clip.write_videofile(type + "_result.mp4", fps=fps)142 143 return type + "_result.mp4"144 145 146def infer(prompt,video_in, control_task, seed_in, trim_value, ddim_steps, scale, gif_import):147 print(f"""148 ———————————————149 {prompt}150 ———————————————""")151 152 # 1. break video into frames and get FPS153 break_vid = get_frames(video_in)154 frames_list= break_vid[0]155 fps = break_vid[1]156 n_frame = int(trim_value*fps)157 158 if n_frame >= len(frames_list):159 print("video is shorter than the cut value")160 n_frame = len(frames_list)161 162 # 2. prepare frames result arrays163 processor_result_frames = []164 result_frames = []165 print("set stop frames to: " + str(n_frame))166 167 for i in frames_list[0:int(n_frame)]:168 controlnet_img = controlnet(i, prompt,control_task, seed_in, ddim_steps, scale)169 #images = controlnet_img[0]170 #rgb_im = images[0].convert("RGB")171 172 # exporting the image173 #rgb_im.save(f"result_img-{i}.jpg")174 processor_result_frames.append(controlnet_img[1])175 result_frames.append(controlnet_img[0])176 print("frame " + i + "/" + str(n_frame) + ": done;")177 178 processor_vid = create_video(processor_result_frames, fps, "processor")179 final_vid = create_video(result_frames, fps, "final")180 181 files = [processor_vid, final_vid]182 if gif_import != None:183 final_gif = VideoFileClip(final_vid)184 final_gif.write_gif("final_result.gif")185 final_gif = "final_result.gif"186 187 files.append(final_gif)188 print("finished !")189 190 return final_vid, gr.Accordion.update(visible=True), gr.Video.update(value=processor_vid, visible=True), gr.File.update(value=files, visible=True), gr.Group.update(visible=True)191 192 193def clean():194 return gr.Accordion.update(visible=False),gr.Video.update(value=None, visible=False), gr.Video.update(value=None), gr.File.update(value=None, visible=False), gr.Group.update(visible=False)195 196title = """197 <div style="text-align: center; max-width: 700px; margin: 0 auto;">198 <div199 style="200 display: inline-flex;201 align-items: center;202 gap: 0.8rem;203 font-size: 1.75rem;204 "205 >206 <h1 style="font-weight: 900; margin-bottom: 7px; margin-top: 5px;">207 ControlNet Video208 </h1>209 </div>210 <p style="margin-bottom: 10px; font-size: 94%">211 Apply ControlNet to a video 212 </p>213 </div>214"""215 216article = """217 218 <div class="footer">219 <p>220 Follow <a href="https://twitter.com/fffiloni" target="_blank">Sylvain Filoni</a> for future updates 🤗221 </p>222 </div>223 <div id="may-like-container" style="display: flex;justify-content: center;flex-direction: column;align-items: center;margin-bottom: 30px;">224 <p>You may also like: </p>225 <div id="may-like-content" style="display:flex;flex-wrap: wrap;align-items:center;height:20px;">226 227 <svg height="20" width="148" style="margin-left:4px;margin-bottom: 6px;"> 228 <a href="https://huggingface.co/spaces/fffiloni/Pix2Pix-Video" target="_blank">229 <image href="https://img.shields.io/badge/🤗 Spaces-Pix2Pix_Video-blue" src="https://img.shields.io/badge/🤗 Spaces-Pix2Pix_Video-blue.png" height="20"/>230 </a>231 </svg>232 233 </div>234 235 </div>236 237"""238 239with gr.Blocks(css='style.css') as demo:240 with gr.Column(elem_id="col-container"):241 gr.HTML(title)242 with gr.Row():243 with gr.Column():244 video_inp = gr.Video(label="Video source", source="upload", type="filepath", elem_id="input-vid")245 video_out = gr.Video(label="ControlNet video result", elem_id="video-output")246 with gr.Accordion("Detailed results", visible=False) as detailed_result:247 prep_video_out = gr.Video(label="Preprocessor video result", visible=False, elem_id="prep-video-output")248 files = gr.File(label="Files can be downloaded ;)", visible=False)249 with gr.Group(elem_id="share-btn-container", visible=False) as share_group:250 community_icon = gr.HTML(community_icon_html)251 loading_icon = gr.HTML(loading_icon_html)252 share_button = gr.Button("Share to community", elem_id="share-btn")253 with gr.Column():254 #status = gr.Textbox()255 256 prompt = gr.Textbox(label="Prompt", placeholder="enter prompt", show_label=True, elem_id="prompt-in")257 with gr.Row():258 control_task = gr.Dropdown(label="Control Task", choices=["Canny", "Depth", "Hed", "Hough", "Normal", "Pose", "Scribble", "Seg"], value="Pose", multiselect=False, elem_id="controltask-in")259 seed_inp = gr.Slider(label="Seed", minimum=0, maximum=2147483647, step=1, value=123456, elem_id="seed-in")260 with gr.Row():261 262 trim_in = gr.Slider(label="Cut video at (s)", minimun=1, maximum=5, step=1, value=1)263 with gr.Accordion("Advanced Options", open=False):264 265 ddim_steps = gr.Slider(label='Steps',266 minimum=1,267 maximum=100,268 value=20,269 step=1)270 scale = gr.Slider(label='Guidance Scale',271 minimum=0.1,272 maximum=30.0,273 value=9.0,274 step=0.1)275 276 gif_import = gr.File(label="import a GIF instead", file_types=['.gif'])277 gif_import.change(convert, gif_import, video_inp, queue=False)278 279 submit_btn = gr.Button("Generate ControlNet video")280 281 gr.HTML("""282 <a style="display:inline-block" href="https://huggingface.co/spaces/fffiloni/Pix2Pix-Video?duplicate=true"><img src="https://img.shields.io/badge/-Duplicate%20Space-blue?labelColor=white&style=flat&logo=data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAABAAAAAQCAYAAAAf8/9hAAAAAXNSR0IArs4c6QAAAP5JREFUOE+lk7FqAkEURY+ltunEgFXS2sZGIbXfEPdLlnxJyDdYB62sbbUKpLbVNhyYFzbrrA74YJlh9r079973psed0cvUD4A+4HoCjsA85X0Dfn/RBLBgBDxnQPfAEJgBY+A9gALA4tcbamSzS4xq4FOQAJgCDwV2CPKV8tZAJcAjMMkUe1vX+U+SMhfAJEHasQIWmXNN3abzDwHUrgcRGmYcgKe0bxrblHEB4E/pndMazNpSZGcsZdBlYJcEL9Afo75molJyM2FxmPgmgPqlWNLGfwZGG6UiyEvLzHYDmoPkDDiNm9JR9uboiONcBXrpY1qmgs21x1QwyZcpvxt9NS09PlsPAAAAAElFTkSuQmCC&logoWidth=14" alt="Duplicate Space"></a> 283 work with longer videos / skip the queue: 284 """, elem_id="duplicate-container")285 286 inputs = [prompt,video_inp,control_task, seed_inp, trim_in, ddim_steps, scale, gif_import]287 outputs = [video_out, detailed_result, prep_video_out, files, share_group]288 #outputs = [status]289 290 291 gr.HTML(article)292 293 submit_btn.click(clean, inputs=[], outputs=[detailed_result, prep_video_out, video_out, files, share_group], queue=False)294 submit_btn.click(infer, inputs, outputs)295 share_button.click(None, [], [], _js=share_js)296 297 298 299demo.queue(max_size=12).launch()