ZeroGPU (#3)

- ZeroGPU (8876d0ea877e4366551925f1061473a83ef89f6e)


Co-authored-by: Charles Bensimon <cbensimon@users.noreply.huggingface.co>

app.py

@@ -46,7 +46,6 @@ default_bas_frame_near = 1
 default_bas_frame_far = 1
 
 
-@spaces.GPU
 def process_image(
     pipe,
     path_input,
@@ -89,7 +88,6 @@ def process_image(
     )
 
 
-@spaces.GPU
 def process_video(
     pipe,
     path_input,
@@ -187,7 +185,6 @@ def process_video(
     )
 
 
-@spaces.GPU
 def process_bas(
     pipe,
     path_input,
@@ -282,9 +279,9 @@ def process_bas(
 
 
 def run_demo_server(pipe):
-    process_pipe_image = functools.partial(process_image, pipe)
-    process_pipe_video = functools.partial(process_video, pipe)
-    process_pipe_bas = functools.partial(process_bas, pipe)
+    process_pipe_image = spaces.GPU(functools.partial(process_image, pipe))
+    process_pipe_video = spaces.GPU(functools.partial(process_video, pipe))
+    process_pipe_bas = spaces.GPU(functools.partial(process_bas, pipe))
     os.environ["GRADIO_ALLOW_FLAGGING"] = "never"
 
     gradio_theme = gr.themes.Default()
@@ -792,11 +789,6 @@ def run_demo_server(pipe):
         )
 
 
-def prefetch_hf_cache(pipe):
-    process_image(pipe, "files/image/bee.jpg", 1, 1, 64)
-    shutil.rmtree("files/image/bee_output")
-
-
 def main():
     CHECKPOINT = "prs-eth/marigold-v1-0"
     CHECKPOINT_UNET_LCM = "prs-eth/marigold-lcm-v1-0"
@@ -821,7 +813,6 @@ def main():
         pass  # run without xformers
 
     pipe = pipe.to(device)
-    prefetch_hf_cache(pipe)
     run_demo_server(pipe)