Update app.py

app.py

@@ -23,8 +23,6 @@ import io
 from io import BytesIO
 from botocore.exceptions import NoCredentialsError, PartialCredentialsError
 import datetime
-from transformers.utils import move_cache
-move_cache()
 
 app = FastAPI()
 
@@ -41,7 +39,7 @@ if torch.cuda.is_available():
 else:
     device = "cpu"
 
-torch.cuda.synchronize()
+# torch.cuda.synchronize()
 
 model = TSR.from_pretrained(
     "stabilityai/TripoSR",
@@ -151,16 +149,16 @@ def check_input_image(input_image):
 
 def preprocess(input_image, do_remove_background, foreground_ratio):
     def fill_background(image):
-        torch.cuda.synchronize()
-        torch.cuda.empty_cache()
+        # torch.cuda.synchronize()
+        # torch.cuda.empty_cache()
         image = np.array(image).astype(np.float32) / 255.0
         image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
         image = Image.fromarray((image * 255.0).astype(np.uint8))
         return image
 
     if do_remove_background:
-        torch.cuda.synchronize()
-        torch.cuda.empty_cache()
+        # torch.cuda.synchronize()
+        # torch.cuda.empty_cache()
         image = input_image.convert("RGB")
         image = remove_background(image, rembg_session)
         image = resize_foreground(image, foreground_ratio)
@@ -173,39 +171,39 @@ def preprocess(input_image, do_remove_background, foreground_ratio):
         image = input_image
         if image.mode == "RGBA":
             image = fill_background(image)
-    torch.cuda.synchronize()  # Wait for all CUDA operations to complete
-    torch.cuda.empty_cache()
+    # torch.cuda.synchronize()  # Wait for all CUDA operations to complete
+    # torch.cuda.empty_cache()
     return image
 
 @spaces.GPU
 def generate(image, mc_resolution, formats=["obj", "glb"]):
-    torch.cuda.synchronize()
+    # torch.cuda.synchronize()
     scene_codes = model(image, device=device)
-    torch.cuda.synchronize()
+    # torch.cuda.synchronize()
     mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0]
-    torch.cuda.synchronize()
+    # torch.cuda.synchronize()
     mesh = to_gradio_3d_orientation(mesh)
-    torch.cuda.synchronize()
+    # torch.cuda.synchronize()
 
     mesh_path_glb = tempfile.NamedTemporaryFile(suffix=f".glb", delete=False)
-    torch.cuda.synchronize()
+    # torch.cuda.synchronize()
     mesh.export(mesh_path_glb.name)
-    torch.cuda.synchronize()
+    # torch.cuda.synchronize()
 
     mesh_path_obj = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False)
-    torch.cuda.synchronize()
+    # torch.cuda.synchronize()
     mesh.apply_scale([-1, 1, 1])  # Otherwise the visualized .obj will be flipped
     mesh.export(mesh_path_obj.name)
-    torch.cuda.synchronize()  # Ensure all CUDA operations are complete before clearing cache
-    torch.cuda.empty_cache()
+    # torch.cuda.synchronize()  # Ensure all CUDA operations are complete before clearing cache
+    # torch.cuda.empty_cache()
     return mesh_path_obj.name, mesh_path_glb.name
 
 def upload_file_to_s3(file_path, bucket_name, object_name=None):    
     s3_client.upload_file(file_path, bucket_name, object_name)
     
     # print(f"File {file_path} uploaded successfully to {bucket_name}/{object_name}.")
-    torch.cuda.synchronize()  # Wait for all CUDA operations to complete
-    torch.cuda.empty_cache()
+    # torch.cuda.synchronize()  # Wait for all CUDA operations to complete
+    # torch.cuda.empty_cache()
     return True
 
     
@@ -238,8 +236,8 @@ async def process_image(
         object_name_2 = f'object_{timestamp}_2.glb'
     
         if upload_file_to_s3(mesh_name_obj, 'framebucket3d',object_name) and upload_file_to_s3(mesh_name_glb, 'framebucket3d',object_name_2):
-            torch.cuda.synchronize()  # Wait for all CUDA operations to complete
-            torch.cuda.empty_cache()
+            # torch.cuda.synchronize()  # Wait for all CUDA operations to complete
+            # torch.cuda.empty_cache()
             return {
                 "obj_path": f"https://framebucket3d.s3.amazonaws.com/{object_name}",
                 "glb_path": f"https://framebucket3d.s3.amazonaws.com/{object_name_2}"