From 6b0f3f93c86dae61b49cc6b3c74f24239ba73b6a Mon Sep 17 00:00:00 2001
From: henryruhs <info@henryruhs.com>
Date: Sun, 4 Jun 2023 22:22:35 +0200
Subject: [PATCH] Fix memory leak thanks to tensorflow

---
 roop/core.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/roop/core.py b/roop/core.py
index 9abb93f..c14bc48 100755
--- a/roop/core.py
+++ b/roop/core.py
@@ -12,6 +12,7 @@ import glob
 import argparse
 import psutil
 import torch
+import tensorflow
 from pathlib import Path
 import multiprocessing as mp
 from opennsfw2 import predict_video_frames, predict_image
@@ -66,6 +67,10 @@ if os.name == "nt":
 
 
 def limit_resources():
+    # prevent tensorflow memory leak
+    gpus = tensorflow.config.experimental.list_physical_devices('GPU')
+    for gpu in gpus:
+        tensorflow.config.experimental.set_memory_growth(gpu, True)
     if args.max_memory:
         memory = args.max_memory * 1024 * 1024 * 1024
         if str(platform.system()).lower() == 'windows':