In exception handlers, clear the torch CUDA cache (if we're using CUDA) to free up memory for other programs using the GPU and to reduce fragmentation. (#2549)

invokeai/backend/invoke_ai_web_server.py

@@ -1208,12 +1208,18 @@ class InvokeAIWebServer:
             )
 
         except KeyboardInterrupt:
+            # Clear the CUDA cache on an exception
+            self.empty_cuda_cache()
             self.socketio.emit("processingCanceled")
             raise
         except CanceledException:
+            # Clear the CUDA cache on an exception
+            self.empty_cuda_cache()
             self.socketio.emit("processingCanceled")
             pass
         except Exception as e:
+            # Clear the CUDA cache on an exception
+            self.empty_cuda_cache()
             print(e)
             self.socketio.emit("error", {"message": (str(e))})
             print("\n")
@@ -1221,6 +1227,12 @@ class InvokeAIWebServer:
             traceback.print_exc()
             print("\n")
 
+    def empty_cuda_cache(self):
+        if self.generate.device.type == "cuda":
+            import torch.cuda
+
+            torch.cuda.empty_cache()
+
     def parameters_to_generated_image_metadata(self, parameters):
         try:
             # top-level metadata minus `image` or `images`