artem/stable-diffusion
1
0
Fork 0
mirror of https://github.com/CompVis/stable-diffusion.git synced 2026-05-09 08:55:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Safety Checker] Add Safety Checker Module

Browse Source
This commit is contained in:
Patrick von Platen 2022-08-19 16:01:56 +00:00
parent 7b8c883b07
commit d0c714ae4a
1 changed files with 24 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
scripts/txt2img.py
View File

@ -16,12 +16,29 @@ from ldm.util import instantiate_from_config
from ldm.models.diffusion.ddim import DDIMSampler from ldm.models.diffusion.ddim import DDIMSampler
from ldm.models.diffusion.plms import PLMSSampler from ldm.models.diffusion.plms import PLMSSampler
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from transformers import AutoFeatureExtractor
feature_extractor = AutoFeatureExtractor.from_pretrained("CompVis/stable-diffusion-v-1-3", use_auth_token=True)
safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-v-1-3", use_auth_token=True)
def chunk(it, size): def chunk(it, size):
it = iter(it) it = iter(it)
return iter(lambda: tuple(islice(it, size)), ()) 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): def load_model_from_config(config, ckpt, verbose=False):
print(f"Loading model from {ckpt}") print(f"Loading model from {ckpt}")
pl_sd = torch.load(ckpt, map_location="cpu") pl_sd = torch.load(ckpt, map_location="cpu")
@ -220,7 +237,9 @@ def main():
if opt.fixed_code: if opt.fixed_code:
start_code = torch.randn([opt.n_samples, opt.C, opt.H // opt.f, opt.W // opt.f], device=device) start_code = torch.randn([opt.n_samples, opt.C, opt.H // opt.f, opt.W // opt.f], device=device)
print("start code", start_code.abs().sum())
precision_scope = autocast if opt.precision=="autocast" else nullcontext precision_scope = autocast if opt.precision=="autocast" else nullcontext
precision_scope = nullcontext
with torch.no_grad(): with torch.no_grad():
with precision_scope("cuda"): with precision_scope("cuda"):
with model.ema_scope(): with model.ema_scope():
@ -269,7 +288,11 @@ def main():
Image.fromarray(grid.astype(np.uint8)).save(os.path.join(outpath, f'grid-{grid_count:04}.png')) Image.fromarray(grid.astype(np.uint8)).save(os.path.join(outpath, f'grid-{grid_count:04}.png'))
grid_count += 1 grid_count += 1
toc = time.time() image = x_samples_ddim.cpu().permute(0, 2, 3, 1).numpy()
# run safety checker
safety_checker_input = pipe.feature_extractor(numpy_to_pil(image), return_tensors="pt")
image, has_nsfw_concept = pipe.safety_checker(images=image, clip_input=safety_checker_input.pixel_values)
print(f"Your samples are ready and waiting for you here: \n{outpath} \n" print(f"Your samples are ready and waiting for you here: \n{outpath} \n"
f" \nEnjoy.") f" \nEnjoy.")