From 6b0f3f93c86dae61b49cc6b3c74f24239ba73b6a Mon Sep 17 00:00:00 2001 From: henryruhs 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':