[Safety Checker] Add Safety Checker Module

scripts/txt2img.py

@@ -16,12 +16,31 @@
 from ldm.models.diffusion.ddim import DDIMSampler
 from ldm.models.diffusion.plms import PLMSSampler
 
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from transformers import AutoFeatureExtractor
+
+# load safety model
+safety_model_id = "CompVis/stable-diffusion-v-1-3"
+safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id, use_auth_token=True)
+safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id, use_auth_token=True)
 
 def chunk(it, size):
  it = iter(it)
  return iter(lambda: tuple(islice(it, size)), ())
 
 
+def numpy_to_pil(images):
+ """
+ Convert a numpy image or a batch of images to a PIL 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 load_model_from_config(config, ckpt, verbose=False):
  print(f"Loading model from {ckpt}")
  pl_sd = torch.load(ckpt, map_location="cpu")
@@ -247,16 +266,23 @@ def main():
 
  x_samples_ddim = model.decode_first_stage(samples_ddim)
  x_samples_ddim = torch.clamp((x_samples_ddim + 1.0) / 2.0, min=0.0, max=1.0)
+ x_samples_ddim = x_samples_ddim.cpu().permute(0, 2, 3, 1).numpy()
+
+ x_image = x_samples_ddim
+ 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)
+
+ x_checked_image_torch = torch.from_numpy(x_checked_image).permute(0, 3, 2, 1)
 
  if not opt.skip_save:
- for x_sample in x_samples_ddim:
+ for x_sample in x_checked_image_torch:
  x_sample = 255. * rearrange(x_sample.cpu().numpy(), 'c h w -> h w c')
  Image.fromarray(x_sample.astype(np.uint8)).save(
  os.path.join(sample_path, f"{base_count:05}.png"))
  base_count += 1
 
  if not opt.skip_grid:
- all_samples.append(x_samples_ddim)
+ all_samples.append(x_checked_image_torch)
 
  if not opt.skip_grid:
  # additionally, save as grid