clear cuda cache

app.py

@@ -56,6 +56,8 @@ control_mode_ids = {
     "low_quality": 6,  # supported
 }
 
+def clear_cuda_cache():
+    torch.cuda.empty_cache()
 
 
 class calculateDuration:
@@ -94,11 +96,6 @@ def calculate_image_dimensions_for_flux(
 
     return new_width, new_height
 
-def is_mask_empty(image: Image.Image) -> bool:
-    gray_img = image.convert("L")
-    pixels = list(gray_img.getdata())
-    return all(pixel == 0 for pixel in pixels)
-
 def process_mask(
     mask: Image.Image,
     mask_inflation: Optional[int] = None,
@@ -121,7 +118,7 @@ def process_mask(
 
     if mask_blur and mask_blur > 0:
         mask = mask.filter(ImageFilter.GaussianBlur(radius=mask_blur))
-
+    clear_cuda_cache()
     return mask
 
 def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
@@ -143,6 +140,7 @@ def upload_image_to_r2(image, account_id, access_key, secret_key, bucket_name):
     buffer.seek(0)
     s3.upload_fileobj(buffer, bucket_name, image_file)
     print("upload finish", image_file)
+    clear_cuda_cache()
     return image_file
 
 
@@ -189,7 +187,7 @@ def run_flux(
             max_sequence_length=256,
             joint_attention_kwargs={"scale": lora_scale}
         ).images[0]
-    
+    clear_cuda_cache()
     return generated_image
 
 @spaces.GPU(duration=120)
@@ -214,7 +212,7 @@ def process(
     bucket:str,
     progress=gr.Progress(track_tqdm=True)
 ):
-    torch.cuda.empty_cache()
+    clear_cuda_cache()
     result = {"status": "false", "message": ""}
     if not image_url:
         gr.Info("please enter image url for inpaiting")
@@ -290,7 +288,7 @@ def process(
     
     control_image = control_image.resize((width, height), Image.LANCZOS)    
     control_mode_id = control_mode_ids[control_mode]
-
+    clear_cuda_cache()
     try:
         generated_image = run_flux(
             image=image,
@@ -323,7 +321,7 @@ def process(
         else:
             result = {"status": "success", "message": "Image generated but not uploaded"}
 
-
+    clear_cuda_cache()
     final_images = []
     final_images.append(image)
     final_images.append(mask)