Update app.py

app.py

@@ -1,5 +1,3 @@
-import os 
-os.system("pip install /home/user/app/third_party/voxelize/udf_ext-0.0.0-cp310-cp310-linux_x86_64.whl") 
 import torch
 import trimesh
 import datetime
@@ -16,7 +14,6 @@ from typing import Any
 
 import gradio as gr
 from gradio.themes.utils import colors, fonts, sizes
-import spaces
 
 # -----------------------------------------------------------------------------
 #  THEME  ▸  a soft glass-like dark theme with a vibrant primary accent
@@ -51,7 +48,18 @@ def check_input_image(input_image):
 # -----------------------------------------------------------------------------
 #  PLACEHOLDER BACK-END HOOKS  ▸  replace with your real logic
 # -----------------------------------------------------------------------------
-@spaces.GPU(duration=600)
+
+
+pipe = Direct3DS2Pipeline.from_pretrained('wushuang98/Direct3D-S2', subfolder="direct3d-s2-v-1-1")
+pipe.to("cuda")
+
+rmbg = BiRefNet(device="cuda")
+
+@spaces.GPU
+def run_rembg(image: Image.Image):
+    return rmbg.run(image)
+    
+@spaces.GPU
 def image2mesh(
     image: Any, 
     resolution: str = '1024', 
@@ -68,10 +76,6 @@ def image2mesh(
     uid = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
     image.save(os.path.join(output_path, uid + '.png'))
 
-    pipe = Direct3DS2Pipeline.from_pretrained('wushuang98/Direct3D-S2', subfolder="direct3d-s2-v-1-1")
-    # pipe.to("cuda:0")
-    pipe.to("cuda")
-
     mesh = pipe(
         image, 
         sdf_resolution=int(resolution), 
@@ -159,13 +163,6 @@ body { background:linear-gradient(215deg,#101113 0%,#0b0c0d 60%,#0d1014 100%) }
         # ---------- Gallery / Examples ----------
         with gr.Column(scale=3):
             gr.Markdown("### Examples", elem_classes="subtitle")
-            # gr.Examples(
-            #     examples=[os.path.join("assets/test", i) for i in os.listdir("assets/test")],
-            #     inputs=[image_input],
-            #     examples_per_page=8,
-            #     label="Gallery",
-            #     elem_id="examples_gallery",
-            # )
             with gr.Tabs(selected='tab_img_gallery') as gallery:
                 with gr.Tab('Image to 3D Gallery', id='tab_img_gallery') as tab_gi:
                     with gr.Row():
@@ -175,25 +172,16 @@ body { background:linear-gradient(215deg,#101113 0%,#0b0c0d 60%,#0d1014 100%) }
                             label=None, 
                             examples_per_page=24
                         )
-            # gallery = gr.Gallery(
-            #     [os.path.join("assets/test", i) for i in os.listdir("assets/test")],
-            #     columns=2,
-            #     object_fit="contain",
-            #     elem_id="examples_gallery",
-            #     allow_preview=False,
-            # )
 
             
     # ▸ callbacks
     outputs = [output_model_obj]
-    # rmbg = BiRefNet(device="cuda:0")
-    rmbg = BiRefNet(device="cuda")
 
     gen_btn.click(
         fn=check_input_image, 
         inputs=[image_input]
     ).success(
-        fn=rmbg.run, 
+        fn=run_rembg, 
         inputs=[image_input],
         outputs=[processed_image]
     ).success(
@@ -209,4 +197,4 @@ if __name__ == "__main__":
     parser.add_argument("--cached_dir", type=str, default="outputs/web")
     args = parser.parse_args()
     
-    demo.queue().launch(share=True, allowed_paths=[args.cached_dir], server_port=7860)
+    demo.queue().launch(share=True, allowed_paths=[args.cached_dir], server_port=7860)