Foreign-Mirrors/Fooocus
1
0
Fork 0
mirror of https://github.com/lllyasviel/Fooocus.git synced 2026-08-16 13:13:16 +02:00
Code Issues Packages Projects Releases Wiki Activity

use config to set cache dir for safety checker

Browse Source
This commit is contained in:
Manuel Schmid
2023-12-16 19:25:02 +01:00
committed by Manuel Schmid
parent ea764855f8
commit cdaaec3e71
4 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/censor.py
+3 -2
View File
@@ -7,6 +7,7 @@ import modules.core as core
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from transformers import AutoFeatureExtractor
from PIL import Image
import modules.config
safety_model_id = "CompVis/stable-diffusion-safety-checker"
safety_feature_extractor = None
@@ -27,8 +28,8 @@ def check_safety(x_image):
global safety_feature_extractor, safety_checker
if safety_feature_extractor is None:
safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id)
safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id)
safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id, cache_dir=modules.config.path_safety_checker_models)
safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id, cache_dir=modules.config.path_safety_checker_models)
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)