update

app.py

@@ -13,9 +13,12 @@ from PIL import Image
 
 import sf3d.utils as sf3d_utils
 from sf3d.system import SF3D
+import devicetorch
 
 rembg_session = rembg.new_session()
 
+DEVICE = devicetorch.get(torch)
+
 COND_WIDTH = 512
 COND_HEIGHT = 512
 COND_DISTANCE = 1.6
@@ -34,7 +37,8 @@ model = SF3D.from_pretrained(
     config_name="config.yaml",
     weight_name="model.safetensors",
 )
-model.eval().cuda()
+#model.eval().cuda()
+model.eval().to(device)
 
 example_files = [
     os.path.join("demo_files/examples", f) for f in os.listdir("demo_files/examples")
@@ -44,9 +48,15 @@ example_files = [
 def run_model(input_image):
     start = time.time()
     with torch.no_grad():
-        with torch.autocast(device_type="cuda", dtype=torch.float16):
+        if DEVICE == "cuda":
+          with torch.autocast(device_type="cuda", dtype=torch.float16):
+              model_batch = create_batch(input_image)
+              model_batch = {k: v.cuda() for k, v in model_batch.items()}
+              trimesh_mesh, _glob_dict = model.generate_mesh(model_batch, 1024)
+              trimesh_mesh = trimesh_mesh[0]
+        else:
             model_batch = create_batch(input_image)
-            model_batch = {k: v.cuda() for k, v in model_batch.items()}
+            model_batch = {k: v.to(DEVICE) for k, v in model_batch.items()}
             trimesh_mesh, _glob_dict = model.generate_mesh(model_batch, 1024)
             trimesh_mesh = trimesh_mesh[0]
 

requirements.txt

@@ -1,5 +1,5 @@
-torch==2.1.2
-torchvision==0.16.2
+#torch==2.1.2
+#torchvision==0.16.2
 einops==0.7.0
 jaxtyping==0.2.31
 omegaconf==2.3.0