zero gpu

app.py

@@ -38,6 +38,7 @@ from annotator.midas import MidasDetector
 import huggingface_hub
 import shutil
 import os
+import spaces
 
 huggingface_hub.hf_hub_download('SingleZombie/FRESCO',
                                 'boxer-punching-towards-camera.mp4',
@@ -170,7 +171,7 @@ def cfg_to_input(filename):
         1.0, cfg['controlnet_type'], 50, 100,
         cfg['num_inference_steps'], 7.5, a_prompt, n_prompt,
         frame_count, cfg['batch_size'], cfg['mininterv'], cfg['maxinterv'],
-        use_constraints, True, True, 4,
+        use_constraints, True,
         1, 1, 1, 1
     ]
     return args
@@ -232,21 +233,17 @@ class GlobalState:
             param.requires_grad = False
 
 
-@torch.no_grad()
-def process(*args):
-    keypath = process1(*args)
-    fullpath = process2(*args)
-    return keypath, fullpath
 
 
 @torch.no_grad()
+@spaces.GPU(duration=240)
 def process1(input_path, prompt, sd_model, seed, image_resolution, control_strength,
              x0_strength, control_type, low_threshold, high_threshold,
              ddpm_steps, scale, a_prompt, n_prompt,
              frame_count, batch_size, mininterv, maxinterv,
              use_constraints, bg_smooth, use_poisson, max_process,
              b1, b2, s1, s2):
-    global global_state
+    global_state = GlobalState()
     global_state.update_controlnet_model(control_type)
     global_state.update_sd_model(sd_model)
     apply_freeu(global_state.pipe, b1=b1, b2=b2, s1=s1, s2=s2)
@@ -392,7 +389,7 @@ def process1(input_path, prompt, sd_model, seed, image_resolution, control_stren
             gc.collect()
             if device == 'cuda':
                 torch.cuda.empty_cache()
-            break
+            break  
 
     writer = imageio.get_writer(os.path.join(save_path, 'key.mp4'), fps=fps)
     file_list = sorted(os.listdir(os.path.join(save_path, 'keys')))
@@ -404,43 +401,9 @@ def process1(input_path, prompt, sd_model, seed, image_resolution, control_stren
         writer.append_data(curImg)
     writer.close()
 
-    global_state.processing_state = ProcessingState.KEY_IMGS
     return os.path.join(save_path, 'key.mp4')
 
 
-@torch.no_grad()
-def process2(input_path, prompt, sd_model, seed, image_resolution, control_strength,
-             x0_strength, control_type, low_threshold, high_threshold,
-             ddpm_steps, scale, a_prompt, n_prompt,
-             frame_count, batch_size, mininterv, maxinterv,
-             use_constraints, bg_smooth, use_poisson, max_process,
-             b1, b2, s1, s2):
-
-    global global_state
-    if global_state.processing_state != ProcessingState.KEY_IMGS:
-        raise gr.Error('Please generate key images before propagation')
-
-    # reset blend dir
-    filename = os.path.splitext(os.path.basename(input_path))[0]
-    blend_dir = os.path.join('output', filename)
-    os.makedirs(blend_dir, exist_ok=True)
-
-    video_cap = cv2.VideoCapture(input_path)
-    fps = int(video_cap.get(cv2.CAP_PROP_FPS))
-    o_video = os.path.join(blend_dir, 'blend.mp4')
-    key_ind = io.StringIO()
-    for k in global_state.keys:
-        print('%d' % (k), end=' ', file=key_ind)
-    ps = '-ps' if use_poisson else ''
-    cmd = (
-        f'python video_blend.py {blend_dir} --key keys '
-        f'--key_ind {key_ind.getvalue()} --output {o_video} --fps {fps} '
-        f'--n_proc {max_process} {ps}')
-    print(cmd)
-    os.system(cmd)
-    return o_video
-
-
 config_dir = 'config'
 filenames = os.listdir(config_dir)
 config_list = []
@@ -448,7 +411,6 @@ for filename in filenames:
     if filename.endswith('yaml'):
         config_list.append(f'{config_dir}/{filename}')
 
-global_state = GlobalState()
 block = gr.Blocks().queue()
 with block:
     with gr.Row():
@@ -456,7 +418,7 @@ with block:
     with gr.Row():
         with gr.Column():
             input_path = gr.Video(label='Input Video',
-                                  source='upload',
+                                  sources='upload',
                                   format='mp4',
                                   visible=True)
             prompt = gr.Textbox(label='Prompt')
@@ -472,10 +434,7 @@ with block:
                              step=1,
                              value=0,
                              randomize=True)
-            run_button = gr.Button(value='Run All')
-            with gr.Row():
-                run_button1 = gr.Button(value='Run Key Frames')
-                run_button2 = gr.Button(value='Run Propagation (Ebsynth)')
+            run_button1 = gr.Button(value='Run Key Frames')
             with gr.Accordion('Advanced options for single frame processing',
                               open=False):
                 image_resolution = gr.Slider(label='Frame resolution',
@@ -595,20 +554,6 @@ with block:
                     value=True,
                     info='Select to smooth background')
 
-            with gr.Accordion(
-                    'Advanced options for the full video translation',
-                    open=False):
-                use_poisson = gr.Checkbox(
-                    label='Gradient blending',
-                    value=True,
-                    info=('Blend the output video in gradient, to reduce'
-                          ' ghosting artifacts (but may increase flickers)'))
-                max_process = gr.Slider(label='Number of parallel processes',
-                                        minimum=1,
-                                        maximum=16,
-                                        value=4,
-                                        step=1)
-
             with gr.Accordion('Example configs', open=True):
 
                 example_list = [cfg_to_input(x) for x in config_list]
@@ -618,7 +563,7 @@ with block:
                     x0_strength, control_type, low_threshold, high_threshold,
                     ddpm_steps, scale, a_prompt, n_prompt,
                     frame_count, batch_size, mininterv, maxinterv,
-                    use_constraints[0], bg_smooth, use_poisson, max_process,
+                    use_constraints[0], bg_smooth,
                     b1, b2, s1, s2
                 ]
 
@@ -631,9 +576,6 @@ with block:
             result_keyframe = gr.Video(label='Output key frame video',
                                        format='mp4',
                                        interactive=False)
-            result_video = gr.Video(label='Output full video',
-                                    format='mp4',
-                                    interactive=False)
 
     def input_changed(path):
         if path is None:
@@ -650,13 +592,13 @@ with block:
         max_interv_c = min(20, frame_count)
         min_interv_r = frame_count
         max_interv_r = frame_count
-        return (gr.Slider.update(minimum=min_interv_l,
+        return (gr.Slider.change(minimum=min_interv_l,
                                  value=min_interv_c,
                                  maximum=min_interv_r),
-                gr.Slider.update(minimum=max_interv_l,
+                gr.Slider.change(minimum=max_interv_l,
                                  value=max_interv_c,
                                  maximum=max_interv_r),
-                gr.Slider.update(minimum=8,
+                gr.Slider.change(minimum=8,
                                  value=frame_count,
                                  maximum=frame_count),
                 )
@@ -666,10 +608,6 @@ with block:
     input_path.upload(input_changed, input_path, [
                       mininterv, maxinterv, frame_count])
 
-    run_button.click(fn=process,
-                     inputs=ips,
-                     outputs=[result_keyframe, result_video])
     run_button1.click(fn=process1, inputs=ips, outputs=[result_keyframe])
-    run_button2.click(fn=process2, inputs=ips, outputs=[result_video])
 
 block.launch()

requirements.txt

@@ -8,4 +8,4 @@ av
 basicsr==1.4.2
 numba==0.57.0
 imageio-ffmpeg
-gradio==3.44.4
+gradio==4.27.0