spaces

app.py

@@ -1,4 +1,5 @@
 import os
+import spaces
 import subprocess
 def install_cuda_toolkit():
     CUDA_TOOLKIT_URL = "https://developer.download.nvidia.com/compute/cuda/12.4.0/local_installers/cuda_12.4.0_550.54.14_linux.run"
@@ -31,10 +32,11 @@ from step1x3d_texture.utils.shape_post_process import (
     DegenerateFaceRemover,
 )
 
-
+@spaces.GPU(duration=240)
 def generate_func(
     input_image_path, guidance_scale, inference_steps, max_facenum, symmetry, edge_type
 ):
+    geometry_model = geometry_model.to("cuda")
     if "Label" in args.geometry_model:
         out = geometry_model(
             input_image_path,
@@ -86,7 +88,7 @@ if __name__ == "__main__":
 
     geometry_model = Step1X3DGeometryPipeline.from_pretrained(
         "stepfun-ai/Step1X-3D", subfolder=args.geometry_model
-    ).to("cuda")
+    )
 
     texture_model = Step1X3DTexturePipeline.from_pretrained("stepfun-ai/Step1X-3D", subfolder=args.texture_model)