NSFW filter setting

This commit is contained in:
cmdr2 2023-02-18 15:01:13 +05:30
parent e7a2dfa57f
commit aa70f2849b
7 changed files with 22 additions and 4 deletions

View File

@ -1,12 +1,11 @@
import os import os
from easydiffusion import app, device_manager from easydiffusion import app
from easydiffusion.types import TaskData from easydiffusion.types import TaskData
from easydiffusion.utils import log from easydiffusion.utils import log
from sdkit import Context from sdkit import Context
from sdkit.models import load_model, unload_model, get_model_info_from_db, scan_model from sdkit.models import load_model, unload_model, scan_model
from sdkit.utils import hash_file_quick
KNOWN_MODEL_TYPES = ["stable-diffusion", "vae", "hypernetwork", "gfpgan", "realesrgan"] KNOWN_MODEL_TYPES = ["stable-diffusion", "vae", "hypernetwork", "gfpgan", "realesrgan"]
MODEL_EXTENSIONS = { MODEL_EXTENSIONS = {
@ -102,6 +101,7 @@ def reload_models_if_necessary(context: Context, task_data: TaskData):
"hypernetwork": task_data.use_hypernetwork_model, "hypernetwork": task_data.use_hypernetwork_model,
"gfpgan": task_data.use_face_correction, "gfpgan": task_data.use_face_correction,
"realesrgan": task_data.use_upscale, "realesrgan": task_data.use_upscale,
"nsfw_checker": True if task_data.block_nsfw else None,
} }
models_to_reload = { models_to_reload = {
model_type: path model_type: path

View File

@ -102,15 +102,20 @@ def generate_images_internal(
def filter_images(task_data: TaskData, images: list, user_stopped): def filter_images(task_data: TaskData, images: list, user_stopped):
if user_stopped or (task_data.use_face_correction is None and task_data.use_upscale is None): if user_stopped:
return images return images
filters_to_apply = [] filters_to_apply = []
if task_data.block_nsfw:
filters_to_apply.append("nsfw_checker")
if task_data.use_face_correction and "gfpgan" in task_data.use_face_correction.lower(): if task_data.use_face_correction and "gfpgan" in task_data.use_face_correction.lower():
filters_to_apply.append("gfpgan") filters_to_apply.append("gfpgan")
if task_data.use_upscale and "realesrgan" in task_data.use_upscale.lower(): if task_data.use_upscale and "realesrgan" in task_data.use_upscale.lower():
filters_to_apply.append("realesrgan") filters_to_apply.append("realesrgan")
if len(filters_to_apply) == 0:
return images
return apply_filters(context, filters_to_apply, images, scale=task_data.upscale_amount) return apply_filters(context, filters_to_apply, images, scale=task_data.upscale_amount)

View File

@ -38,6 +38,7 @@ class TaskData(BaseModel):
use_hypernetwork_model: str = None use_hypernetwork_model: str = None
show_only_filtered_image: bool = False show_only_filtered_image: bool = False
block_nsfw: bool = False
output_format: str = "jpeg" # or "png" output_format: str = "jpeg" # or "png"
output_quality: int = 75 output_quality: int = 75
metadata_output_format: str = "txt" # or "json" metadata_output_format: str = "txt" # or "json"

View File

@ -30,6 +30,7 @@ const SETTINGS_IDS_LIST = [
"gfpgan_model", "gfpgan_model",
"use_upscale", "use_upscale",
"upscale_amount", "upscale_amount",
"block_nsfw",
"show_only_filtered_image", "show_only_filtered_image",
"upscale_model", "upscale_model",
"preview-image", "preview-image",

View File

@ -741,6 +741,7 @@
"stream_progress_updates": true, "stream_progress_updates": true,
"stream_image_progress": true, "stream_image_progress": true,
"show_only_filtered_image": true, "show_only_filtered_image": true,
"block_nsfw": false,
"output_format": "png", "output_format": "png",
"output_quality": 75, "output_quality": 75,
} }

View File

@ -43,6 +43,7 @@ let hypernetworkModelField = new ModelDropdown(document.querySelector('#hypernet
let hypernetworkStrengthSlider = document.querySelector('#hypernetwork_strength_slider') let hypernetworkStrengthSlider = document.querySelector('#hypernetwork_strength_slider')
let hypernetworkStrengthField = document.querySelector('#hypernetwork_strength') let hypernetworkStrengthField = document.querySelector('#hypernetwork_strength')
let outputFormatField = document.querySelector('#output_format') let outputFormatField = document.querySelector('#output_format')
let blockNSFWField = document.querySelector('#block_nsfw')
let showOnlyFilteredImageField = document.querySelector("#show_only_filtered_image") let showOnlyFilteredImageField = document.querySelector("#show_only_filtered_image")
let updateBranchLabel = document.querySelector("#updateBranchLabel") let updateBranchLabel = document.querySelector("#updateBranchLabel")
let streamImageProgressField = document.querySelector("#stream_image_progress") let streamImageProgressField = document.querySelector("#stream_image_progress")
@ -967,6 +968,7 @@ function getCurrentUserRequest() {
stream_progress_updates: true, stream_progress_updates: true,
stream_image_progress: (numOutputsTotal > 50 ? false : streamImageProgressField.checked), stream_image_progress: (numOutputsTotal > 50 ? false : streamImageProgressField.checked),
show_only_filtered_image: showOnlyFilteredImageField.checked, show_only_filtered_image: showOnlyFilteredImageField.checked,
block_nsfw: blockNSFWField.checked,
output_format: outputFormatField.value, output_format: outputFormatField.value,
output_quality: parseInt(outputQualityField.value), output_quality: parseInt(outputQualityField.value),
metadata_output_format: metadataOutputFormatField.value, metadata_output_format: metadataOutputFormatField.value,

View File

@ -78,6 +78,14 @@ var PARAMETERS = [
} }
], ],
}, },
{
id: "block_nsfw",
type: ParameterType.checkbox,
label: "Block NSFW images",
note: "blurs out NSFW images",
icon: "fa-land-mine-on",
default: false,
},
{ {
id: "sound_toggle", id: "sound_toggle",
type: ParameterType.checkbox, type: ParameterType.checkbox,