Update app.py

app.py

@@ -31,7 +31,8 @@ from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orienta
 
 
 HEADER = """FRAME AI"""
-torch.cuda.empty_cache()
+
+
 if torch.cuda.is_available():
     device = "cuda:0"
 else:
@@ -150,6 +151,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
@@ -165,6 +167,7 @@ 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