| import os | |
| 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") | |
| 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, 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 | |
| default_censor = Censor().censor | |