upagge/roop
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Fix memory leak thanks to tensorflow

Browse Source
This commit is contained in:
henryruhs 2023-06-04 22:22:35 +02:00
parent 9c88540629
commit 6b0f3f93c8
1 changed files with 5 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
roop/core.py
View File

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