diff --git a/prompts.json b/prompts.json index f714815..6a30a2e 100644 --- a/prompts.json +++ b/prompts.json @@ -1,20 +1,18 @@ { "Animagine": { "positive": "(masterpiece, exceptional, best aesthetic, best quality, masterpiece, extremely detailed:1.1)", - "positive_nsfw": "(masterpiece, exceptional, best aesthetic, best quality, masterpiece, extremely detailed, uncensored:1.1)", - "negative": "(lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts:1.15), (signature, watermark, username, gumroad username, patreon username, artist name:1.1), multiple views, text, error, blurry", - "negative_sfw": "(lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts:1.15), (rating:explicit, rating:questionable, rating:sensitive, signature, watermark, username, gumroad username, patreon username, artist name:1.1), multiple views, text, error, blurry" + "negative": "(lowres, bad anatomy, bad hands, missing fingers, extra digit, fewer digits, cropped, worst quality, low quality, normal quality, jpeg artifacts:1.15), (signature, watermark, username, gumroad username, patreon username, artist name:1.1), multiple views, text, error, blurry" }, "Pony": { "positive": "(score_9, score_8_up, score_7_up, score_6_up, source_anime, BREAK:1.1)", - "positive_nsfw": "(score_9, score_8_up, score_7_up, score_6_up, source_anime, BREAK, rating_explicit, uncensored:1.1)", - "negative": "score_4, score_3, score_2, score_1, score_0, censored, monochrome, 3d, signature, watermark, username, speech bubble", - "negative_sfw": "score_4, score_3, score_2, score_1, score_0, monochrome, 3d, signature, watermark, username, speech bubble, nude, explicit" + "negative": "score_4, score_3, score_2, score_1, score_0, censored, monochrome, 3d, signature, watermark, username, speech bubble" }, "Illust": { - "positive": "masterpiece, best quality, newest, absurdres, highres", - "positive_nsfw": "masterpiece, best quality, newest, absurdres, highres, uncensored", - "negative": "worst quality, old, early, low quality, lowres, signature, username, logo, bad hands, mutated hands, deformed hands, mammal, anthro, furry, ambiguous form, feral, semi-anthro", - "negative_sfw": "worst quality, old, early, low quality, lowres, signature, username, logo, bad hands, mutated hands, deformed hands, mammal, anthro, furry, ambiguous form, feral, semi-anthro, nude, explicit" + "positive": "masterpiece, best quality, absurdres, highres", + "negative": "worst quality, old, early, low quality, lowres, signature, username, logo, bad hands, mutated hands, deformed hands, mammal, anthro, furry, ambiguous form, feral, semi-anthro" + }, + "Anima": { + "positive": "masterpiece, best quality, score_7", + "negative": "worst quality, low quality, score_1, score_2, score_3, deviant art, artist name, bad hands, mutated hands, deformed hands" } } diff --git a/src/character_select.py b/src/character_select.py index 6e244e9..7972710 100644 --- a/src/character_select.py +++ b/src/character_select.py @@ -176,7 +176,7 @@ class CharacterDB: for char in custom_list: raw_name = char.get("name", "").strip() name = raw_name.lower() - copyright = char.get("copyright", "custom").strip().lower() + copyright = char.get("copyright", "custom").split(",")[0].strip().lower() if not raw_name: continue diff --git a/src/prompt_selector.py b/src/prompt_selector.py index a17edf2..ff584d7 100644 --- a/src/prompt_selector.py +++ b/src/prompt_selector.py @@ -1,6 +1,5 @@ import os import json -from transformers import pipeline class PromptSelectNode: @@ -18,11 +17,6 @@ class PromptSelectNode: with open(json_path, "r", encoding="utf-8") as file: cls.prompts = json.load(file) - def get_classifier(self): - if self.classifier is None: - self.classifier = pipeline("sentiment-analysis", model="michellejieli/NSFW_text_classifier") - return self.classifier - @classmethod def INPUT_TYPES(cls): cls.load_prompts() @@ -31,31 +25,58 @@ class PromptSelectNode: "required": { "prompt": ("STRING", {"default": ""}), "model": (model_names,), - "skip_classification": ("BOOLEAN", {"default": False}), + "separator": ("STRING", {"default": ", "}), }, } + NSFW_WORDS = [ + "uncensored", + "penis", + "sex", + "completely nude", + "genitals", + "vagina", + "pussy", + "dick", + "cock", + "asshole", + "masturbat", + "orgasm", + "ejaculat", + "threesome", + "gangbang", + "orgy", + "incest", + "squirt", + "creampie", + "futanari", + "erection", + "dildo", + "vibrator", + "pegging", + ] + + def is_nsfw(self, prompt): + return any(word.lower() in prompt.lower() for word in self.NSFW_WORDS if word) + # Define static properties for the class RETURN_TYPES = ("STRING", "STRING") RETURN_NAMES = ("POSITIVE", "NEGATIVE") FUNCTION = "find_chosen_prompts" CATEGORY = "sn0w" - def find_chosen_prompts(self, prompt, model, skip_classification): - if skip_classification: - positive_prompt = self.prompts.get(model, {}).get("positive", "") - negative_prompt = self.prompts.get(model, {}).get("negative", "") - return (positive_prompt, negative_prompt) + def find_chosen_prompts(self, prompt, model, separator: str): + positive_prompt = self.prompts.get(model, {}).get("positive", "") + negative_prompt = self.prompts.get(model, {}).get("negative", "") - classifier = self.get_classifier() - result = classifier(prompt) - - is_nsfw = result[0]["label"] == "NSFW" + is_nsfw = self.is_nsfw(prompt) if is_nsfw: - positive_prompt = self.prompts.get(model, {}).get("positive_nsfw", "") - negative_prompt = self.prompts.get(model, {}).get("negative", "") - else: - positive_prompt = self.prompts.get(model, {}).get("positive", "") - negative_prompt = self.prompts.get(model, {}).get("negative_sfw", "") + if "uncensored" not in prompt: + positive_prompt = positive_prompt + separator + "uncensored" + negative_prompt = ( + negative_prompt + + separator + + separator.join(["censored", "bar censor", "blur censor", "mosaic censoring"]) + ) return (positive_prompt, negative_prompt)