Update app.py

app.py

@@ -36,19 +36,10 @@ os.environ["CUDA_VISIBLE_DEVICES"] = "1"
 
 HEADER = """FRAME AI"""
 
-# if torch.cuda.is_available():
-#     device = "cuda:0"
-# else:
-#     device = "cpu"
-
-# torch.cuda.set_device(1)
-
-# CUDA_LAUNCH_BLOCKING=1
-
-
 if torch.cuda.is_available():
-    torch.cuda.set_device(1)
-
+    device = "cuda:0"
+else:
+    device = "cpu"
     
 
 model = TSR.from_pretrained(
@@ -168,6 +159,7 @@ def check_input_image(input_image):
         raise gr.Error("No image uploaded!")
 
 def preprocess(input_image, do_remove_background, foreground_ratio):
+    torch.cuda.synchronize()
     def fill_background(image):
         image = np.array(image).astype(np.float32) / 255.0
         image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
@@ -185,10 +177,12 @@ def preprocess(input_image, do_remove_background, foreground_ratio):
         image = input_image
         if image.mode == "RGBA":
             image = fill_background(image)
+    torch.cuda.synchronize()
     return image
 
 # @spaces.GPU
 def generate(image, mc_resolution, formats=["obj", "glb"]):
+    torch.cuda.synchronize()
     torch.cuda.empty_cache()
     scene_codes = model(image, device=device)
     mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0]
@@ -200,6 +194,8 @@ def generate(image, mc_resolution, formats=["obj", "glb"]):
     mesh_path_obj = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False)
     mesh.apply_scale([-1, 1, 1])  # Otherwise the visualized .obj will be flipped
     mesh.export(mesh_path_obj.name)
+
+    torch.cuda.synchronize()
     
     return mesh_path_obj.name, mesh_path_glb.name