From aa25788e6a426329da0f5017fcc099cbdd8c17a8 Mon Sep 17 00:00:00 2001 From: Andres Berejnoi Date: Sun, 30 Oct 2022 14:31:10 -0400 Subject: [PATCH] Update txt2img.py --- scripts/txt2img.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/scripts/txt2img.py b/scripts/txt2img.py index 59c16a1db..02729c333 100644 --- a/scripts/txt2img.py +++ b/scripts/txt2img.py @@ -226,6 +226,13 @@ def main(): choices=["full", "autocast"], default="autocast" ) + + parser.add_argument( + "--nsfw_filter", + action="store_true", + help="Boolean Flag to enable the built-in safety check and NSFW filter" + ) + opt = parser.parse_args() if opt.laion400m: @@ -306,7 +313,10 @@ def main(): 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_checked_image, has_nsfw_concept = check_safety(x_samples_ddim) + if opt.nsfw_filter: + x_checked_image, has_nsfw_concept = check_safety(x_samples_ddim) + else: + x_checked_image = x_samples_ddim x_checked_image_torch = torch.from_numpy(x_checked_image).permute(0, 3, 1, 2)