diff --git a/scripts/txt2img.py b/scripts/txt2img.py index da77e1a03e97908953dc60880966db358b6a79a0..1b8888fa0305c78d638ff2dc68e8e5d3d9357e56 100644 --- a/scripts/txt2img.py +++ b/scripts/txt2img.py @@ -16,12 +16,31 @@ from ldm.util import instantiate_from_config 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