diff --git a/README.md b/README.md index b65f4a7..d61d4c4 100644 --- a/README.md +++ b/README.md @@ -25,16 +25,23 @@ Don't touch the FPS checkbox unless you know what you are doing. Additional command line arguments are given below: ``` --h, --help show this help message and exit --f SOURCE_IMG, --face SOURCE_IMG +options: + -h, --help show this help message and exit + -f SOURCE_IMG, --face SOURCE_IMG use this face --t TARGET_PATH, --target TARGET_PATH + -t TARGET_PATH, --target TARGET_PATH replace this face --o OUTPUT_FILE, --output OUTPUT_FILE - save output to this file ---keep-fps keep original fps ---gpu use gpu ---keep-frames don't delete frames directory + -o OUTPUT_FILE, --output OUTPUT_FILE + save output to this file + --keep-fps maintain original fps + --gpu use gpu + --keep-frames keep frames directory + --max-memory MAX_MEMORY + set max memory + --max-cpu-cores MAX_CPU_CORES + set max cpu cores + --max-cpu-usage MAX_CPU_USAGE + set cpu usage in percent ``` Looking for a CLI mode? Using the -f/--face argument will make the program in cli mode. diff --git a/run.py b/run.py index 300daab..9fabf17 100644 --- a/run.py +++ b/run.py @@ -1,4 +1,5 @@ #!/usr/bin/env python3 +import multiprocessing import sys import time import torch @@ -20,6 +21,7 @@ import psutil import cv2 import threading from PIL import Image, ImageTk +import resource pool = None args = {} @@ -31,6 +33,9 @@ parser.add_argument('-o', '--output', help='save output to this file', dest='out parser.add_argument('--keep-fps', help='maintain original fps', dest='keep_fps', action='store_true', default=False) parser.add_argument('--gpu', help='use gpu', dest='gpu', action='store_true', default=False) parser.add_argument('--keep-frames', help='keep frames directory', dest='keep_frames', action='store_true', default=False) +parser.add_argument('--max-memory', help='set max memory', default=16, type=int) +parser.add_argument('--max-cpu-cores', help='set max cpu cores', default=multiprocessing.cpu_count(), type=int) +parser.add_argument('--max-cpu-usage', help='set cpu usage in percent', default=100, type=int) for name, value in vars(parser.parse_args()).items(): args[name] = value @@ -40,6 +45,16 @@ if os.name == "nt": sep = "\\" +def limit_resources(): + current_cpu_usage, current_cpu_cores = resource.getrlimit(resource.RLIMIT_CPU) + if args['max_memory'] < 1: + resource.setrlimit(resource.RLIMIT_DATA, (args['max_memory'] * 1024 * 1024 * 1024, -1)) + if args['max_cpu_usage'] < 1: + resource.setrlimit(resource.RLIMIT_CPU, (args['max_cpu_usage'], current_cpu_cores)) + if args['max_cpu_cores'] < multiprocessing.cpu_count(): + resource.setrlimit(resource.RLIMIT_CPU, (current_cpu_usage, args['max_cpu_cores'])) + + def pre_check(): if not shutil.which('ffmpeg'): quit('ffmpeg is not installed!') @@ -199,6 +214,8 @@ if __name__ == "__main__": global status_label, window pre_check() + limit_resources() + if args['source_img']: start() quit()