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 <hello@jamdon2.dev>
This commit is contained in:
Divided by Zer0 2022-10-05 14:55:04 +02:00 committed by GitHub
parent c1e97c8bcc
commit fe9fb596e5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 32 additions and 6 deletions

View File

@ -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
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 = []

View File

@ -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: