Update for ZeroGPU

app_image_to_3d.py

@@ -5,6 +5,8 @@ import shlex
 import subprocess
 
 import gradio as gr
+import PIL.Image
+import spaces
 
 from model import Model
 from settings import CACHE_EXAMPLES, MAX_SEED
@@ -20,9 +22,14 @@ def create_demo(model: Model) -> gr.Blocks:
         )
     examples = ["corgi.png"]
 
+    @spaces.GPU
     def process_example_fn(image_path: str) -> str:
         return model.run_image(image_path)
 
+    @spaces.GPU
+    def run(image: PIL.Image.Image, seed: int, guidance_scale: float, num_inference_steps: int) -> str:
+        return model.run_image(image, seed, guidance_scale, num_inference_steps)
+
     with gr.Blocks() as demo:
         with gr.Box():
             image = gr.Image(label="Input image", show_label=False, type="pil")
@@ -74,7 +81,7 @@ def create_demo(model: Model) -> gr.Blocks:
             queue=False,
             api_name=False,
         ).then(
-            fn=model.run_image,
+            fn=run,
             inputs=inputs,
             outputs=result,
             api_name="image-to-3d",

app_text_to_3d.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python
 
 import gradio as gr
+import spaces
 
 from model import Model
 from settings import CACHE_EXAMPLES, MAX_SEED
@@ -20,9 +21,14 @@ def create_demo(model: Model) -> gr.Blocks:
         "A bowl of vegetables",
     ]
 
+    @spaces.GPU
     def process_example_fn(prompt: str) -> str:
         return model.run_text(prompt)
 
+    @spaces.GPU
+    def run(prompt: str, seed: int, guidance_scale: float, num_inference_steps: int) -> str:
+        return model.run_text(prompt, seed, guidance_scale, num_inference_steps)
+
     with gr.Blocks() as demo:
         with gr.Box():
             with gr.Row(elem_id="prompt-container"):
@@ -80,7 +86,7 @@ def create_demo(model: Model) -> gr.Blocks:
             queue=False,
             api_name=False,
         ).then(
-            fn=model.run_text,
+            fn=run,
             inputs=inputs,
             outputs=result,
             api_name=False,
@@ -92,7 +98,7 @@ def create_demo(model: Model) -> gr.Blocks:
             queue=False,
             api_name=False,
         ).then(
-            fn=model.run_text,
+            fn=run,
             inputs=inputs,
             outputs=result,
             api_name="text-to-3d",

requirements.txt

@@ -1,6 +1,6 @@
 diffusers==0.21.2
 gradio==3.44.4
-torch==2.0.1
-torchvision==0.15.2
+torch==2.0.0
+torchvision==0.15.1
 transformers==4.33.2
 trimesh==3.23.5