artem/stable-diffusion
1
0
Fork 0
mirror of https://github.com/CompVis/stable-diffusion.git synced 2026-04-02 15:12:32 +00:00
Code Issues Packages Projects Releases Wiki Activity

update safety model id

Browse Source
This commit is contained in:
Patrick Esser 2022-08-22 10:58:42 +00:00
parent a6e2f3b120
commit 673b0ab3a3
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
scripts/txt2img.py
View File

@ -20,9 +20,9 @@ from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionS
from transformers import AutoFeatureExtractor from transformers import AutoFeatureExtractor
# load safety model # load safety model
safety_model_id = "CompVis/stable-diffusion-v-1-3" safety_model_id = "CompVis/stable-diffusion-safety-checker"
safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id, use_auth_token=True) safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id)
safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id, use_auth_token=True) safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id)
def chunk(it, size): def chunk(it, size):
it = iter(it) it = iter(it)
@ -272,7 +272,7 @@ def main():
safety_checker_input = safety_feature_extractor(numpy_to_pil(x_image), return_tensors="pt") 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, 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) x_checked_image_torch = torch.from_numpy(x_checked_image).permute(0, 3, 1, 2)
if not opt.skip_save: if not opt.skip_save:
for x_sample in x_checked_image_torch: for x_sample in x_checked_image_torch: