58 lines
1.8 KiB
Python
58 lines
1.8 KiB
Python
from typing import List, Union, Dict, Set, Tuple
|
|
|
|
from diffusers.pipelines.stable_diffusion.safety_checker import (
|
|
StableDiffusionSafetyChecker,
|
|
)
|
|
from transformers import AutoFeatureExtractor
|
|
import torch
|
|
from PIL import Image, ImageFilter
|
|
import numpy as np
|
|
|
|
safety_model_id: str = "CompVis/stable-diffusion-safety-checker"
|
|
safety_feature_extractor: AutoFeatureExtractor = None
|
|
safety_checker: StableDiffusionSafetyChecker = None
|
|
|
|
|
|
def numpy_to_pil(images: np.ndarray) -> List[Image.Image]:
|
|
if images.ndim == 3:
|
|
images = images[None, ...]
|
|
images = (images * 255).round().astype("uint8")
|
|
pil_images = [Image.fromarray(image) for image in images]
|
|
|
|
return pil_images
|
|
|
|
|
|
def check_image(x_image: np.ndarray) -> Tuple[np.ndarray, List[bool]]:
|
|
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(
|
|
images=numpy_to_pil(x_image), return_tensors="pt"
|
|
)
|
|
x_checked_image, hs = safety_checker(
|
|
images=x_image, clip_input=safety_checker_input.pixel_values
|
|
)
|
|
|
|
return x_checked_image, hs
|
|
|
|
|
|
def check_batch(x: torch.Tensor) -> torch.Tensor:
|
|
x_samples_ddim_numpy = x.cpu().permute(0, 2, 3, 1).numpy()
|
|
x_checked_image, _ = check_image(x_samples_ddim_numpy)
|
|
x = torch.from_numpy(x_checked_image).permute(0, 3, 1, 2)
|
|
return x
|
|
|
|
|
|
def convert_to_sd(img: Image) -> Image:
|
|
_, hs = check_image(np.array(img))
|
|
if any(hs):
|
|
img = (
|
|
img.resize((int(img.width * 0.1), int(img.height * 0.1)))
|
|
.resize(img.size, Image.BOX)
|
|
.filter(ImageFilter.BLUR)
|
|
)
|
|
return img
|