Upload NAIA_generation.py
Browse files- NAIA_generation.py +1124 -0
NAIA_generation.py
ADDED
@@ -0,0 +1,1124 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import requests
|
2 |
+
import random
|
3 |
+
from PIL import Image, ImageOps, ImageDraw, ImageFont
|
4 |
+
from datetime import datetime
|
5 |
+
from pathlib import Path
|
6 |
+
import io, re, base64, os
|
7 |
+
import zipfile
|
8 |
+
import math
|
9 |
+
|
10 |
+
BASE_URL="https://image.novelai.net"
|
11 |
+
|
12 |
+
def process_xargs(xargs):
|
13 |
+
processed_xargs = []
|
14 |
+
skip_until = None # ์ด ๋ณ์๋ ํ์ฌ ์ฒ๋ฆฌ ์ค์ธ '<...>' ๊ตฌ๋ฌธ์ ๋ ์ธ๋ฑ์ค๋ฅผ ์ถ์ ํฉ๋๋ค.
|
15 |
+
|
16 |
+
for i, elem in enumerate(xargs):
|
17 |
+
if skip_until is not None and i <= skip_until:
|
18 |
+
continue # ์ด๋ฏธ ์ฒ๋ฆฌ๋ ๋ฒ์๋ฅผ ๊ฑด๋๋ฐ๊ธฐ
|
19 |
+
|
20 |
+
if "add_negative:<" in elem or "rem_negative:<" in elem:
|
21 |
+
# '<'๋ก ์์ํ์ฌ '>'๋ก ๋๋๋ ๊ตฌ๋ฌธ ์ฒ๋ฆฌ
|
22 |
+
if elem.endswith(">"):
|
23 |
+
processed_xargs.append(elem)
|
24 |
+
else:
|
25 |
+
combined_elem = elem
|
26 |
+
for j in range(i + 1, len(xargs)):
|
27 |
+
combined_elem += ',' + xargs[j] # ','๋ฅผ ํฌํจํ์ฌ ์์๋ค์ ๊ฒฐํฉ
|
28 |
+
if xargs[j].endswith(">"):
|
29 |
+
skip_until = j # '>'๋ก ๋๋๋ ์์๋ฅผ ์ฐพ์ผ๋ฉด, ์ฒ๋ฆฌํ ๋ฒ์์ ๋ ์ธ๋ฑ์ค๋ฅผ ์
๋ฐ์ดํธ
|
30 |
+
break
|
31 |
+
processed_xargs.append(combined_elem)
|
32 |
+
else:
|
33 |
+
processed_xargs.append(elem) # ์ผ๋ฐ์ ์ธ ์์๋ ๊ทธ๋๋ก ๋ฆฌ์คํธ์ ์ถ๊ฐ
|
34 |
+
|
35 |
+
return processed_xargs
|
36 |
+
|
37 |
+
def event_cond_negative(_prompt, negative, user_input, rating):
|
38 |
+
def insert_text_after_keyword(negative, user_keyword, user_additional_keyword):
|
39 |
+
if user_keyword in negative:
|
40 |
+
index = negative.index(user_keyword) + 1
|
41 |
+
negative.insert(index, user_additional_keyword)
|
42 |
+
return negative
|
43 |
+
def parse_conditional_command(command):
|
44 |
+
match = re.match(r"\((.*?)\)\:(.*)", command)
|
45 |
+
if match:
|
46 |
+
return match.groups()
|
47 |
+
return '', command
|
48 |
+
def check_condition(prompt, condition, rating):
|
49 |
+
if not condition:
|
50 |
+
return True
|
51 |
+
sub_conditions = re.split(r'\)\s*&\s*\(', condition)
|
52 |
+
sub_conditions = [re.sub(r'^\(|\)$', '', cond) for cond in sub_conditions]
|
53 |
+
|
54 |
+
results = []
|
55 |
+
for sub_cond in sub_conditions:
|
56 |
+
if '&' in sub_cond:
|
57 |
+
results.append(all(check_condition(prompt, cond, rating) for cond in sub_cond.split('&')))
|
58 |
+
elif '|' in sub_cond:
|
59 |
+
results.append(any(check_condition(prompt, cond, rating) for cond in sub_cond.split('|')))
|
60 |
+
else:
|
61 |
+
if sub_cond in ['e', 'q', 's', 'g']:
|
62 |
+
results.append(sub_cond == rating)
|
63 |
+
elif sub_cond in ['~e', '~q', '~s', '~g']:
|
64 |
+
results.append(sub_cond != rating)
|
65 |
+
# PM
|
66 |
+
elif sub_cond.startswith('*'):
|
67 |
+
results.append(sub_cond[1:] in prompt)
|
68 |
+
# NOT IN
|
69 |
+
elif sub_cond.startswith('~!'):
|
70 |
+
results.append(sub_cond[2:] not in prompt)
|
71 |
+
elif sub_cond.startswith('~'):
|
72 |
+
results.append(any(sub_cond[1:] not in element for element in prompt))
|
73 |
+
# CONTAIN
|
74 |
+
else:
|
75 |
+
results.append(any(sub_cond in element for element in prompt))
|
76 |
+
return all(results)
|
77 |
+
def execute_command(negative, command):
|
78 |
+
if '+=' in command:
|
79 |
+
keyword, addition = command.split('+=', 1)
|
80 |
+
addition = addition.replace('^', ', ')
|
81 |
+
return insert_text_after_keyword(negative, keyword, addition)
|
82 |
+
elif command.startswith('add '):
|
83 |
+
keyword = command[4:]
|
84 |
+
keyword = keyword.replace('^', ', ')
|
85 |
+
keys = keyword.split(',')
|
86 |
+
keys = [key.strip() for key in keys]
|
87 |
+
for key in keys:
|
88 |
+
if key not in negative:
|
89 |
+
negative.append(key)
|
90 |
+
return negative
|
91 |
+
elif command.startswith('rem '):
|
92 |
+
keyword = command[4:]
|
93 |
+
keyword = keyword.replace('^', ', ')
|
94 |
+
keys = keyword.split(',')
|
95 |
+
keys = [key.strip() for key in keys]
|
96 |
+
for key in keys:
|
97 |
+
if key in negative:
|
98 |
+
negative.remove(key)
|
99 |
+
return negative
|
100 |
+
elif '=' in command:
|
101 |
+
keyword, replacement = command.split('=', 1)
|
102 |
+
if keyword in negative:
|
103 |
+
replacement = replacement.replace('^', ', ')
|
104 |
+
index = negative.index(keyword)
|
105 |
+
negative[index] = replacement
|
106 |
+
return negative
|
107 |
+
negative = negative.split(',')
|
108 |
+
negative = [neg.strip() for neg in negative]
|
109 |
+
prompt = _prompt.split(',')
|
110 |
+
prompt = [key.strip() for key in prompt]
|
111 |
+
commands = [cmd.strip() for cmd in user_input.split(',') if not cmd.strip().startswith('#')]
|
112 |
+
for command in commands:
|
113 |
+
condition, cmd = parse_conditional_command(command)
|
114 |
+
if check_condition(prompt, condition, rating):
|
115 |
+
negative = execute_command(negative, cmd)
|
116 |
+
return ', '.join(negative)
|
117 |
+
|
118 |
+
def image_to_base64(image):
|
119 |
+
image_bytesIO = io.BytesIO()
|
120 |
+
image.save(image_bytesIO, format="PNG")
|
121 |
+
return base64.b64encode(image_bytesIO.getvalue()).decode()
|
122 |
+
|
123 |
+
def make_turbo_prompt(gen_request):
|
124 |
+
lines = gen_request['prompt']
|
125 |
+
result = {
|
126 |
+
"boys": False,
|
127 |
+
"girls": False,
|
128 |
+
"1girl": False,
|
129 |
+
"1boy": False,
|
130 |
+
"1other": False,
|
131 |
+
"others": False
|
132 |
+
}
|
133 |
+
state = {
|
134 |
+
"nude,": False,
|
135 |
+
"pov,": False,
|
136 |
+
"cum,": False,
|
137 |
+
"after ": False,
|
138 |
+
"pussy juice": False,
|
139 |
+
"barefoot": False,
|
140 |
+
"breasts": False,
|
141 |
+
"ejaculation": False,
|
142 |
+
}
|
143 |
+
|
144 |
+
def insert_spaces(source_list, reference_list):
|
145 |
+
modified_list = source_list.copy()
|
146 |
+
for index, keyword in enumerate(reference_list):
|
147 |
+
if keyword not in source_list:
|
148 |
+
space_count = len(keyword) # ํค์๋ ๊ธธ์ด๋งํผ์ ๊ณต๋ฐฑ ๋ฌธ์
|
149 |
+
modified_list.insert(index, ' ' * space_count)
|
150 |
+
return modified_list
|
151 |
+
|
152 |
+
keywords = gen_request['prompt'].split(', ')
|
153 |
+
filtered_keywords = []
|
154 |
+
removed_indices = []
|
155 |
+
positive0, positive1, positive2, positive3 = gen_request.copy(),gen_request.copy(),gen_request.copy(),gen_request.copy()
|
156 |
+
|
157 |
+
for word in result.keys():
|
158 |
+
if word in lines:
|
159 |
+
result[word] = True
|
160 |
+
for word in state.keys():
|
161 |
+
if word in gen_request['prompt']:
|
162 |
+
state[word] = True
|
163 |
+
|
164 |
+
key_index = int((len(keywords)/2)-1)
|
165 |
+
|
166 |
+
if(result["1boy"]) or (result["boys"]):
|
167 |
+
if(result["1girl"]):
|
168 |
+
if(', sex' in gen_request['prompt'] or 'group sex' in gen_request['prompt']):
|
169 |
+
sex_pos_keywords = ['stomach bulge','insertion', 'fucked silly', 'x-ray', 'orgasm', 'cross-section', 'uterus', 'overflow', 'rape', 'vaginal', 'anal']
|
170 |
+
facial_keywords = ['tongue','ahegao']
|
171 |
+
temp_sex_pos = []
|
172 |
+
temp_facial = []
|
173 |
+
cum_events = []
|
174 |
+
explicit_check = []
|
175 |
+
if 'open mouth' in keywords: keywords.remove('open mouth')
|
176 |
+
if 'closed mouth' in keywords: keywords.remove('closed mouth')
|
177 |
+
if 'after rape' in keywords:
|
178 |
+
keywords.remove('after rape')
|
179 |
+
explicit_check.append('after rape')
|
180 |
+
if 'used condom' in keywords:
|
181 |
+
keywords.remove('used condom')
|
182 |
+
explicit_check.append('used condom')
|
183 |
+
for keyword in keywords:
|
184 |
+
if ('sex' not in keyword and 'cum' not in keyword and 'ejaculation' not in keyword and 'vaginal' not in keyword and 'penetration' not in keyword) and all(sex_pos not in keyword for sex_pos in sex_pos_keywords) and all(facial not in keyword for facial in facial_keywords):
|
185 |
+
filtered_keywords.append(keyword)
|
186 |
+
elif 'sex' in keyword:
|
187 |
+
removed_indices.append(keyword)
|
188 |
+
elif 'penetration' in keyword:
|
189 |
+
removed_indices.append(keyword)
|
190 |
+
elif 'cum' in keyword and keyword != 'cum':
|
191 |
+
cum_events.append(keyword)
|
192 |
+
elif any(sex_pos in keyword for sex_pos in sex_pos_keywords):
|
193 |
+
for sex_pos in sex_pos_keywords:
|
194 |
+
if sex_pos in keyword:
|
195 |
+
temp_sex_pos.append(sex_pos)
|
196 |
+
elif any(facial not in keyword for facial in facial_keywords):
|
197 |
+
for facial in facial_keywords:
|
198 |
+
if facial in keyword:
|
199 |
+
temp_facial.append(facial)
|
200 |
+
filtered_keywords.insert(int((len(filtered_keywords)/2)-1), ' no penetration, imminent penetration')
|
201 |
+
filtered_keywords_positive0 = filtered_keywords.copy()
|
202 |
+
filtered_keywords.remove(' no penetration, imminent penetration')
|
203 |
+
#0 imminent penetration, imminent sex
|
204 |
+
if 'condom' in filtered_keywords and 'condom on penis' not in filtered_keywords:
|
205 |
+
t_index = filtered_keywords.index('condom')
|
206 |
+
rand_num = random.randint(0,2)
|
207 |
+
if rand_num == 1: filtered_keywords.insert(t_index, 'condom on penis')
|
208 |
+
for i, keyword in enumerate(filtered_keywords):
|
209 |
+
if 'pantyhose' in keyword:
|
210 |
+
filtered_keywords[i] = 'torn ' + filtered_keywords[i]
|
211 |
+
#1 default
|
212 |
+
key_index = int((len(filtered_keywords)/2)-1)
|
213 |
+
if 'pussy' in filtered_keywords: key_index = filtered_keywords.index('pussy')
|
214 |
+
if 'penis' in filtered_keywords: key_index = filtered_keywords.index('penis')
|
215 |
+
filtered_keywords[key_index:key_index] = ['motion lines', 'surprised']
|
216 |
+
for keyword in removed_indices:
|
217 |
+
if 'cum' not in keyword and 'ejaculation' not in keyword:
|
218 |
+
filtered_keywords.insert(key_index,keyword)
|
219 |
+
if(temp_sex_pos): filtered_keywords[key_index:key_index] = temp_sex_pos
|
220 |
+
if 'group sex' in filtered_keywords and 'sex' not in filtered_keywords:
|
221 |
+
t_index = filtered_keywords.index('group sex')
|
222 |
+
filtered_keywords.insert(t_index, 'sex')
|
223 |
+
if('clothed sex' in filtered_keywords and not 'bottomless' in filtered_keywords): filtered_keywords.insert(filtered_keywords.index('clothed sex')+1, 'bottomless')
|
224 |
+
pos1_copied_keywords = filtered_keywords.copy()
|
225 |
+
for i, keyword in enumerate(pos1_copied_keywords):
|
226 |
+
if 'closed eyes' in keyword:
|
227 |
+
rand_num = random.randint(0,2)
|
228 |
+
if(rand_num == 0): pos1_copied_keywords[i] = 'half-' + pos1_copied_keywords[i]
|
229 |
+
elif(rand_num == 1 and 'closed eyes' in pos1_copied_keywords):
|
230 |
+
pos1_copied_keywords.remove('closed eyes')
|
231 |
+
filtered_keywords[i] = 'half-closed eyes'
|
232 |
+
filtered_keywords_positive1 = pos1_copied_keywords.copy()
|
233 |
+
#2 ejaculation,cum in pussy
|
234 |
+
key_index = filtered_keywords.index('surprised')
|
235 |
+
filtered_keywords.remove('surprised')
|
236 |
+
filtered_keywords[key_index:key_index] = ["ejaculation","cum"] if "condom on penis" not in filtered_keywords else ["twitching penis", "[[[[orgasm]]]]"]
|
237 |
+
for keyword in removed_indices:
|
238 |
+
if 'cum' in keyword:
|
239 |
+
filtered_keywords.insert(key_index,keyword)
|
240 |
+
if(temp_facial): filtered_keywords[key_index:key_index] =temp_facial
|
241 |
+
filtered_keywords_positive2 = filtered_keywords.copy()
|
242 |
+
#3 after sex, after ejaculation
|
243 |
+
for i, keyword in enumerate(filtered_keywords):
|
244 |
+
if 'closed eyes' in keyword:
|
245 |
+
rand_num = random.randint(0,2)
|
246 |
+
if(rand_num == 0 and filtered_keywords[i] != 'half-closed eyes'): filtered_keywords[i] = 'half-' + filtered_keywords[i]
|
247 |
+
elif(rand_num == 1): filtered_keywords[i] = 'empty eyes'
|
248 |
+
else: filtered_keywords[i] = 'empty eyes, half-closed eyes'
|
249 |
+
if 'sex' in filtered_keywords:
|
250 |
+
key_index = filtered_keywords.index('sex')
|
251 |
+
elif 'group sex' in filtered_keywords:
|
252 |
+
key_index = filtered_keywords.index('group sex')
|
253 |
+
if "condom on penis" not in filtered_keywords: filtered_keywords.remove('ejaculation')
|
254 |
+
else:
|
255 |
+
filtered_keywords.remove('twitching penis')
|
256 |
+
filtered_keywords.remove('[[[[orgasm]]]]')
|
257 |
+
filtered_keywords[key_index:key_index] = ['cum drip', 'erection'] + cum_events if "condom on penis" not in filtered_keywords else ["used condom", "{{used condom on penis}}"]
|
258 |
+
if(explicit_check): filtered_keywords[key_index:key_index] = explicit_check
|
259 |
+
if 'sex' in filtered_keywords and 'group sex' not in filtered_keywords:
|
260 |
+
if('pussy' in filtered_keywords and not 'anal' in filtered_keywords):
|
261 |
+
if "condom on penis" not in filtered_keywords: filtered_keywords.insert(filtered_keywords.index('sex')+1, 'after vaginal, spread pussy')
|
262 |
+
else: filtered_keywords.insert(filtered_keywords.index('sex')+1, 'after vaginal, spread pussy, pussy juice puddle')
|
263 |
+
elif('anal' in filtered_keywords):
|
264 |
+
if "condom on penis" not in filtered_keywords: filtered_keywords.insert(filtered_keywords.index('sex')+1, 'after anal, cum in ass')
|
265 |
+
else: filtered_keywords.insert(filtered_keywords.index('sex')+1, 'after anal, pussy juice')
|
266 |
+
filtered_keywords.insert(filtered_keywords.index('sex'), 'after sex')
|
267 |
+
filtered_keywords.remove('sex')
|
268 |
+
elif 'group sex' in filtered_keywords:
|
269 |
+
if('vaginal' in filtered_keywords and not 'anal' in filtered_keywords):
|
270 |
+
filtered_keywords.insert(filtered_keywords.index('group sex')+1, 'after vaginal, spread pussy')
|
271 |
+
if 'multiple penises' in filtered_keywords:
|
272 |
+
if "condom on penis" not in filtered_keywords: filtered_keywords.insert(filtered_keywords.index('group sex')+3, 'cum on body, bukkake')
|
273 |
+
else: filtered_keywords.insert(filtered_keywords.index('group sex')+3, 'pussy juice, pussy juice puddle')
|
274 |
+
elif('anal' in filtered_keywords):
|
275 |
+
if "condom on penis" not in filtered_keywords: filtered_keywords.insert(filtered_keywords.index('group sex')+1, 'after anus, cum in ass')
|
276 |
+
else: filtered_keywords.insert(filtered_keywords.index('group sex')+1, 'after anus')
|
277 |
+
if 'multiple penises' in filtered_keywords:
|
278 |
+
if "condom on penis" not in filtered_keywords: filtered_keywords.insert(filtered_keywords.index('group sex')+3, 'cum on body, bukkake')
|
279 |
+
else: filtered_keywords.insert(filtered_keywords.index('group sex')+3, 'pussy juice')
|
280 |
+
else: filtered_keywords.insert(filtered_keywords.index('group sex')+1, 'cum on body, {bukkake}')
|
281 |
+
temp_post_keyword = []
|
282 |
+
for keyword in sex_pos_keywords:
|
283 |
+
if not (keyword == 'orgasm' or keyword == 'overflow'):
|
284 |
+
if keyword in filtered_keywords:
|
285 |
+
temp_post_keyword.append(keyword)
|
286 |
+
for keyword in temp_post_keyword:
|
287 |
+
filtered_keywords.remove(keyword)
|
288 |
+
|
289 |
+
positive0['prompt'] = ', '.join(insert_spaces(filtered_keywords_positive0, filtered_keywords)).strip()
|
290 |
+
positive1['prompt'] = ', '.join(insert_spaces(filtered_keywords_positive1, filtered_keywords)).strip()
|
291 |
+
positive2['prompt'] = ', '.join(insert_spaces(filtered_keywords_positive2, filtered_keywords)).strip()
|
292 |
+
positive3['prompt'] = ', '.join(filtered_keywords).strip()
|
293 |
+
positive0["type"] = "turbo"
|
294 |
+
positive1["type"] = "turbo"
|
295 |
+
positive2["type"] = "turbo"
|
296 |
+
positive3["type"] = "turbo"
|
297 |
+
return positive0, positive1, positive2, positive3
|
298 |
+
|
299 |
+
def generate_image(access_token, prompt, model, action, parameters):
|
300 |
+
if re.match(r'^http[s]?://', access_token):
|
301 |
+
return generate_image_webui(access_token, prompt, model, action, parameters)
|
302 |
+
return generate_image_NAI(access_token, prompt, model, action, parameters)
|
303 |
+
|
304 |
+
def generate_image_NAI(access_token, prompt, model, action, parameters):
|
305 |
+
data = {
|
306 |
+
"input": prompt,
|
307 |
+
"model": model,
|
308 |
+
"action": action,
|
309 |
+
"parameters": parameters,
|
310 |
+
}
|
311 |
+
|
312 |
+
if data['parameters']['qualityToggle'] == True:
|
313 |
+
if "nai-diffusion-furry-3" not in data['model']:
|
314 |
+
data['input'] += ', best quality, amazing quality, very aesthetic, absurdres'
|
315 |
+
else:
|
316 |
+
data['input'] += ', {best quality}, {amazing quality}'
|
317 |
+
|
318 |
+
|
319 |
+
if data['parameters']['sampler'] not in ["k_euler", "k_euler_ancestral", "k_dpmpp_sde", "k_dpmpp_2s_ancestral", "k_dpmpp_2m", "k_dpmpp_2m_sde"]:
|
320 |
+
data['parameters']['sampler'] = "k_euler_ancestral"
|
321 |
+
|
322 |
+
if data['parameters']['cfg_rescale'] > 1:
|
323 |
+
data['parameters']['cfg_rescale'] = 0
|
324 |
+
|
325 |
+
response = requests.post(f"{BASE_URL}/ai/generate-image", json=data, headers={ "Authorization": f"Bearer {access_token}" })
|
326 |
+
# catch any errors
|
327 |
+
return response.content
|
328 |
+
|
329 |
+
def augment_image_NAI(gen_request):
|
330 |
+
def resize_and_fill(image, max_size=None):
|
331 |
+
if max_size is None:
|
332 |
+
max_size = gen_request["user_screen_size"]
|
333 |
+
original_width, original_height = image.size
|
334 |
+
if original_width > max_size or original_height > max_size:
|
335 |
+
# ๋น์จ์ ์ ์งํ๋ฉด์ ํฌ๊ธฐ ์กฐ์
|
336 |
+
image.thumbnail((max_size, max_size))
|
337 |
+
|
338 |
+
# ์ ์ด๋ฏธ์ง ํฌ๊ธฐ ๊ณ์ฐ
|
339 |
+
width, height = image.size
|
340 |
+
new_image = Image.new("RGB", (max_size, max_size), "black")
|
341 |
+
new_image.paste(image, ((max_size - width) // 2, (max_size - height) // 2))
|
342 |
+
return new_image
|
343 |
+
else:
|
344 |
+
return image
|
345 |
+
def log_error(e, output_file_path="output_file_path"):
|
346 |
+
# ํ์ฌ ์๊ฐ์ ์ป์ต๋๋ค
|
347 |
+
current_time = datetime.now().strftime("%m/%d %H:%M:%S")
|
348 |
+
|
349 |
+
# ์๋ฌ ๋ก๊ทธ ๋ฉ์์ง
|
350 |
+
error_message = f"#### Error occured at {current_time} ####\nError: {e}\n############################################\n"
|
351 |
+
|
352 |
+
# ์ง์ ๋ ์ถ๋ ฅ ํด๋์ error_log.txt ํ์ผ์ ์ฐ๊ธฐ
|
353 |
+
with open(f"error_log.txt", "a") as file:
|
354 |
+
file.write(error_message)
|
355 |
+
access_token = gen_request["access_token"]
|
356 |
+
mode = gen_request["mode"]
|
357 |
+
defry = gen_request["defry"]
|
358 |
+
prompt = gen_request["prompt"]
|
359 |
+
iw = gen_request["width"]
|
360 |
+
ih = gen_request["height"]
|
361 |
+
image = gen_request["image"]
|
362 |
+
if mode in ["declutter", "lineart", "sketch", "colorize"]: _mode = mode
|
363 |
+
else:
|
364 |
+
_mode = "emotion"
|
365 |
+
mode = mode.lower()
|
366 |
+
data = {
|
367 |
+
"req_type": _mode,
|
368 |
+
"width": iw,
|
369 |
+
"height": ih,
|
370 |
+
"image" : image_to_base64(image)
|
371 |
+
}
|
372 |
+
dfval = {
|
373 |
+
"Normal" : 0,
|
374 |
+
"Slightly Weak" : 1,
|
375 |
+
"Weak" : 2,
|
376 |
+
"Even Weaker" : 3,
|
377 |
+
"Very Weak" : 4,
|
378 |
+
"Weakest" : 5
|
379 |
+
}
|
380 |
+
if _mode == "emotion":
|
381 |
+
try: data["prompt"] = mode+";;"+prompt+";"
|
382 |
+
except: data["prompt"] = mode+";"
|
383 |
+
try: data["defry"] = dfval[defry]
|
384 |
+
except: data["defry"] = 0
|
385 |
+
prompt = data["prompt"]
|
386 |
+
elif _mode == "colorize":
|
387 |
+
data["prompt"] = prompt
|
388 |
+
try: data["defry"] = dfval[defry]
|
389 |
+
except: data["defry"] = 0
|
390 |
+
else:
|
391 |
+
prompt = ""
|
392 |
+
aug_response = requests.post(f"{BASE_URL}/ai/augment-image", json=data, headers={ "Authorization": f"Bearer {access_token}" })
|
393 |
+
save_folder = gen_request["save_folder"]
|
394 |
+
additional_folder = ''
|
395 |
+
if gen_request["png_rule"] == "count":
|
396 |
+
additional_folder = "/" + gen_request["start_time"]
|
397 |
+
if "additional_save_folder" in gen_request:
|
398 |
+
if gen_request["additional_save_folder"]["command1"] != "":
|
399 |
+
additional_folder += "/" + gen_request["additional_save_folder"]["command1"].replace('/',"_")
|
400 |
+
if gen_request["additional_save_folder"]["command2"] != "":
|
401 |
+
additional_folder += "/" + gen_request["additional_save_folder"]["command2"].replace('/',"_")
|
402 |
+
forbidden_chars = r'[\:*?"<>|]'
|
403 |
+
additional_folder = re.sub(forbidden_chars, '_', additional_folder)
|
404 |
+
additional_folder += "/director"
|
405 |
+
d = Path(save_folder + additional_folder)
|
406 |
+
d.mkdir(parents=True, exist_ok=True)
|
407 |
+
try:
|
408 |
+
zipped = zipfile.ZipFile(io.BytesIO(aug_response.content))
|
409 |
+
result_list = []
|
410 |
+
for idx, file_info in enumerate(zipped.infolist()):
|
411 |
+
image_bytes = zipped.read(file_info)
|
412 |
+
if gen_request["png_rule"] == "count":
|
413 |
+
_count = gen_request["count"]
|
414 |
+
if "batch_size" in gen_request: filename = (d / f"{_count:05}_{idx}.png" )
|
415 |
+
else: filename = (d / f"{_count:05}.png" )
|
416 |
+
else:
|
417 |
+
if "batch_size" in gen_request: filename = (d / f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{idx}.png" )
|
418 |
+
else: filename = (d / f"{datetime.now().strftime('%Y%m%d_%H%M%S')}.png" )
|
419 |
+
filename.write_bytes(image_bytes)
|
420 |
+
i = Image.open(io.BytesIO(image_bytes))
|
421 |
+
i = ImageOps.exif_transpose(i).convert("RGB")
|
422 |
+
if "artist" not in gen_request:
|
423 |
+
i_resized = resize_and_fill(i)
|
424 |
+
next = i_resized, prompt, 0, i.info, str(filename)
|
425 |
+
result_list.append(next)
|
426 |
+
return i_resized, prompt, 0, i.info, str(filename) if len(result_list) ==1 else result_list
|
427 |
+
except Exception as e:
|
428 |
+
try:
|
429 |
+
if aug_response.content is None:
|
430 |
+
raise ValueError("Connection broken (Protocol Error)")
|
431 |
+
error_message = aug_response.decode('utf-8')[2:-2]
|
432 |
+
except Exception as inner_exception:
|
433 |
+
error_message = str(inner_exception)
|
434 |
+
log_error(error_message, "path_to_output_folder")
|
435 |
+
return None, error_message, 0, None, None
|
436 |
+
|
437 |
+
def upscale_NAI(access_token, image, parameters):
|
438 |
+
img_str = image_to_base64(image)
|
439 |
+
_width, _height = image.size
|
440 |
+
data = {
|
441 |
+
"image": img_str,
|
442 |
+
"width": _width,
|
443 |
+
"height": _height,
|
444 |
+
"scale": 4,
|
445 |
+
}
|
446 |
+
response = requests.post(f"https://api.novelai.net/ai/upscale", json=data, headers={ "Authorization": f"Bearer {access_token}" })
|
447 |
+
return response.content
|
448 |
+
|
449 |
+
|
450 |
+
def convert_prompt(prompt):
|
451 |
+
keywords = [keyword.strip() for keyword in prompt.split(',')]
|
452 |
+
converted_keywords = []
|
453 |
+
for keyword in keywords:
|
454 |
+
#if 'artist:' in keyword:
|
455 |
+
# keyword = keyword.replace('artist:', '')
|
456 |
+
#if '(' in keyword:
|
457 |
+
# keyword = keyword.replace('(', '\(').replace(')', '\)')
|
458 |
+
if '{' in keyword:
|
459 |
+
keyword = keyword.replace('{', '(').replace('}', ')')
|
460 |
+
converted_keywords.append(keyword)
|
461 |
+
return ', '.join(converted_keywords)
|
462 |
+
|
463 |
+
def convert_prompt_ad(prompt, data):
|
464 |
+
keywords = [keyword.strip() for keyword in prompt.split(',')]
|
465 |
+
converted_keywords = []
|
466 |
+
closed_eyes_check = False if "closed eyes" not in keywords else True
|
467 |
+
for idx, keyword in enumerate(keywords):
|
468 |
+
if idx > 4 and (keyword in data.bag_of_tags):
|
469 |
+
if ("eyes" in keyword or "pupils" in keyword) and closed_eyes_check:
|
470 |
+
continue
|
471 |
+
keywords[idx] = "((" + keyword + "))"
|
472 |
+
if 'artist:' in keyword:
|
473 |
+
keyword = keyword.replace('artist:', '').replace('[','').replace(']','')
|
474 |
+
keyword = keyword
|
475 |
+
if '(' in keyword:
|
476 |
+
keyword = keyword.replace('(', '\(').replace(')', '\)')
|
477 |
+
if '{' in keyword:
|
478 |
+
keyword = keyword.replace('{', '(').replace('}', ')')
|
479 |
+
if '}' in keyword:
|
480 |
+
keyword = keyword.replace('}', ')')
|
481 |
+
converted_keywords.append(keyword)
|
482 |
+
return ', '.join(converted_keywords)
|
483 |
+
|
484 |
+
def interrogate_webui(img, access_token):
|
485 |
+
img_str = image_to_base64(img)
|
486 |
+
|
487 |
+
req_img = {
|
488 |
+
"image": img_str,
|
489 |
+
"threshold": 0.35,
|
490 |
+
"model": "wd-v1-4-moat-tagger.v2",
|
491 |
+
"queue": ""
|
492 |
+
}
|
493 |
+
|
494 |
+
try:
|
495 |
+
res = requests.post(f"{access_token}/tagger/v1/interrogate", json=req_img)
|
496 |
+
if res.status_code != 200:
|
497 |
+
raise Exception(f"Error code: {res.status_code}")
|
498 |
+
|
499 |
+
text = res.json().get('caption', {}).get('tag', {})
|
500 |
+
text_list = [txt.replace("_", " ") for txt in text.keys()]
|
501 |
+
_rating = res.json().get('caption', {}).get('rating', {})
|
502 |
+
rating = max(_rating, key=_rating.get)
|
503 |
+
if rating == "sensitive": rv = "s"
|
504 |
+
elif rating == "questionable": rv = "q"
|
505 |
+
elif rating == "explicit": rv = "e"
|
506 |
+
else: rv = "g"
|
507 |
+
return ", ".join(text_list), rv
|
508 |
+
except Exception as e:
|
509 |
+
return f"{e} : WEBUI์ Extension์ stable-diffusion-webui-wd14-tagger ๊ฐ ์ ์์ ์ผ๋ก ์ค์น๋์ด์๋์ง ํ์ธํ์ธ์. (WD14 ๊ธฐ์ค๋ฒ์ : e72d984b, 2023-11-04)", None
|
510 |
+
|
511 |
+
def generate_image_webui(access_token, prompt, model, action, parameters):
|
512 |
+
samplers = {
|
513 |
+
"k_euler": "Euler",
|
514 |
+
"k_euler_ancestral": "Euler a",
|
515 |
+
"k_dpmpp_2s_ancestral": "DPM++ 2S a",
|
516 |
+
"k_dpmpp_sde": "DPM++ SDE",
|
517 |
+
"k_dpm_3m_sde": "DPM++ 3M SDE"
|
518 |
+
}
|
519 |
+
|
520 |
+
# matching data format
|
521 |
+
data = {
|
522 |
+
"input": prompt,
|
523 |
+
"model": model,
|
524 |
+
"action": action,
|
525 |
+
"parameters": parameters,
|
526 |
+
}
|
527 |
+
|
528 |
+
if data['parameters']['sampler'] in samplers:
|
529 |
+
data['parameters']['sampler'] = samplers[data['parameters']['sampler']]
|
530 |
+
|
531 |
+
params = {
|
532 |
+
"prompt": convert_prompt(data['input']) if "nai_enable_AD" not in parameters else convert_prompt_ad(data['input'], parameters["ad_data"]),
|
533 |
+
"negative_prompt": convert_prompt(data['parameters']['negative_prompt']) if "nai_enable_AD" not in parameters else convert_prompt_ad(data['parameters']['negative_prompt'], parameters["ad_data"]),
|
534 |
+
"steps": math.floor(data['parameters']['cfg_rescale'] / 0.5) * 0.5,
|
535 |
+
"width": data['parameters']['width'],
|
536 |
+
"height": data['parameters']['height'],
|
537 |
+
"cfg_scale": data['parameters']['scale'],
|
538 |
+
"sampler_index": data['parameters']['sampler'],
|
539 |
+
"seed": data['parameters']['seed'],
|
540 |
+
"seed_resize_from_h": -1,
|
541 |
+
"seed_resize_from_w": -1,
|
542 |
+
"denoising_strength": None,
|
543 |
+
"n_iter": "1",
|
544 |
+
"batch_size": data['parameters']['n_samples']
|
545 |
+
}
|
546 |
+
|
547 |
+
if 'scheduler' in parameters:
|
548 |
+
params["scheduler"] = data['parameters']['scheduler']
|
549 |
+
|
550 |
+
if data['parameters']['enable_hr'] == True:
|
551 |
+
params['enable_hr'] = True
|
552 |
+
params["hr_upscaler"] = data['parameters']["hr_upscaler"]
|
553 |
+
params["hr_scale"] = data['parameters']["hr_scale"]
|
554 |
+
params["hr_second_pass_steps"] = data['parameters']["hr_second_pass_steps"]
|
555 |
+
params["denoising_strength"] = data['parameters']["denoising_strength"]
|
556 |
+
|
557 |
+
if data['parameters']['enable_AD'] == True:
|
558 |
+
params["alwayson_scripts"] = {"ADetailer":
|
559 |
+
{
|
560 |
+
"args": [
|
561 |
+
True,
|
562 |
+
False if "nai_enable_AD" not in parameters else True,
|
563 |
+
{
|
564 |
+
"ad_model": "face_yolov8n.pt",
|
565 |
+
"ad_prompt": params["prompt"],
|
566 |
+
"ad_negative_prompt": params["negative_prompt"],
|
567 |
+
"ad_denoising_strength": 0.4 if "nai_enable_AD" not in parameters else parameters["ad_data_str"],
|
568 |
+
"ad_inpaint_only_masked": True,
|
569 |
+
"ad_inpaint_only_masked_padding": 32,
|
570 |
+
"ad_sampler" : data['parameters']['sampler']
|
571 |
+
#"ad_scheduler": "Automatic"
|
572 |
+
}
|
573 |
+
]
|
574 |
+
}
|
575 |
+
}
|
576 |
+
if "nai_enable_AD" in parameters:
|
577 |
+
params["steps"] = 28
|
578 |
+
else:
|
579 |
+
params["alwayson_scripts"] = {}
|
580 |
+
if 'enable_TV' in data['parameters'] and data['parameters']['enable_TV'] == True:
|
581 |
+
params["alwayson_scripts"]["Tiled VAE"] = {}
|
582 |
+
params["alwayson_scripts"]["Tiled VAE"]["args"] = data['parameters']["tiled_vae_args"]
|
583 |
+
|
584 |
+
if 'pag' in parameters:
|
585 |
+
params["alwayson_scripts"]["Incantations"] = {}
|
586 |
+
pag_pre = math.floor(float(data['parameters']['pag']) / 0.5) * 0.5
|
587 |
+
params["alwayson_scripts"]["Incantations"]["args"] = [True,pag_pre, 0, 150, False, "Constant", 0, 100, False,False,2,0.1,0.5,0,"",0,25, 1, False,False,False,"BREAK","-",0.2,10]
|
588 |
+
|
589 |
+
if "image" in data['parameters']:
|
590 |
+
params['init_images'] = [data['parameters']['image']]
|
591 |
+
params['include_init_images'] = True
|
592 |
+
|
593 |
+
if "strength" in data['parameters']:
|
594 |
+
params['denoising_strength'] = data['parameters']['strength']
|
595 |
+
|
596 |
+
if "mask" in data['parameters']:
|
597 |
+
params['mask'] = data['parameters']['mask']
|
598 |
+
params['inpainting_fill'] = 1
|
599 |
+
params['inpaint_full_res'] = data['parameters']['add_original_image']
|
600 |
+
params['inpaint_full_res_padding'] = 32
|
601 |
+
if 'denoising_strength' not in params: params['denoising_strength'] = 0.7
|
602 |
+
|
603 |
+
res = requests.post(f"{access_token}/sdapi/v1/img2img", json=params)
|
604 |
+
else: res = requests.post(f"{access_token}/sdapi/v1/txt2img", json=params)
|
605 |
+
imageb64s = res.json()['images']
|
606 |
+
content = None
|
607 |
+
for b64 in imageb64s:
|
608 |
+
img = b64.encode()
|
609 |
+
content = base64.b64decode(img)
|
610 |
+
|
611 |
+
s = io.BytesIO()
|
612 |
+
zf = zipfile.ZipFile(s, "w")
|
613 |
+
for idx, b64 in enumerate(imageb64s):
|
614 |
+
content = base64.b64decode(b64)
|
615 |
+
zf.writestr(f"generated_{idx}.png", content)
|
616 |
+
zf.close()
|
617 |
+
return s.getvalue()
|
618 |
+
|
619 |
+
def generate_guide_image_webui(access_token, parameters):
|
620 |
+
samplers = {
|
621 |
+
"k_euler": "Euler",
|
622 |
+
"k_euler_ancestral": "Euler a",
|
623 |
+
"k_dpmpp_2s_ancestral": "DPM++ 2S a",
|
624 |
+
"k_dpmpp_sde": "DPM++ SDE",
|
625 |
+
"k_dpm_3m_sde": "DPM++ 3M SDE"
|
626 |
+
}
|
627 |
+
|
628 |
+
params = {
|
629 |
+
"prompt": convert_prompt(parameters['prompt']),
|
630 |
+
"negative_prompt": convert_prompt(parameters['negative prompt']),
|
631 |
+
"steps": parameters['steps'],
|
632 |
+
"width": parameters["width"],
|
633 |
+
"height": parameters["height"],
|
634 |
+
"cfg_scale": parameters['cfg_scale'],
|
635 |
+
"sampler_index": samplers[parameters['sampler']] if parameters['sampler'] in samplers else parameters['sampler'],
|
636 |
+
"seed": parameters['seed'],
|
637 |
+
"seed_resize_from_h": -1,
|
638 |
+
"seed_resize_from_w": -1,
|
639 |
+
"denoising_strength": None,
|
640 |
+
"n_iter": "1",
|
641 |
+
"batch_size": 1
|
642 |
+
}
|
643 |
+
additional_folder = "/guide"
|
644 |
+
if "start_time" in parameters and parameters["png_rule"] == "count":
|
645 |
+
additional_folder = "/" + parameters["start_time"] + "/guide"
|
646 |
+
if "save_folder" in parameters:
|
647 |
+
save_folder = parameters["save_folder"]
|
648 |
+
try:
|
649 |
+
res = requests.post(f"{access_token}/sdapi/v1/txt2img", json=params)
|
650 |
+
imageb64s = res.json()['images']
|
651 |
+
image_data = base64.b64decode(imageb64s[0])
|
652 |
+
image_file = io.BytesIO(image_data)
|
653 |
+
i = Image.open(image_file)
|
654 |
+
if "save_folder" in parameters:
|
655 |
+
s = io.BytesIO()
|
656 |
+
zf = zipfile.ZipFile(s, "w")
|
657 |
+
for idx, b64 in enumerate(imageb64s):
|
658 |
+
img_data = base64.b64decode(b64)
|
659 |
+
img = Image.open(io.BytesIO(img_data))
|
660 |
+
jpeg_buffer = io.BytesIO()
|
661 |
+
try: img.save(jpeg_buffer, format="JPEG", exif=img.info.get('parameters').encode('utf-8'))
|
662 |
+
except: img.save(jpeg_buffer, format="JPEG")
|
663 |
+
jpeg_content = jpeg_buffer.getvalue()
|
664 |
+
zf.writestr(f"generated_{idx}.jpg", jpeg_content)
|
665 |
+
zf.close()
|
666 |
+
d = Path(save_folder + additional_folder)
|
667 |
+
d.mkdir(parents=True, exist_ok=True)
|
668 |
+
zipped = zipfile.ZipFile(io.BytesIO(s.getvalue()))
|
669 |
+
result_list = []
|
670 |
+
for idx, file_info in enumerate(zipped.infolist()):
|
671 |
+
image_bytes = zipped.read(file_info)
|
672 |
+
filename = (d / f"{datetime.now().strftime('%Y%m%d_%H%M%S')}.png" )
|
673 |
+
filename.write_bytes(image_bytes)
|
674 |
+
except:
|
675 |
+
i = Image.new('RGB', (768, 768), 'white')
|
676 |
+
return i
|
677 |
+
|
678 |
+
def generate_typer_image_webui(access_token, parameters):
|
679 |
+
params = {
|
680 |
+
"prompt": convert_prompt(parameters['prompt']),
|
681 |
+
"negative_prompt": convert_prompt(parameters['negative prompt']),
|
682 |
+
"steps": parameters['steps'],
|
683 |
+
"width": parameters["width"],
|
684 |
+
"height": parameters["height"],
|
685 |
+
"cfg_scale": parameters['cfg_scale'],
|
686 |
+
"sampler_index": parameters['sampler'],
|
687 |
+
"seed": parameters['seed'],
|
688 |
+
"seed_resize_from_h": -1,
|
689 |
+
"seed_resize_from_w": -1,
|
690 |
+
"denoising_strength": None,
|
691 |
+
"n_iter": "1",
|
692 |
+
"batch_size": 1
|
693 |
+
}
|
694 |
+
if "scheduler" in parameters:
|
695 |
+
params["scheduler"] = parameters["scheduler"]
|
696 |
+
try:
|
697 |
+
res = requests.post(f"{access_token}/sdapi/v1/txt2img", json=params)
|
698 |
+
imageb64s = res.json()['images']
|
699 |
+
image_data = base64.b64decode(imageb64s[0])
|
700 |
+
image_file = io.BytesIO(image_data)
|
701 |
+
i = Image.open(image_file)
|
702 |
+
except:
|
703 |
+
i = Image.new('RGB', (768, 768), 'white')
|
704 |
+
return i
|
705 |
+
|
706 |
+
def generate_dtg(access_token, parameters):
|
707 |
+
params = {
|
708 |
+
"prompt": parameters['prompt'],
|
709 |
+
"negative_prompt": parameters['negative prompt'],
|
710 |
+
"steps": 1,
|
711 |
+
"width": 64,
|
712 |
+
"height": 64,
|
713 |
+
"cfg_scale": 1,
|
714 |
+
"sampler_index": "Euler a",
|
715 |
+
"seed": 1,
|
716 |
+
"seed_resize_from_h": -1,
|
717 |
+
"seed_resize_from_w": -1,
|
718 |
+
"denoising_strength": None,
|
719 |
+
"n_iter": "1",
|
720 |
+
"batch_size": 1,
|
721 |
+
}
|
722 |
+
if parameters["format"] == "Animagine":
|
723 |
+
format = "<|special|>, <|characters|>, <|copyrights|>, <|artist|>, <|general|>, <|quality|>, <|meta|>, <|rating|>"
|
724 |
+
elif parameters["format"] == "Pony":
|
725 |
+
format = "<|quality|>, <|special|>, <|characters|>, <|copyrights|>, <|artist|>, <|general|>, <|meta|>, <|rating|>"
|
726 |
+
|
727 |
+
params["alwayson_scripts"] = {}
|
728 |
+
params["alwayson_scripts"]["DanTagGen"] = {
|
729 |
+
"args" : [
|
730 |
+
True,
|
731 |
+
"After applying other prompt processings",
|
732 |
+
-1, #random.randint(0, 4294967295 - params["seed"])
|
733 |
+
parameters["length"],
|
734 |
+
params["negative_prompt"],
|
735 |
+
format,
|
736 |
+
parameters["temp"],
|
737 |
+
0.95,
|
738 |
+
100,
|
739 |
+
"KBlueLeaf/DanTagGen-delta-rev2",
|
740 |
+
False,
|
741 |
+
False
|
742 |
+
]
|
743 |
+
}
|
744 |
+
try:
|
745 |
+
res = requests.post(f"{access_token}/sdapi/v1/txt2img", json=params)
|
746 |
+
imageb64s = res.json()['images']
|
747 |
+
image_data = base64.b64decode(imageb64s[0])
|
748 |
+
image_file = io.BytesIO(image_data)
|
749 |
+
i = Image.open(image_file)
|
750 |
+
except:
|
751 |
+
return ""
|
752 |
+
return i.info
|
753 |
+
|
754 |
+
def generate(gen_request, _naia):
|
755 |
+
def parse_and_execute_commands(_prompt, negative, user_input, rating):
|
756 |
+
negative = negative.split(',')
|
757 |
+
negative = [neg.strip() for neg in negative]
|
758 |
+
prompt = _prompt.split(',')
|
759 |
+
prompt = [key.strip() for key in prompt]
|
760 |
+
commands = [cmd.strip() for cmd in user_input.split(',') if not cmd.strip().startswith('#')]
|
761 |
+
for command in commands:
|
762 |
+
condition, cmd = parse_conditional_command(command)
|
763 |
+
if check_condition(prompt, condition, rating):
|
764 |
+
negative = execute_command(negative, cmd)
|
765 |
+
return ', '.join(negative)
|
766 |
+
|
767 |
+
def parse_conditional_command(command):
|
768 |
+
match = re.match(r"\((.*?)\)\:(.*)", command)
|
769 |
+
if match:
|
770 |
+
return match.groups()
|
771 |
+
return '', command
|
772 |
+
|
773 |
+
def check_condition(prompt, condition, rating):
|
774 |
+
if not condition:
|
775 |
+
return True
|
776 |
+
sub_conditions = re.split(r'\)\s*&\s*\(', condition)
|
777 |
+
sub_conditions = [re.sub(r'^\(|\)$', '', cond) for cond in sub_conditions]
|
778 |
+
|
779 |
+
results = []
|
780 |
+
for sub_cond in sub_conditions:
|
781 |
+
if '&' in sub_cond:
|
782 |
+
results.append(all(check_condition(prompt, cond, rating) for cond in sub_cond.split('&')))
|
783 |
+
elif '|' in sub_cond:
|
784 |
+
results.append(any(check_condition(prompt, cond, rating) for cond in sub_cond.split('|')))
|
785 |
+
else:
|
786 |
+
if sub_cond in ['e', 'q', 's', 'g']:
|
787 |
+
results.append(sub_cond == rating)
|
788 |
+
elif sub_cond in ['~e', '~q', '~s', '~g']:
|
789 |
+
results.append(sub_cond != rating)
|
790 |
+
# PM
|
791 |
+
elif sub_cond.startswith('*'):
|
792 |
+
results.append(sub_cond[1:] in prompt)
|
793 |
+
# NOT IN
|
794 |
+
elif sub_cond.startswith('~!'):
|
795 |
+
results.append(sub_cond[2:] not in prompt)
|
796 |
+
elif sub_cond.startswith('~'):
|
797 |
+
results.append(any(sub_cond[1:] not in element for element in prompt))
|
798 |
+
# CONTAIN
|
799 |
+
else:
|
800 |
+
results.append(any(sub_cond in element for element in prompt))
|
801 |
+
return all(results)
|
802 |
+
|
803 |
+
def execute_command(negative, command):
|
804 |
+
if '+=' in command:
|
805 |
+
keyword, addition = command.split('+=', 1)
|
806 |
+
addition = addition.replace('^', ', ')
|
807 |
+
return insert_text_after_keyword(negative, keyword, addition)
|
808 |
+
elif command.startswith('add '):
|
809 |
+
keyword = command[4:]
|
810 |
+
keyword = keyword.replace('^', ', ')
|
811 |
+
keys = keyword.split(',')
|
812 |
+
keys = [key.strip() for key in keys]
|
813 |
+
for key in keys:
|
814 |
+
if key not in negative:
|
815 |
+
negative.append(key)
|
816 |
+
return negative
|
817 |
+
elif command.startswith('rem '):
|
818 |
+
keyword = command[4:]
|
819 |
+
keyword = keyword.replace('^', ', ')
|
820 |
+
keys = keyword.split(',')
|
821 |
+
keys = [key.strip() for key in keys]
|
822 |
+
for key in keys:
|
823 |
+
if key in negative:
|
824 |
+
negative.remove(key)
|
825 |
+
return negative
|
826 |
+
elif '=' in command:
|
827 |
+
keyword, replacement = command.split('=', 1)
|
828 |
+
if keyword in negative:
|
829 |
+
replacement = replacement.replace('^', ', ')
|
830 |
+
index = negative.index(keyword)
|
831 |
+
negative[index] = replacement
|
832 |
+
return negative
|
833 |
+
|
834 |
+
def insert_text_after_keyword(negative, user_keyword, user_additional_keyword):
|
835 |
+
if user_keyword in negative:
|
836 |
+
index = negative.index(user_keyword) + 1
|
837 |
+
negative.insert(index, user_additional_keyword)
|
838 |
+
return negative
|
839 |
+
|
840 |
+
def open_random_png(folderpath):
|
841 |
+
files = os.listdir(folderpath)
|
842 |
+
png_files = [f for f in files if f.endswith('.png')]
|
843 |
+
|
844 |
+
if not png_files:
|
845 |
+
return None
|
846 |
+
|
847 |
+
random_png = random.choice(png_files)
|
848 |
+
img = Image.open(os.path.join(folderpath, random_png))
|
849 |
+
|
850 |
+
return img
|
851 |
+
|
852 |
+
try: seed = int(gen_request["seed"])
|
853 |
+
except: seed = random.randint(0,9999999999)
|
854 |
+
|
855 |
+
try:
|
856 |
+
width = int(gen_request["width"])
|
857 |
+
height = int(gen_request["height"])
|
858 |
+
except:
|
859 |
+
width, height= 1024, 1024
|
860 |
+
|
861 |
+
params = {
|
862 |
+
"legacy": False,
|
863 |
+
"qualityToggle": True if gen_request["quality_toggle"] == 1 else False,
|
864 |
+
"width": width,
|
865 |
+
"height": height,
|
866 |
+
"n_samples": 1 if "batch_size" not in gen_request else gen_request["batch_size"],
|
867 |
+
"seed": seed,
|
868 |
+
"extra_noise_seed": random.randint(0,9999999999),
|
869 |
+
"sampler": gen_request["sampler"],
|
870 |
+
"steps": 28 if "steps" not in gen_request and gen_request["type"]!="upper" else gen_request["steps"],
|
871 |
+
"scale": gen_request["scale"],
|
872 |
+
"negative_prompt": ', '.join([keyword.strip() for keyword in gen_request["negative"].split(',') if not keyword.strip().startswith('#')]),
|
873 |
+
"sm" : True if gen_request["sema"] == 1 else False,
|
874 |
+
"sm_dyn" : True if gen_request["sema_dyn"] == 1 else False,
|
875 |
+
"dynamic_thresholding": True if ("dynamic_thresholding" in gen_request and gen_request["dynamic_thresholding"] == True) else False,
|
876 |
+
"controlnet_strength": 1.0,
|
877 |
+
"add_original_image": False,
|
878 |
+
"cfg_rescale": gen_request["cfg_rescale"],
|
879 |
+
"noise_schedule": gen_request["noise_schedule"],
|
880 |
+
"enable_hr" : gen_request["enable_hr"],
|
881 |
+
"enable_AD" : gen_request["enable_AD"]
|
882 |
+
}
|
883 |
+
|
884 |
+
if "skip_cfg_above_sigma" in gen_request and gen_request["skip_cfg_above_sigma"] == True:
|
885 |
+
if float(gen_request["uncond_scale"]) == 19.19:
|
886 |
+
params["skip_cfg_above_sigma"] = 19.191344202730882
|
887 |
+
else:
|
888 |
+
try:
|
889 |
+
_scale = float(gen_request["uncond_scale"])
|
890 |
+
params["skip_cfg_above_sigma"] = _scale
|
891 |
+
except:
|
892 |
+
params["skip_cfg_above_sigma"] = 19.191344202730882
|
893 |
+
|
894 |
+
if params["enable_hr"] == True:
|
895 |
+
params["hr_upscaler"] = gen_request["hr_upscaler"]
|
896 |
+
params["hr_scale"] = gen_request["hr_scale"]
|
897 |
+
params["hr_second_pass_steps"] = gen_request["hr_second_pass_steps"]
|
898 |
+
params["denoising_strength"] = gen_request["denoising_strength"]
|
899 |
+
if "enable_TV" in gen_request and gen_request["enable_TV"] == True:
|
900 |
+
params["enable_TV"] = True
|
901 |
+
params["tiled_vae_args"] = gen_request["tiled_vae_args"]
|
902 |
+
|
903 |
+
if "reference_image" in gen_request:
|
904 |
+
params["reference_image_multiple"] = gen_request["reference_image"] if isinstance(gen_request["reference_image"], list) else [gen_request["reference_image"]]
|
905 |
+
params["reference_information_extracted_multiple"] = gen_request["reference_information_extracted"] if isinstance(gen_request["reference_information_extracted"], list) else [gen_request["reference_information_extracted"]]
|
906 |
+
params["reference_strength_multiple"] = gen_request["reference_strength"] if isinstance(gen_request["reference_strength"], list) else [gen_request["reference_strength"]]
|
907 |
+
|
908 |
+
if gen_request["type"]=="inpaint":
|
909 |
+
if "mask" in gen_request:
|
910 |
+
params["mask"] = gen_request["mask"]
|
911 |
+
params['add_original_image'] = gen_request['add_original_image']
|
912 |
+
|
913 |
+
positive = gen_request["prompt"]
|
914 |
+
positive = re.sub(r'#.*?(\n|,|$)', '', positive)
|
915 |
+
keywords = [key.strip() for key in positive.split(',')]
|
916 |
+
|
917 |
+
if "cond_negative" in gen_request and gen_request["cond_negative"]:
|
918 |
+
user_input = gen_request["cond_negative"]
|
919 |
+
rating = gen_request["rating"]
|
920 |
+
params["negative_prompt"] = parse_and_execute_commands(positive, params["negative_prompt"], user_input, rating)
|
921 |
+
|
922 |
+
if "repeat" in gen_request:
|
923 |
+
max = gen_request["repeat_max"]
|
924 |
+
|
925 |
+
for i, key in enumerate(keywords):
|
926 |
+
if "->" in key:
|
927 |
+
instant_keyword = [k for k in key.split('->')]
|
928 |
+
if len(instant_keyword) > gen_request["repeat"]:
|
929 |
+
current_key = instant_keyword[gen_request["repeat"]]
|
930 |
+
else:
|
931 |
+
current_key = instant_keyword[gen_request["repeat"] % len(instant_keyword)]
|
932 |
+
keywords[i] = ', '.join(current_key.split('^'))
|
933 |
+
|
934 |
+
filename_rule = gen_request["png_rule"]
|
935 |
+
save_folder = gen_request["save_folder"]
|
936 |
+
|
937 |
+
access_token = gen_request["access_token"]
|
938 |
+
additional_folder = ""
|
939 |
+
|
940 |
+
request_type = "generate"
|
941 |
+
|
942 |
+
if "*i2i:(" in gen_request["prompt"] and "image" not in gen_request:
|
943 |
+
try:
|
944 |
+
start_index = gen_request["prompt"].index('*i2i:(') + len('*i2i:(')
|
945 |
+
end_index = gen_request["prompt"].index(')', start_index)
|
946 |
+
i2i_content = gen_request["prompt"][start_index:end_index]
|
947 |
+
_img, _str = [item.strip() for item in i2i_content.split(':')]
|
948 |
+
_str = float(_str)
|
949 |
+
_img = open_random_png(_img)
|
950 |
+
if _img:
|
951 |
+
gen_request["image"] = image_to_base64(_img)
|
952 |
+
gen_request["strength"] = _str
|
953 |
+
gen_request["noise"] = 0
|
954 |
+
except:
|
955 |
+
pass
|
956 |
+
|
957 |
+
if "nai_enable_AD" in gen_request:
|
958 |
+
params["nai_enable_AD"] = True
|
959 |
+
params["ad_data"] = gen_request["ad_data"]
|
960 |
+
params["ad_data_str"] = gen_request["ad_data_str"]
|
961 |
+
|
962 |
+
if "scheduler" in gen_request:
|
963 |
+
params["scheduler"] = gen_request["scheduler"]
|
964 |
+
if "pag" in gen_request:
|
965 |
+
params["pag"] = gen_request["pag"]
|
966 |
+
|
967 |
+
if "image" in gen_request:
|
968 |
+
params["image"] = gen_request["image"]
|
969 |
+
if "strength" in gen_request:
|
970 |
+
params["strength"] = gen_request["strength"]
|
971 |
+
params["noise"] = gen_request["noise"]
|
972 |
+
params["sm"] = False
|
973 |
+
params["sm_dyn"] = False
|
974 |
+
request_type = "img2img" if "mask" not in gen_request else "infill"
|
975 |
+
|
976 |
+
temp_del = []
|
977 |
+
for key in keywords:
|
978 |
+
if key.startswith('*'):
|
979 |
+
temp_del.append(key)
|
980 |
+
for key in temp_del:
|
981 |
+
if key in keywords:
|
982 |
+
keywords.remove(key)
|
983 |
+
|
984 |
+
positive = ', '.join(keywords)
|
985 |
+
|
986 |
+
def resize_and_fill(image, max_size=None):
|
987 |
+
if max_size is None:
|
988 |
+
max_size = gen_request["user_screen_size"]
|
989 |
+
original_width, original_height = image.size
|
990 |
+
if original_width > max_size or original_height > max_size:
|
991 |
+
# ๋น์จ์ ์ ์งํ๋ฉด์ ํฌ๊ธฐ ์กฐ์
|
992 |
+
image.thumbnail((max_size, max_size))
|
993 |
+
|
994 |
+
# ์ ์ด๋ฏธ์ง ํฌ๊ธฐ ๊ณ์ฐ
|
995 |
+
width, height = image.size
|
996 |
+
new_image = Image.new("RGB", (max_size, max_size), "black")
|
997 |
+
new_image.paste(image, ((max_size - width) // 2, (max_size - height) // 2))
|
998 |
+
return new_image
|
999 |
+
else:
|
1000 |
+
return image
|
1001 |
+
|
1002 |
+
def resize_and_fill_artist(image, max_size=None, _text=""):
|
1003 |
+
if _text != "":
|
1004 |
+
if "_" in _text:
|
1005 |
+
_text0 = _text.split("_")[0]
|
1006 |
+
_text1 = _text.split("_")[1]
|
1007 |
+
_text = _text1+ " : "+_text0
|
1008 |
+
else:
|
1009 |
+
_text1 = _text
|
1010 |
+
if max_size is None:
|
1011 |
+
max_size = gen_request["user_screen_size"]
|
1012 |
+
original_width, original_height = image.size
|
1013 |
+
max_size = (max_size, max_size)
|
1014 |
+
if original_width > max_size[0] or original_height > max_size[1]:
|
1015 |
+
# ๋น์จ์ ์ ์งํ๋ฉด์ ํฌ๊ธฐ ์กฐ์
|
1016 |
+
image.thumbnail(max_size)
|
1017 |
+
|
1018 |
+
# ์ ์ด๋ฏธ์ง ํฌ๊ธฐ ๊ณ์ฐ
|
1019 |
+
width, height = image.size
|
1020 |
+
new_image = Image.new("RGB", (max_size[0], max_size[1]), "black") # ํ์ ์์๋ฅผ ์ํด ๋์ด์ 100 ์ถ๊ฐ
|
1021 |
+
new_image.paste(image, ((max_size[0] - width) // 2, (max_size[1] - height) // 2))
|
1022 |
+
|
1023 |
+
# ํ
์คํธ ์ถ๊ฐ
|
1024 |
+
draw = ImageDraw.Draw(new_image)
|
1025 |
+
font_size = 24
|
1026 |
+
font = ImageFont.truetype("arial.ttf", font_size) # ํฐํธ ํ์ผ ๊ฒฝ๋ก ํ์ธ ํ์
|
1027 |
+
# textbbox ์ฌ์ฉํ์ฌ ํ
์คํธ ๋ฐ์ด๋ฉ ๋ฐ์ค ๊ณ์ฐ
|
1028 |
+
text_x = (max_size[0]) // 2
|
1029 |
+
text_y = (max_size[1] - 40) + (font_size) // 2
|
1030 |
+
bbox = draw.textbbox((text_x, text_y), _text, font=font, anchor="mm")
|
1031 |
+
text_width = bbox[2] - bbox[0]
|
1032 |
+
text_height = bbox[3] - bbox[1]
|
1033 |
+
|
1034 |
+
# ํ
์คํธ ๋ฐฐ๊ฒฝ ๊ทธ๋ฆฌ๊ธฐ (ํ์ ์์)
|
1035 |
+
draw.rectangle([text_x - text_width // 2 - 10, text_y - text_height // 2 - 10, text_x + text_width // 2 + 10, text_y + text_height // 2 + 10], fill="white")
|
1036 |
+
|
1037 |
+
# ํ
์คํธ ๊ทธ๋ฆฌ๊ธฐ
|
1038 |
+
draw.text((text_x, text_y), _text, fill="black", font=font, anchor="mm")
|
1039 |
+
|
1040 |
+
return new_image, _text1
|
1041 |
+
|
1042 |
+
|
1043 |
+
def log_error(e, output_file_path="output_file_path"):
|
1044 |
+
# ํ์ฌ ์๊ฐ์ ์ป์ต๋๋ค
|
1045 |
+
current_time = datetime.now().strftime("%m/%d %H:%M:%S")
|
1046 |
+
|
1047 |
+
# ์๋ฌ ๋ก๊ทธ ๋ฉ์์ง
|
1048 |
+
error_message = f"#### Error occured at {current_time} ####\nError: {e}\n############################################\n"
|
1049 |
+
|
1050 |
+
# ์ง์ ๋ ์ถ๋ ฅ ํด๋์ error_log.txt ํ์ผ์ ์ฐ๊ธฐ
|
1051 |
+
with open(f"error_log.txt", "a") as file:
|
1052 |
+
file.write(error_message)
|
1053 |
+
|
1054 |
+
if _naia == "NAID3-Furry":
|
1055 |
+
_m1 = "nai-diffusion-furry-3"
|
1056 |
+
_m2 = "nai-diffusion-furry-3-inpainting"
|
1057 |
+
else:
|
1058 |
+
_m1 = "nai-diffusion-3"
|
1059 |
+
_m2 = "nai-diffusion-3-inpainting"
|
1060 |
+
|
1061 |
+
try:
|
1062 |
+
zipped_bytes = generate_image(access_token, positive, _m1 if "mask" not in params else _m2, request_type, params)
|
1063 |
+
if gen_request["png_rule"] == "count":
|
1064 |
+
additional_folder = "/" + gen_request["start_time"]
|
1065 |
+
if "additional_save_folder" in gen_request:
|
1066 |
+
if gen_request["additional_save_folder"]["command1"] != "":
|
1067 |
+
additional_folder += "/" + gen_request["additional_save_folder"]["command1"].replace('/',"_")
|
1068 |
+
if gen_request["additional_save_folder"]["command2"] != "":
|
1069 |
+
additional_folder += "/" + gen_request["additional_save_folder"]["command2"].replace('/',"_")
|
1070 |
+
forbidden_chars = r'[\:*?"<>|]'
|
1071 |
+
additional_folder = re.sub(forbidden_chars, '_', additional_folder)
|
1072 |
+
if gen_request["type"] == "turbo":
|
1073 |
+
additional_folder += "/turbo"
|
1074 |
+
if ("hires" in gen_request and gen_request["hires"]) or ("enable_hr" in gen_request and gen_request["enable_hr"] and "http" in access_token) or (int(gen_request["width"]) * int(gen_request["height"]) > 1048576):
|
1075 |
+
additional_folder += "/hires"
|
1076 |
+
if "auto_i2i" in gen_request and gen_request["auto_i2i"]:
|
1077 |
+
if gen_request["png_rule"] == "count": additional_folder = "/" + gen_request["start_time"] + "/autoi2i"
|
1078 |
+
else: additional_folder = "/autoi2i"
|
1079 |
+
if "nai_enable_AD" in gen_request and gen_request["nai_enable_AD"]:
|
1080 |
+
additional_folder += "/Adetailer"
|
1081 |
+
if "webui_nai_i2i" in gen_request and gen_request["webui_nai_i2i"] and "http" not in access_token:
|
1082 |
+
additional_folder += "/webui_nai_i2i"
|
1083 |
+
d = Path(save_folder + additional_folder)
|
1084 |
+
d.mkdir(parents=True, exist_ok=True)
|
1085 |
+
zipped = zipfile.ZipFile(io.BytesIO(zipped_bytes))
|
1086 |
+
result_list = []
|
1087 |
+
for idx, file_info in enumerate(zipped.infolist()):
|
1088 |
+
image_bytes = zipped.read(file_info)
|
1089 |
+
if gen_request["png_rule"] == "count":
|
1090 |
+
_count = gen_request["count"]
|
1091 |
+
if "batch_size" in gen_request: filename = (d / f"{_count:05}_{idx}.png" )
|
1092 |
+
else: filename = (d / f"{_count:05}.png" )
|
1093 |
+
else:
|
1094 |
+
if "batch_size" in gen_request: filename = (d / f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{idx}.png" )
|
1095 |
+
else: filename = (d / f"{datetime.now().strftime('%Y%m%d_%H%M%S')}.png" )
|
1096 |
+
filename.write_bytes(image_bytes)
|
1097 |
+
i = Image.open(io.BytesIO(image_bytes))
|
1098 |
+
i = ImageOps.exif_transpose(i).convert("RGB")
|
1099 |
+
if "artist" not in gen_request:
|
1100 |
+
i_resized = resize_and_fill(i)
|
1101 |
+
else:
|
1102 |
+
i_resized, _i_temp = resize_and_fill_artist(i, None, gen_request["artist"])
|
1103 |
+
if not os.path.exists(f"artist_thumb"):
|
1104 |
+
os.makedirs(f"artist_thumb")
|
1105 |
+
model_name = gen_request["artist_model"]
|
1106 |
+
if not os.path.exists(f"artist_thumb/{model_name}"):
|
1107 |
+
os.makedirs(f"artist_thumb/{model_name}")
|
1108 |
+
_counts = gen_request["artist"].split("_")[0]
|
1109 |
+
i_resized.save(f"artist_thumb/{model_name}/{_i_temp}.jpg")
|
1110 |
+
if not os.path.exists(f"artist_thumb/{model_name}/add_count"):
|
1111 |
+
os.makedirs(f"artist_thumb/{model_name}/add_count")
|
1112 |
+
i_resized.save(f"artist_thumb/{model_name}/add_count/{_counts}_{_i_temp}.jpg")
|
1113 |
+
next = i_resized, positive, params['seed'], i.info, str(filename)
|
1114 |
+
result_list.append(next)
|
1115 |
+
return i_resized, positive, params['seed'], i.info, str(filename) if len(result_list) ==1 else result_list
|
1116 |
+
except Exception as e:
|
1117 |
+
try:
|
1118 |
+
if zipped_bytes is None:
|
1119 |
+
raise ValueError("Connection broken (Protocol Error)")
|
1120 |
+
error_message = zipped_bytes.decode('utf-8')[2:-2]
|
1121 |
+
except Exception as inner_exception:
|
1122 |
+
error_message = str(inner_exception)
|
1123 |
+
log_error(error_message, "path_to_output_folder")
|
1124 |
+
return None, error_message, params['seed'], None, None
|