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 argparse
import psutil import psutil
import torch import torch
import tensorflow
from pathlib import Path from pathlib import Path
import multiprocessing as mp import multiprocessing as mp
from opennsfw2 import predict_video_frames, predict_image from opennsfw2 import predict_video_frames, predict_image
@ -66,6 +67,10 @@ if os.name == "nt":
def limit_resources(): 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: if args.max_memory:
memory = args.max_memory * 1024 * 1024 * 1024 memory = args.max_memory * 1024 * 1024 * 1024
if str(platform.system()).lower() == 'windows': if str(platform.system()).lower() == 'windows':