baqu2213 commited on
Commit
7b89563
โ€ข
1 Parent(s): afcfbf2

Upload NAIA_generation.py

Browse files
Files changed (1) hide show
  1. 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