From fe9fb596e572f266b02431caf6a1daf5f81eeb4d Mon Sep 17 00:00:00 2001 From: Divided by Zer0 Date: Wed, 5 Oct 2022 14:55:04 +0200 Subject: [PATCH] Adds more NSFW controls to the horde worker (#1426) Allows worker to specify a word blacklist which it will refuse to pickup Allows the worker to specify a word censorlist, against which it will always apply the NSFW filter, even if the worker accepts NSFW. Co-authored-by: hlky <106811348+hlky@users.noreply.github.com> Co-authored-by: lukas5450 <46075099+lukas5450@users.noreply.github.com> Co-authored-by: JamDon2 --- scripts/bridgeData_template.py | 6 +++++- scripts/webui.py | 32 +++++++++++++++++++++++++++----- 2 files changed, 32 insertions(+), 6 deletions(-) diff --git a/scripts/bridgeData_template.py b/scripts/bridgeData_template.py index 365ce98..483ab29 100644 --- a/scripts/bridgeData_template.py +++ b/scripts/bridgeData_template.py @@ -14,4 +14,8 @@ horde_priority_usernames = [] # If you put this to 32, it is equivalent to 1024x1024 pixels horde_max_power = 8 # Set this to false, if you do not want your worker to receive requests for NSFW generations -horde_nsfw = True \ No newline at end of file +horde_nsfw = True +# A list of words which you do not want to your worker to accept +horde_blacklist = [] +# A list of words for which you always want to allow the NSFW censor filter, even when this worker is in NSFW mode +horde_censorlist = [] diff --git a/scripts/webui.py b/scripts/webui.py index 7c373e4..263f9d7 100644 --- a/scripts/webui.py +++ b/scripts/webui.py @@ -71,7 +71,10 @@ parser.add_argument('--horde_name', action="store", required=False, type=str, he parser.add_argument('--horde_url', action="store", required=False, type=str, help="The SH Horde URL. Where the bridge will pickup prompts and send the finished generations.") parser.add_argument('--horde_priority_usernames',type=str, action='append', required=False, help="Usernames which get priority use in this horde instance. The owner's username is always in this list.") parser.add_argument('--horde_max_power',type=int, required=False, help="How much power this instance has to generate pictures. Min: 2") -parser.add_argument('--horde_nsfw', action='store_true', required=False, help="Set to false if you do not want this worker generating NSFW images.") +parser.add_argument('--horde_sfw', action='store_true', required=False, help="Set to true if you do not want this worker generating NSFW images.") +parser.add_argument('--horde_blacklist', nargs='+', required=False, help="List the words that you want to blacklist.") +parser.add_argument('--horde_censorlist', nargs='+', required=False, help="List the words that you want to censor.") +parser.add_argument('--horde_censor_nsfw', action='store_true', required=False, help="Set to true if you want this bridge worker to censor NSFW images.") opt = parser.parse_args() #Should not be needed anymore @@ -2649,7 +2652,7 @@ def run_headless(): print() @logger.catch -def run_bridge(interval, api_key, horde_name, horde_url, priority_usernames, horde_max_pixels, horde_nsfw): +def run_bridge(interval, api_key, horde_name, horde_url, priority_usernames, horde_max_pixels, horde_nsfw, horde_censor_nsfw, horde_blacklist, horde_censorlist): current_id = None current_payload = None loop_retry = 0 @@ -2667,6 +2670,7 @@ def run_bridge(interval, api_key, horde_name, horde_url, priority_usernames, hor "max_pixels": horde_max_pixels, "priority_usernames": priority_usernames, "nsfw": horde_nsfw, + "blacklist": horde_blacklist, } headers = {"apikey": api_key} if current_id: @@ -2719,9 +2723,15 @@ def run_bridge(interval, api_key, horde_name, horde_url, priority_usernames, hor loop_retry = 0 time.sleep(10) continue + current_payload['toggles'] = current_payload.get('toggles', [1,4]) # In bridge-mode, matrix is prepared on the horde and split in multiple nodes - if 'toggles' in current_payload and 0 in current_payload['toggles']: + if 0 in current_payload['toggles']: current_payload['toggles'].remove(0) + if 8 not in current_payload['toggles']: + if horde_censor_nsfw and not horde_nsfw: + current_payload['toggles'].append(8) + elif any(word in current_payload['prompt'] for word in horde_censorlist): + current_payload['toggles'].append(8) images, seed, info, stats = txt2img(**current_payload) buffer = BytesIO() # We send as WebP to avoid using all the horde bandwidth @@ -2795,15 +2805,27 @@ if __name__ == '__main__': horde_priority_usernames = opt.horde_priority_usernames if opt.horde_priority_usernames else cd.horde_priority_usernames horde_max_power = opt.horde_max_power if opt.horde_max_power else cd.horde_max_power try: - horde_nsfw = opt.horde_nsfw if opt.horde_nsfw else cd.horde_nsfw + horde_nsfw = not opt.horde_sfw if opt.horde_sfw else cd.horde_nsfw except AttributeError: horde_nsfw = True + try: + horde_censor_nsfw = opt.horde_censor_nsfw if opt.horde_censor_nsfw else cd.horde_censor_nsfw + except AttributeError: + horde_censor_nsfw = False + try: + horde_blacklist = opt.horde_blacklist if opt.horde_blacklist else cd.horde_blacklist + except AttributeError: + horde_blacklist = [] + try: + horde_censorlist = opt.horde_censorlist if opt.horde_censorlist else cd.horde_censorlist + except AttributeError: + horde_censorlist = [] if horde_max_power < 2: horde_max_power = 2 horde_max_pixels = 64*64*8*horde_max_power logger.info(f"Joining Horde with parameters: API Key '{horde_api_key}'. Server Name '{horde_name}'. Horde URL '{horde_url}'. Max Pixels {horde_max_pixels}") try: - run_bridge(1, horde_api_key, horde_name, horde_url, horde_priority_usernames, horde_max_pixels, horde_nsfw) + run_bridge(1, horde_api_key, horde_name, horde_url, horde_priority_usernames, horde_max_pixels, horde_nsfw, horde_censor_nsfw, horde_blacklist, horde_censorlist) except KeyboardInterrupt: logger.info(f"Keyboard Interrupt Received. Ending Bridge") else: