Update app.py

app.py

@@ -22,9 +22,10 @@ import io
 from io import BytesIO
 from botocore.exceptions import NoCredentialsError, PartialCredentialsError
 
+torch.cuda.empty_cache()
+
 app = FastAPI()
 
-torch.cuda.empty_cache()
 
 subprocess.run(shlex.split('pip install wheel/torchmcubes-0.1.0-cp310-cp310-linux_x86_64.whl'))
 
@@ -168,6 +169,7 @@ def preprocess(input_image, do_remove_background, foreground_ratio):
         image = resize_foreground(image, foreground_ratio)
         image = fill_background(image)
         print("do_remove_back")
+        torch.cuda.empty_cache()
     else:
         image = input_image
         if image.mode == "RGBA":
@@ -176,6 +178,7 @@ def preprocess(input_image, do_remove_background, foreground_ratio):
 
 # @spaces.GPU
 def generate(image, mc_resolution, formats=["obj", "glb"]):
+    torch.cuda.empty_cache()
     scene_codes = model(image, device=device)
     mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0]
     mesh = to_gradio_3d_orientation(mesh)