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':