Re-enabled the NSFW function (turn-off by default).

@refer to the PR #237 in the `experimental` branch.

Signed-off-by: Vic P <vic4key@gmail.com>
This commit is contained in:
Vic P 2024-08-21 02:02:00 +07:00
parent e4b494174d
commit 7313a332c8
5 changed files with 46 additions and 26 deletions

View file

@ -39,6 +39,7 @@ def parse_args() -> None:
program.add_argument('--keep-audio', help='keep original audio', dest='keep_audio', action='store_true', default=True)
program.add_argument('--keep-frames', help='keep temporary frames', dest='keep_frames', action='store_true', default=False)
program.add_argument('--many-faces', help='process every face', dest='many_faces', action='store_true', default=False)
program.add_argument('--nsfw-filter', help='filter the NSFW image or video', dest='nsfw_filter', action='store_true', default=False)
program.add_argument('--video-encoder', help='adjust output video encoder', dest='video_encoder', default='libx264', choices=['libx264', 'libx265', 'libvpx-vp9'])
program.add_argument('--video-quality', help='adjust output video quality', dest='video_quality', type=int, default=18, choices=range(52), metavar='[0-51]')
program.add_argument('--max-memory', help='maximum amount of RAM in GB', dest='max_memory', type=int, default=suggest_max_memory())
@ -63,6 +64,7 @@ def parse_args() -> None:
modules.globals.keep_audio = args.keep_audio
modules.globals.keep_frames = args.keep_frames
modules.globals.many_faces = args.many_faces
modules.globals.nsfw_filter = args.nsfw_filter
modules.globals.video_encoder = args.video_encoder
modules.globals.video_quality = args.video_quality
modules.globals.max_memory = args.max_memory
@ -75,8 +77,6 @@ def parse_args() -> None:
else:
modules.globals.fp_ui['face_enhancer'] = False
modules.globals.nsfw = False
# translate deprecated args
if args.source_path_deprecated:
print('\033[33mArgument -f and --face are deprecated. Use -s and --source instead.\033[0m')
@ -169,13 +169,15 @@ def start() -> None:
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
if not frame_processor.pre_start():
return
update_status('Processing...')
# process image to image
if has_image_extension(modules.globals.target_path):
if modules.globals.nsfw == False:
from modules.predicter import predict_image
if predict_image(modules.globals.target_path):
destroy()
shutil.copy2(modules.globals.target_path, modules.globals.output_path)
if modules.globals.nsfw_filter and ui.check_and_ignore_nsfw(modules.globals.target_path, destroy):
return
try:
shutil.copy2(modules.globals.target_path, modules.globals.output_path)
except Exception as e:
print("Error copying file:", str(e))
for frame_processor in get_frame_processors_modules(modules.globals.frame_processors):
update_status('Progressing...', frame_processor.NAME)
frame_processor.process_image(modules.globals.source_path, modules.globals.output_path, modules.globals.output_path)
@ -186,10 +188,8 @@ def start() -> None:
update_status('Processing to image failed!')
return
# process image to videos
if modules.globals.nsfw == False:
from modules.predicter import predict_video
if predict_video(modules.globals.target_path):
destroy()
if modules.globals.nsfw_filter and ui.check_and_ignore_nsfw(modules.globals.target_path, destroy):
return
update_status('Creating temp resources...')
create_temp(modules.globals.target_path)
update_status('Extracting frames...')
@ -225,10 +225,10 @@ def start() -> None:
update_status('Processing to video failed!')
def destroy() -> None:
def destroy(to_quit=True) -> None:
if modules.globals.target_path:
clean_temp(modules.globals.target_path)
quit()
if to_quit: quit()
def run() -> None: