feat: optimize model management of image censoring (#2960)

now follows general Fooocus model management principles + includes code optimisations for reusability
2026-08-16 13:13:16 +02:00 · 2024-05-19 18:36:47 +02:00
parent dad228907e
commit 35b74dfa64
2 changed files with 58 additions and 53 deletions
@@ -1,56 +1,60 @@
-# modified version of https://github.com/AUTOMATIC1111/stable-diffusion-webui-nsfw-censor/blob/master/scripts/censor.py
-import numpy as np
 import os

-from extras.safety_checker.models.safety_checker import StableDiffusionSafetyChecker
-from transformers import CLIPFeatureExtractor, CLIPConfig
-from PIL import Image
+import numpy as np
+import torch
+from transformers import CLIPConfig, CLIPImageProcessor
+
+import ldm_patched.modules.model_management as model_management
 import modules.config
+from extras.safety_checker.models.safety_checker import StableDiffusionSafetyChecker
+from ldm_patched.modules.model_patcher import ModelPatcher

 safety_checker_repo_root = os.path.join(os.path.dirname(__file__), 'safety_checker')
 config_path = os.path.join(safety_checker_repo_root, "configs", "config.json")
 preprocessor_config_path = os.path.join(safety_checker_repo_root, "configs", "preprocessor_config.json")

-safety_feature_extractor = None
-safety_checker = None
+
+class Censor:
+    def __init__(self):
+        self.safety_checker_model: ModelPatcher | None = None
+        self.clip_image_processor: CLIPImageProcessor | None = None
+        self.load_device = torch.device('cpu')
+        self.offload_device = torch.device('cpu')
+
+    def init(self):
+        if self.safety_checker_model is None and self.clip_image_processor is None:
+            safety_checker_model = modules.config.downloading_safety_checker_model()
+            self.clip_image_processor = CLIPImageProcessor.from_json_file(preprocessor_config_path)
+            clip_config = CLIPConfig.from_json_file(config_path)
+            model = StableDiffusionSafetyChecker.from_pretrained(safety_checker_model, config=clip_config)
+            model.eval()
+
+            self.load_device = model_management.text_encoder_device()
+            self.offload_device = model_management.text_encoder_offload_device()
+
+            model.to(self.offload_device)
+
+            self.safety_checker_model = ModelPatcher(model, load_device=self.load_device, offload_device=self.offload_device)
+
+    def censor(self, images: list | np.ndarray) -> list | np.ndarray:
+        self.init()
+        model_management.load_model_gpu(self.safety_checker_model)
+
+        single = False
+        if not isinstance(images, list) or isinstance(images, np.ndarray):
+            images = [images]
+            single = True
+
+        safety_checker_input = self.clip_image_processor(images, return_tensors="pt")
+        safety_checker_input.to(device=self.load_device)
+        checked_images, has_nsfw_concept = self.safety_checker_model.model(images=images,
+                                                                           clip_input=safety_checker_input.pixel_values)
+        checked_images = [image.astype(np.uint8) for image in checked_images]
+
+        if single:
+            checked_images = checked_images[0]
+
+        return checked_images


-def numpy_to_pil(image):
-    image = (image * 255).round().astype("uint8")
-    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 or safety_checker is None:
-        safety_checker_model = modules.config.downloading_safety_checker_model()
-        safety_feature_extractor = CLIPFeatureExtractor.from_json_file(preprocessor_config_path)
-        clip_config = CLIPConfig.from_json_file(config_path)
-        safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_checker_model, config=clip_config)
-
-    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
+default_censor = Censor().censor