Update app.py

app.py

@@ -165,7 +165,8 @@ def preprocess(input_image, do_remove_background, foreground_ratio):
         image = input_image
         if image.mode == "RGBA":
             image = fill_background(image)
-    torch.cuda.synchronize()
+    torch.cuda.synchronize()  # Wait for all CUDA operations to complete
+    torch.cuda.empty_cache()
     return image
 
 # @spaces.GPU
@@ -182,7 +183,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()
+    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):
@@ -239,7 +241,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()
             return {
                 "obj_path": f"https://framebucket3d.s3.amazonaws.com/{object_name}",
                 "glb_path": f"https://framebucket3d.s3.amazonaws.com/{object_name_2}"