Update handler.py

handler.py

@@ -665,8 +665,20 @@ class EndpointHandler():
 
         with_big_tits = data.get("with_big_tits", False)
 
-        nude_pils = get_nude(get_mask_function=just_fucking_get_sd_mask, cfg_scale=data.get("cfg_scale"), generate_max_size=data.get("generate_max_size"), original_max_size=data.get(
+        nude_pils = []
+
+        try:
+            nude_pils = get_nude(get_mask_function=just_fucking_get_sd_mask, cfg_scale=data.get("cfg_scale"), generate_max_size=data.get("generate_max_size"), original_max_size=data.get(
             "original_max_size"), original_pil=original_pil, positive_prompt=data.get("positive_prompt"), steps=data.get("steps"), with_small_tits=with_small_tits, with_big_tits=with_big_tits)
+        except RuntimeError as e:
+            if 'out of memory' in str(e):
+                torch.cuda.empty_cache()
+                nude_pils = get_nude(get_mask_function=just_fucking_get_sd_mask, cfg_scale=data.get("cfg_scale"), generate_max_size=data.get("generate_max_size"), original_max_size=data.get(
+                "original_max_size"), original_pil=original_pil, positive_prompt=data.get("positive_prompt"), steps=data.get("steps"), with_small_tits=with_small_tits, with_big_tits=with_big_tits)
+                print("CUDA Out of Memory, clearing cache")
+                # Optionally, you can retry your operation here, or handle the exception further
+            else:
+                raise
 
         filenames = []