rebase changes of main for easier handling

modules/async_worker.py

@@ -13,7 +13,6 @@ async_tasks = []
 
 def worker():
     global async_tasks
-
     import traceback
     import math
     import numpy as np
@@ -35,6 +34,8 @@ 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
@@ -55,10 +56,14 @@ 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(async_task, progressbar_index, 'Checking for NSFW content ...')
+            imgs = censor_batch(imgs)
+
         async_task.results = async_task.results + imgs
 
         if do_not_show_finished_images:
@@ -652,7 +657,7 @@ def worker():
             done_steps = current_task_id * steps + step
             async_task.yields.append(['preview', (
                 int(15.0 + 85.0 * float(done_steps) / float(all_steps)),
-                f'Step {step}/{total_steps} in the {current_task_id + 1}-th Sampling',
+                f'Sampling Image {current_task_id + 1}/{image_number}, Step {step + 1}/{total_steps} ...',
                 y)])
 
         for current_task_id, task in enumerate(tasks):
@@ -720,11 +725,14 @@ 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)
+                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')
                     continue
+                elif shared.last_stop == 'stop_previous':
+                    print('Previous task stopped')
+                    break
                 else:
                     print('User stopped')
                     break

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

modules/config.py

@@ -243,10 +243,15 @@ default_advanced_checkbox = get_config_item_or_set_default(
     default_value=False,
     validator=lambda x: isinstance(x, bool)
 )
+default_max_image_number = get_config_item_or_set_default(
+    key='default_max_image_number',
+    default_value=4,
+    validator=lambda x: isinstance(x, int) and x >= 1 and x <= 32
+)
 default_image_number = get_config_item_or_set_default(
     key='default_image_number',
     default_value=2,
-    validator=lambda x: isinstance(x, int) and 1 <= x <= 32
+    validator=lambda x: isinstance(x, int) and 1 <= x <= default_max_image_number
 )
 checkpoint_downloads = get_config_item_or_set_default(
     key='checkpoint_downloads',
@@ -303,6 +308,16 @@ 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)
+)
+default_hide_preview_if_black_out_nsfw = get_config_item_or_set_default(
+    key='default_hide_preview_if_black_out_nsfw',
+    default_value=True,
+    validator=lambda x: isinstance(x, bool)
+)
 
 config_dict["default_loras"] = default_loras = default_loras[:5] + [['None', 1.0] for _ in range(5 - len(default_loras))]