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

feat: extract safety checker, remove dependency to diffusers

Browse Source
This commit is contained in:
Manuel Schmid
2024-05-17 23:56:02 +02:00
parent 270ce71dbe
commit f338d5fc16
3 changed files with 129 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/censor.py
+3 -8
View File
@@ -1,10 +1,7 @@
# modified version of https://github.com/AUTOMATIC1111/stable-diffusion-webui-nsfw-censor/blob/master/scripts/censor.py
import numpy as np
import torch
import modules.core as core
from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from extras.diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
from transformers import AutoFeatureExtractor
from PIL import Image
import modules.config
@@ -16,8 +13,6 @@ safety_checker = None
def numpy_to_pil(image):
image = (image * 255).round().astype("uint8")
#pil_image = Image.fromarray(image, 'RGB')
pil_image = Image.fromarray(image)
return pil_image
@@ -27,7 +22,7 @@ def numpy_to_pil(image):
def check_safety(x_image):
global safety_feature_extractor, safety_checker
if safety_feature_extractor is None:
if safety_feature_extractor is None or safety_checker is None:
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)
@@ -52,4 +47,4 @@ def censor_single(x):
def censor_batch(images):
images = [censor_single(image) for image in images]
return images
return images