ZeroGPU

- By wrapping GPU functions like this the pipeline doesn't have to be transferred from the main process to the GPU worker and GPU coldstart should be much faster (measured between 3s and 4s on this Space)
- prefetch_hf_cache actually calls the pipe, which can't be done outside of decorated functions on ZeroGPU.
When no coldstart happens, execution time does not seem to change, with or without prefetch_hf_cache

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)