Upload app.py

app.py

@@ -1,8 +1,21 @@
 import gradio as gr
-import spaces
 from gradio_litmodel3d import LitModel3D
 
 import os
+if os.environ.get("SPACES_ZERO_GPU") is not None:
+    import spaces
+else:
+    class spaces:
+        @staticmethod
+        def GPU(func):
+            def wrapper(*args, **kwargs):
+                return func(*args, **kwargs)
+            return wrapper
+
+@spaces.GPU
+def fake_gpu():
+  pass
+
 import shutil
 os.environ['SPCONV_ALGO'] = 'native'
 from typing import *
@@ -97,6 +110,7 @@ def segment_and_overlay(image, points, sam_predictor):
     return overlaid, visible_mask
 
 
+@spaces.GPU
 def image_to_3d(
     images: List[tuple],
     masks: List[np.ndarray],
@@ -133,33 +147,38 @@ def image_to_3d(
     video_path = os.path.join(user_dir, 'sample.mp4')
     imageio.mimsave(video_path, video, fps=15)
     state = pack_state(outputs['gaussian'][0], outputs['mesh'][0])
+    torch.cuda.empty_cache()
     return state, video_path
 
-# def extract_glb(
-#     state: dict,
-#     mesh_simplify: float,
-#     texture_size: int,
-#     req: gr.Request,
-# ) -> tuple:
-#     """
-#     从生成的 3D 模型中提取 GLB 文件。
-#     """
-#     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-#     gs, mesh = unpack_state(state)
-#     glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
-#     glb_path = os.path.join(user_dir, 'sample.glb')
-#     glb.export(glb_path)
-#     return glb_path, glb_path
-
-# def extract_gaussian(state: dict, req: gr.Request) -> tuple:
-#     """
-#     从生成的 3D 模型中提取 Gaussian 文件。
-#     """
-#     user_dir = os.path.join(TMP_DIR, str(req.session_hash))
-#     gs, _ = unpack_state(state)
-#     gaussian_path = os.path.join(user_dir, 'sample.ply')
-#     gs.save_ply(gaussian_path)
-#     return gaussian_path, gaussian_path
+@spaces.GPU(duration=90)
+def extract_glb(
+    state: dict,
+    mesh_simplify: float,
+    texture_size: int,
+    req: gr.Request,
+) -> tuple:
+    """
+    从生成的 3D 模型中提取 GLB 文件。
+    """
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    gs, mesh = unpack_state(state)
+    glb = postprocessing_utils.to_glb(gs, mesh, simplify=mesh_simplify, texture_size=texture_size, verbose=False)
+    glb_path = os.path.join(user_dir, 'sample.glb')
+    glb.export(glb_path)
+    torch.cuda.empty_cache()
+    return glb_path, glb_path
+
+@spaces.GPU
+def extract_gaussian(state: dict, req: gr.Request) -> tuple:
+    """
+    从生成的 3D 模型中提取 Gaussian 文件。
+    """
+    user_dir = os.path.join(TMP_DIR, str(req.session_hash))
+    gs, _ = unpack_state(state)
+    gaussian_path = os.path.join(user_dir, 'sample.ply')
+    gs.save_ply(gaussian_path)
+    torch.cuda.empty_cache()
+    return gaussian_path, gaussian_path
 
 
 def pack_state(gs: Gaussian, mesh: MeshExtractResult) -> dict: