From 52ae346c9dfedaf0df61a6f9d376ec9c5856e497 Mon Sep 17 00:00:00 2001 From: Manuel Schmid Date: Wed, 15 Nov 2023 22:00:28 +0100 Subject: [PATCH 1/3] add nsfw image censoring activatable via config, uses CompVis/stable-diffusion-safety-checker --- modules/async_worker.py | 9 +++++-- modules/censor.py | 54 +++++++++++++++++++++++++++++++++++++++++ modules/config.py | 5 ++++ 3 files changed, 66 insertions(+), 2 deletions(-) create mode 100644 modules/censor.py diff --git a/modules/async_worker.py b/modules/async_worker.py index a6807547..362cddde 100644 --- a/modules/async_worker.py +++ b/modules/async_worker.py @@ -30,6 +30,7 @@ def worker(): import fooocus_extras.ip_adapter as ip_adapter import fooocus_extras.face_crop + from modules.censor import censor_batch from modules.sdxl_styles import apply_style, apply_wildcards, fooocus_expansion from modules.private_logger import log from modules.expansion import safe_str @@ -50,12 +51,16 @@ def worker(): print(f'[Fooocus] {text}') outputs.append(['preview', (number, text, None)]) - def yield_result(imgs, do_not_show_finished_images=False): + def yield_result(imgs, do_not_show_finished_images=False, progressbar_index=13): global global_results if not isinstance(imgs, list): imgs = [imgs] + if modules.config.default_black_out_nsfw: + progressbar(progressbar_index, 'Checking for NSFW content ...') + imgs = censor_batch(imgs) + global_results = global_results + imgs if do_not_show_finished_images: @@ -711,7 +716,7 @@ def worker(): d.append((f'LoRA [{n}] weight', w)) log(x, d, single_line_number=3) - yield_result(imgs, do_not_show_finished_images=len(tasks) == 1) + yield_result(imgs, do_not_show_finished_images=len(tasks) == 1, progressbar_index=int(15.0 + 85.0 * float((current_task_id + 1) * steps) / float(all_steps)))) except fcbh.model_management.InterruptProcessingException as e: if shared.last_stop == 'skip': print('User skipped') diff --git a/modules/censor.py b/modules/censor.py new file mode 100644 index 00000000..fac6db09 --- /dev/null +++ b/modules/censor.py @@ -0,0 +1,54 @@ +# modified version of https://github.com/AUTOMATIC1111/stable-diffusion-webui-nsfw-censor/blob/master/scripts/censor.py + +import numpy as np +import torch +import modules.core as core + +from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker +from transformers import AutoFeatureExtractor +from PIL import Image + +safety_model_id = "CompVis/stable-diffusion-safety-checker" +safety_feature_extractor = None +safety_checker = None + + +def numpy_to_pil(image): + image = (image * 255).round().astype("uint8") + + #pil_image = Image.fromarray(image, 'RGB') + pil_image = Image.fromarray(image) + + return pil_image + + +# check and replace nsfw content +def check_safety(x_image): + global safety_feature_extractor, safety_checker + + if safety_feature_extractor is None: + safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id) + safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id) + + safety_checker_input = safety_feature_extractor(numpy_to_pil(x_image), return_tensors="pt") + x_checked_image, has_nsfw_concept = safety_checker(images=x_image, clip_input=safety_checker_input.pixel_values) + + return x_checked_image, has_nsfw_concept + + +def censor_single(x): + x_checked_image, has_nsfw_concept = check_safety(x) + + # replace image with black pixels, keep dimensions + # workaround due to different numpy / pytorch image matrix format + if has_nsfw_concept[0]: + imageshape = x_checked_image.shape + x_checked_image = np.zeros((imageshape[0], imageshape[1], 3), dtype = np.uint8) + + return x_checked_image + + +def censor_batch(images): + images = [censor_single(image) for image in images] + + return images diff --git a/modules/config.py b/modules/config.py index 7216a26d..ac81ef0b 100644 --- a/modules/config.py +++ b/modules/config.py @@ -268,6 +268,11 @@ default_overwrite_switch = get_config_item_or_set_default( default_value=-1, validator=lambda x: isinstance(x, int) ) +default_black_out_nsfw = get_config_item_or_set_default( + key='default_black_out_nsfw', + default_value=False, + validator=lambda x: isinstance(x, bool) +) def add_ratio(x): From 3dc69a5293b045cd560467baef388836c7b66dd6 Mon Sep 17 00:00:00 2001 From: Manuel Schmid Date: Sat, 18 Nov 2023 12:14:33 +0100 Subject: [PATCH 2/3] fix progressbar call for nsfw output --- modules/async_worker.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/modules/async_worker.py b/modules/async_worker.py index 28089554..c57a12ae 100644 --- a/modules/async_worker.py +++ b/modules/async_worker.py @@ -56,12 +56,12 @@ def worker(): print(f'[Fooocus] {text}') async_task.yields.append(['preview', (number, text, None)]) - def yield_result(async_task, imgs, do_not_show_finished_images=False): + def yield_result(async_task, imgs, do_not_show_finished_images=False, progressbar_index=13): if not isinstance(imgs, list): imgs = [imgs] if modules.config.default_black_out_nsfw: - progressbar(progressbar_index, 'Checking for NSFW content ...') + progressbar(async_task, progressbar_index, 'Checking for NSFW content ...') imgs = censor_batch(imgs) async_task.results = async_task.results + imgs @@ -725,7 +725,7 @@ def worker(): d.append((f'LoRA [{n}] weight', w)) log(x, d, single_line_number=3) - yield_result(async_task, imgs, do_not_show_finished_images=len(tasks) == 1, progressbar_index=int(15.0 + 85.0 * float((current_task_id + 1) * steps) / float(all_steps)))) + yield_result(async_task, imgs, do_not_show_finished_images=len(tasks) == 1, progressbar_index=int(15.0 + 85.0 * float((current_task_id + 1) * steps) / float(all_steps))) except fcbh.model_management.InterruptProcessingException as e: if shared.last_stop == 'skip': print('User skipped') From cdaaec3e71866c0ceedcd368c2ca329a223681a0 Mon Sep 17 00:00:00 2001 From: Manuel Schmid Date: Sat, 16 Dec 2023 18:54:46 +0100 Subject: [PATCH 3/3] use config to set cache dir for safety checker --- .gitignore | 1 + models/safety_checker_models/put_safety_checker_models_here | 0 modules/censor.py | 5 +++-- modules/config.py | 1 + 4 files changed, 5 insertions(+), 2 deletions(-) create mode 100644 models/safety_checker_models/put_safety_checker_models_here diff --git a/.gitignore b/.gitignore index 05ce1df8..da9cf974 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ config.txt config_modification_tutorial.txt user_path_config.txt user_path_config-deprecated.txt +/models/safety_checker_models /modules/*.png /repositories /venv diff --git a/models/safety_checker_models/put_safety_checker_models_here b/models/safety_checker_models/put_safety_checker_models_here new file mode 100644 index 00000000..e69de29b diff --git a/modules/censor.py b/modules/censor.py index fac6db09..e2352218 100644 --- a/modules/censor.py +++ b/modules/censor.py @@ -7,6 +7,7 @@ import modules.core as core from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker from transformers import AutoFeatureExtractor from PIL import Image +import modules.config safety_model_id = "CompVis/stable-diffusion-safety-checker" safety_feature_extractor = None @@ -27,8 +28,8 @@ def check_safety(x_image): global safety_feature_extractor, safety_checker if safety_feature_extractor is None: - safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id) - safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id) + safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id, cache_dir=modules.config.path_safety_checker_models) + safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id, cache_dir=modules.config.path_safety_checker_models) safety_checker_input = safety_feature_extractor(numpy_to_pil(x_image), return_tensors="pt") x_checked_image, has_nsfw_concept = safety_checker(images=x_image, clip_input=safety_checker_input.pixel_values) diff --git a/modules/config.py b/modules/config.py index ec0b49c3..c6515bd9 100644 --- a/modules/config.py +++ b/modules/config.py @@ -126,6 +126,7 @@ path_controlnet = get_dir_or_set_default('path_controlnet', '../models/controlne path_clip_vision = get_dir_or_set_default('path_clip_vision', '../models/clip_vision/') path_fooocus_expansion = get_dir_or_set_default('path_fooocus_expansion', '../models/prompt_expansion/fooocus_expansion') path_outputs = get_dir_or_set_default('path_outputs', '../outputs/') +path_safety_checker_models = get_dir_or_set_default('path_safety_checker_models', '../models/safety_checker_models/') def get_config_item_or_set_default(key, default_value, validator, disable_empty_as_none=False):