Update app.py

app.py

@@ -151,40 +151,56 @@ 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):
+        torch.cuda.synchronize()  # Ensure previous CUDA operations are complete
         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()
         image = input_image.convert("RGB")
         image = remove_background(image, rembg_session)
         image = resize_foreground(image, foreground_ratio)
         image = fill_background(image)
+        
+        torch.cuda.synchronize()
     else:
         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
 def generate(image, mc_resolution, formats=["obj", "glb"]):
+    torch.cuda.synchronize()
     scene_codes = model(image, device=device)
+    torch.cuda.synchronize()
     mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0]
+    torch.cuda.synchronize()
     mesh = to_gradio_3d_orientation(mesh)
-
+    torch.cuda.synchronize()
+    
     mesh_path_glb = tempfile.NamedTemporaryFile(suffix=f".glb", delete=False)
+    torch.cuda.synchronize()
     mesh.export(mesh_path_glb.name)
-
+    torch.cuda.synchronize()
+    
     mesh_path_obj = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False)
-    mesh.apply_scale([-1, 1, 1])  # Otherwise the visualized .obj will be flipped
+    torch.cuda.synchronize()
+    mesh.apply_scale([-1, 1, 1])
     mesh.export(mesh_path_obj.name)
-    
+    torch.cuda.synchronize()
+    torch.cuda.empty_cache()
     return mesh_path_obj.name, mesh_path_glb.name
 
+    
+
 def run_example(text_prompt,seed ,do_remove_background, foreground_ratio, mc_resolution):
     image_pil = generate_image_from_text(text_prompt, seed)