cymic/Waifu_Diffusion_Webui

1import torch2from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker3from transformers import AutoFeatureExtractor4from PIL import Image5 6import modules.shared as shared7 8safety_model_id = "CompVis/stable-diffusion-safety-checker"9safety_feature_extractor = None10safety_checker = None11 12def numpy_to_pil(images):13    """14    Convert a numpy image or a batch of images to a PIL image.15    """16    if images.ndim == 3:17        images = images[None, ...]18    images = (images * 255).round().astype("uint8")19    pil_images = [Image.fromarray(image) for image in images]20 21    return pil_images22 23# check and replace nsfw content24def check_safety(x_image):25    global safety_feature_extractor, safety_checker26 27    if safety_feature_extractor is None:28        safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id)29        safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id)30 31    safety_checker_input = safety_feature_extractor(numpy_to_pil(x_image), return_tensors="pt")32    x_checked_image, has_nsfw_concept = safety_checker(images=x_image, clip_input=safety_checker_input.pixel_values)33 34    return x_checked_image, has_nsfw_concept35 36 37def censor_batch(x):38    x_samples_ddim_numpy = x.cpu().permute(0, 2, 3, 1).numpy()39    x_checked_image, has_nsfw_concept = check_safety(x_samples_ddim_numpy)40    x = torch.from_numpy(x_checked_image).permute(0, 3, 1, 2)41 42    return x43