a

app.py

@@ -42,6 +42,7 @@ from src.data import DemoData
 from src.models import LiNo_UniPS
 from torch.utils.data import DataLoader
 import pytorch_lightning as pl
+import spaces
 
 MAX_SEED = np.iinfo(np.int32).max
 TMP_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'tmp')
@@ -49,12 +50,8 @@ WEIGHTS_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'weights'
 os.makedirs(TMP_DIR, exist_ok=True)
 os.makedirs(WEIGHTS_DIR, exist_ok=True)
 
-is_gpu_available = torch.cuda.is_available()
 
 
-if is_gpu_available:
-    print("✅ NVIDIA GPU。")
-
 
 def cache_weights(weights_dir: str) -> dict:
     import os
@@ -88,7 +85,7 @@ def preprocess_mesh(mesh_prompt):
     trimesh_mesh = trimesh.load_mesh(mesh_prompt)
     trimesh_mesh.export(mesh_prompt+'.glb')
     return mesh_prompt+'.glb'
-
+@spaces.GPU
 def generate_3d(image, seed=-1,  
                 ss_guidance_strength=3, ss_sampling_steps=50,
                 slat_guidance_strength=3, slat_sampling_steps=6,normal_bridge=None):
@@ -136,7 +133,7 @@ def generate_3d(image, seed=-1,
     trimesh_mesh.export(mesh_path)
 
     return mesh_path, mesh_path
-
+@spaces.GPU
 def predict_normal(input_images,input_mask):
     test_dataset = DemoData(input_imgs_list=input_images,input_mask=input_mask)
     test_loader = DataLoader(test_dataset, batch_size=1)