# modified version of https://github.com/AUTOMATIC1111/stable-diffusion-webui-nsfw-censor/blob/master/scripts/censor.py
import numpy as np

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
safety_checker = None


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:
        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)
    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