import requests import random from PIL import Image, ImageOps, ImageDraw, ImageFont from datetime import datetime from pathlib import Path import io, re, base64, os import zipfile import math import time import piexif from imghdr import what BASE_URL="https://image.novelai.net" def extract_prompt_info(info, app_mode=""): import string # EXIF UserComment 디코딩 함수 def decode_user_comment(exif_dict): if 'Exif' in exif_dict and 37510 in exif_dict['Exif']: user_comment = exif_dict['Exif'][37510] if user_comment.startswith(b'UNICODE\x00\x00'): try: # UNICODE 헤더(8바이트) 제거 후 UTF-16으로 디코딩 unicode_data = user_comment[8:] decoded_text = unicode_data.decode('utf-16').replace('\x00', '') decoded_text = decoded_text.replace("UNICODE", "") return re.sub(f"[^{string.printable}]", "", decoded_text) except UnicodeDecodeError: return None return None try: # 1. NAI 모드인 경우 Comment 필드 확인 if app_mode == "NAI": return info.get('Comment', '') # 2. parameters 필드 확인 if 'parameters' in info: return info['parameters'] # 3. EXIF 데이터 확인 if 'exif' in info: exif_dict = piexif.load(info['exif']) # UserComment에서 프롬프트 추출 시도 decoded_comment = decode_user_comment(exif_dict) if decoded_comment: return decoded_comment # 다른 EXIF 필드에서도 프롬프트 찾기 시도 for ifd in exif_dict: if ifd == 'thumbnail': continue if isinstance(exif_dict[ifd], dict): for tag, value in exif_dict[ifd].items(): # 문자열이나 바이트 형태의 데이터가 있는 경우 if isinstance(value, (bytes, str)) and value: try: if isinstance(value, bytes): decoded = value.decode('utf-8', errors='ignore') else: decoded = value if decoded and len(decoded) > 10: # 의미 있는 길이의 텍스트만 고려 # 영어 및 숫자 이외의 문자를 제거하여 필터링 return re.sub(f"[^{string.printable}]", "", decoded) except: continue # 4. 아무것도 찾지 못한 경우 기본값 반환 return "webp or jpeg prompts" except Exception as e: print(f"프롬프트 추출 중 오류 발생: {e}") return "webp or jpeg prompts" def find_max_resolution(width, height, max_pixels=1048576, multiple_of=64): ratio = int(width) / int(height) max_width = int((max_pixels * ratio)**0.5) max_height = int((max_pixels / ratio)**0.5) max_width = (max_width // multiple_of) * multiple_of max_height = (max_height // multiple_of) * multiple_of while max_width * max_height > max_pixels: max_width -= multiple_of max_height = int(max_width / ratio) max_height = (max_height // multiple_of) * multiple_of if max_pixels == 1048576: attempts = 0 max_attempts = 10 while (max_width % multiple_of + max_height % multiple_of) < 32 and attempts < max_attempts: if random.choice([True, False]): if (max_width + multiple_of) * max_height <= max_pixels: max_width += multiple_of else: if max_width * (max_height + multiple_of) <= max_pixels: max_height += multiple_of attempts += 1 return (str(max_width), str(max_height)) def process_xargs(xargs): processed_xargs = [] skip_until = None # 이 변수는 현재 처리 중인 '<...>' 구문의 끝 인덱스를 추적합니다. for i, elem in enumerate(xargs): if skip_until is not None and i <= skip_until: continue # 이미 처리된 범위를 건너뛰기 if "add_negative:<" in elem or "rem_negative:<" in elem: # '<'로 시작하여 '>'로 끝나는 구문 처리 if elem.endswith(">"): processed_xargs.append(elem) else: combined_elem = elem for j in range(i + 1, len(xargs)): combined_elem += ',' + xargs[j] # ','를 포함하여 원소들을 결합 if xargs[j].endswith(">"): skip_until = j # '>'로 끝나는 원소를 찾으면, 처리할 범위의 끝 인덱스를 업데이트 break processed_xargs.append(combined_elem) else: processed_xargs.append(elem) # 일반적인 원소는 그대로 리스트에 추가 return processed_xargs def event_cond_negative(_prompt, negative, user_input, rating): def insert_text_after_keyword(negative, user_keyword, user_additional_keyword): if user_keyword in negative: index = negative.index(user_keyword) + 1 negative.insert(index, user_additional_keyword) return negative def parse_conditional_command(command): match = re.match(r"\((.*?)\)\:(.*)", command) if match: return match.groups() return '', command def check_condition(prompt, condition, rating): if not condition: return True sub_conditions = re.split(r'\)\s*&\s*\(', condition) sub_conditions = [re.sub(r'^\(|\)$', '', cond) for cond in sub_conditions] results = [] for sub_cond in sub_conditions: if '&' in sub_cond: results.append(all(check_condition(prompt, cond, rating) for cond in sub_cond.split('&'))) elif '|' in sub_cond: results.append(any(check_condition(prompt, cond, rating) for cond in sub_cond.split('|'))) else: if sub_cond in ['e', 'q', 's', 'g']: results.append(sub_cond == rating) elif sub_cond in ['~e', '~q', '~s', '~g']: results.append(sub_cond != rating) # PM elif sub_cond.startswith('*'): results.append(sub_cond[1:] in prompt) # NOT IN elif sub_cond.startswith('~!'): results.append(sub_cond[2:] not in prompt) elif sub_cond.startswith('~'): results.append(any(sub_cond[1:] not in element for element in prompt)) # CONTAIN else: results.append(any(sub_cond in element for element in prompt)) return all(results) def execute_command(negative, command): if '+=' in command: keyword, addition = command.split('+=', 1) addition = addition.replace('^', ', ') return insert_text_after_keyword(negative, keyword, addition) elif command.startswith('add '): keyword = command[4:] keyword = keyword.replace('^', ', ') keys = keyword.split(',') keys = [key.strip() for key in keys] for key in keys: if key not in negative: negative.append(key) return negative elif command.startswith('rem '): keyword = command[4:] keyword = keyword.replace('^', ', ') keys = keyword.split(',') keys = [key.strip() for key in keys] for key in keys: if key in negative: negative.remove(key) return negative elif '=' in command: keyword, replacement = command.split('=', 1) if keyword in negative: replacement = replacement.replace('^', ', ') index = negative.index(keyword) negative[index] = replacement return negative negative = negative.split(',') negative = [neg.strip() for neg in negative] prompt = _prompt.split(',') prompt = [key.strip() for key in prompt] commands = [cmd.strip() for cmd in user_input.split(',') if not cmd.strip().startswith('#')] for command in commands: condition, cmd = parse_conditional_command(command) if check_condition(prompt, condition, rating): negative = execute_command(negative, cmd) return ', '.join(negative) def image_to_base64(image): image_bytesIO = io.BytesIO() image.save(image_bytesIO, format="PNG") return base64.b64encode(image_bytesIO.getvalue()).decode() def make_turbo_prompt(gen_request): lines = gen_request['prompt'] result = { "boys": False, "girls": False, "1girl": False, "1boy": False, "1other": False, "others": False } state = { "nude,": False, "pov,": False, "cum,": False, "after ": False, "pussy juice": False, "barefoot": False, "breasts": False, "ejaculation": False, } def insert_spaces(source_list, reference_list): modified_list = source_list.copy() for index, keyword in enumerate(reference_list): if keyword not in source_list: space_count = len(keyword) # 키워드 길이만큼의 공백 문자 modified_list.insert(index, ' ' * space_count) return modified_list keywords = gen_request['prompt'].split(', ') filtered_keywords = [] removed_indices = [] positive0, positive1, positive2, positive3 = gen_request.copy(),gen_request.copy(),gen_request.copy(),gen_request.copy() for word in result.keys(): if word in lines: result[word] = True for word in state.keys(): if word in gen_request['prompt']: state[word] = True key_index = int((len(keywords)/2)-1) if(result["1boy"]) or (result["boys"]): if(result["1girl"]): if(', sex' in gen_request['prompt'] or 'group sex' in gen_request['prompt']): sex_pos_keywords = ['stomach bulge','insertion', 'fucked silly', 'x-ray', 'orgasm', 'cross-section', 'uterus', 'overflow', 'rape', 'vaginal', 'anal'] facial_keywords = ['tongue','ahegao'] temp_sex_pos = [] temp_facial = [] cum_events = [] explicit_check = [] if 'open mouth' in keywords: keywords.remove('open mouth') if 'closed mouth' in keywords: keywords.remove('closed mouth') if 'after rape' in keywords: keywords.remove('after rape') explicit_check.append('after rape') if 'used condom' in keywords: keywords.remove('used condom') explicit_check.append('used condom') for keyword in keywords: 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): filtered_keywords.append(keyword) elif 'sex' in keyword: removed_indices.append(keyword) elif 'penetration' in keyword: removed_indices.append(keyword) elif 'cum' in keyword and keyword != 'cum': cum_events.append(keyword) elif any(sex_pos in keyword for sex_pos in sex_pos_keywords): for sex_pos in sex_pos_keywords: if sex_pos in keyword: temp_sex_pos.append(sex_pos) elif any(facial not in keyword for facial in facial_keywords): for facial in facial_keywords: if facial in keyword: temp_facial.append(facial) filtered_keywords.insert(int((len(filtered_keywords)/2)-1), ' no penetration, imminent penetration') filtered_keywords_positive0 = filtered_keywords.copy() filtered_keywords.remove(' no penetration, imminent penetration') #0 imminent penetration, imminent sex if 'condom' in filtered_keywords and 'condom on penis' not in filtered_keywords: t_index = filtered_keywords.index('condom') rand_num = random.randint(0,2) if rand_num == 1: filtered_keywords.insert(t_index, 'condom on penis') for i, keyword in enumerate(filtered_keywords): if 'pantyhose' in keyword: filtered_keywords[i] = 'torn ' + filtered_keywords[i] #1 default key_index = int((len(filtered_keywords)/2)-1) if 'pussy' in filtered_keywords: key_index = filtered_keywords.index('pussy') if 'penis' in filtered_keywords: key_index = filtered_keywords.index('penis') filtered_keywords[key_index:key_index] = ['motion lines', 'surprised'] for keyword in removed_indices: if 'cum' not in keyword and 'ejaculation' not in keyword: filtered_keywords.insert(key_index,keyword) if(temp_sex_pos): filtered_keywords[key_index:key_index] = temp_sex_pos if 'group sex' in filtered_keywords and 'sex' not in filtered_keywords: t_index = filtered_keywords.index('group sex') filtered_keywords.insert(t_index, 'sex') if('clothed sex' in filtered_keywords and not 'bottomless' in filtered_keywords): filtered_keywords.insert(filtered_keywords.index('clothed sex')+1, 'bottomless') pos1_copied_keywords = filtered_keywords.copy() for i, keyword in enumerate(pos1_copied_keywords): if 'closed eyes' in keyword: rand_num = random.randint(0,2) if(rand_num == 0): pos1_copied_keywords[i] = 'half-' + pos1_copied_keywords[i] elif(rand_num == 1 and 'closed eyes' in pos1_copied_keywords): pos1_copied_keywords.remove('closed eyes') filtered_keywords[i] = 'half-closed eyes' filtered_keywords_positive1 = pos1_copied_keywords.copy() #2 ejaculation,cum in pussy key_index = filtered_keywords.index('surprised') filtered_keywords.remove('surprised') filtered_keywords[key_index:key_index] = ["ejaculation","cum"] if "condom on penis" not in filtered_keywords else ["twitching penis", "[[[[orgasm]]]]"] for keyword in removed_indices: if 'cum' in keyword: filtered_keywords.insert(key_index,keyword) if(temp_facial): filtered_keywords[key_index:key_index] =temp_facial filtered_keywords_positive2 = filtered_keywords.copy() #3 after sex, after ejaculation for i, keyword in enumerate(filtered_keywords): if 'closed eyes' in keyword: rand_num = random.randint(0,2) if(rand_num == 0 and filtered_keywords[i] != 'half-closed eyes'): filtered_keywords[i] = 'half-' + filtered_keywords[i] elif(rand_num == 1): filtered_keywords[i] = 'empty eyes' else: filtered_keywords[i] = 'empty eyes, half-closed eyes' if 'sex' in filtered_keywords: key_index = filtered_keywords.index('sex') elif 'group sex' in filtered_keywords: key_index = filtered_keywords.index('group sex') if "condom on penis" not in filtered_keywords: filtered_keywords.remove('ejaculation') else: filtered_keywords.remove('twitching penis') filtered_keywords.remove('[[[[orgasm]]]]') 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}}"] if(explicit_check): filtered_keywords[key_index:key_index] = explicit_check if 'sex' in filtered_keywords and 'group sex' not in filtered_keywords: if('pussy' in filtered_keywords and not 'anal' in filtered_keywords): if "condom on penis" not in filtered_keywords: filtered_keywords.insert(filtered_keywords.index('sex')+1, 'after vaginal, spread pussy') else: filtered_keywords.insert(filtered_keywords.index('sex')+1, 'after vaginal, spread pussy, pussy juice puddle') elif('anal' in filtered_keywords): if "condom on penis" not in filtered_keywords: filtered_keywords.insert(filtered_keywords.index('sex')+1, 'after anal, cum in ass') else: filtered_keywords.insert(filtered_keywords.index('sex')+1, 'after anal, pussy juice') filtered_keywords.insert(filtered_keywords.index('sex'), 'after sex') filtered_keywords.remove('sex') elif 'group sex' in filtered_keywords: if('vaginal' in filtered_keywords and not 'anal' in filtered_keywords): filtered_keywords.insert(filtered_keywords.index('group sex')+1, 'after vaginal, spread pussy') if 'multiple penises' in filtered_keywords: if "condom on penis" not in filtered_keywords: filtered_keywords.insert(filtered_keywords.index('group sex')+3, 'cum on body, bukkake') else: filtered_keywords.insert(filtered_keywords.index('group sex')+3, 'pussy juice, pussy juice puddle') elif('anal' in filtered_keywords): if "condom on penis" not in filtered_keywords: filtered_keywords.insert(filtered_keywords.index('group sex')+1, 'after anus, cum in ass') else: filtered_keywords.insert(filtered_keywords.index('group sex')+1, 'after anus') if 'multiple penises' in filtered_keywords: if "condom on penis" not in filtered_keywords: filtered_keywords.insert(filtered_keywords.index('group sex')+3, 'cum on body, bukkake') else: filtered_keywords.insert(filtered_keywords.index('group sex')+3, 'pussy juice') else: filtered_keywords.insert(filtered_keywords.index('group sex')+1, 'cum on body, {bukkake}') temp_post_keyword = [] for keyword in sex_pos_keywords: if not (keyword == 'orgasm' or keyword == 'overflow'): if keyword in filtered_keywords: temp_post_keyword.append(keyword) for keyword in temp_post_keyword: filtered_keywords.remove(keyword) positive0['prompt'] = ', '.join(insert_spaces(filtered_keywords_positive0, filtered_keywords)).strip() positive1['prompt'] = ', '.join(insert_spaces(filtered_keywords_positive1, filtered_keywords)).strip() positive2['prompt'] = ', '.join(insert_spaces(filtered_keywords_positive2, filtered_keywords)).strip() positive3['prompt'] = ', '.join(filtered_keywords).strip() positive0["type"] = "turbo" positive1["type"] = "turbo" positive2["type"] = "turbo" positive3["type"] = "turbo" return positive0, positive1, positive2, positive3 def generate_image(access_token, prompt, model, action, parameters): if re.match(r'^http[s]?://', access_token): return generate_image_webui(access_token, prompt, model, action, parameters) return generate_image_NAI(access_token, prompt, model, action, parameters) def generate_image_NAI(access_token, prompt, model, action, parameters): data = { "input": prompt, "model": model, "action": action, "parameters": parameters, } if "fts1" in parameters: data["input"] = parameters["fts1"] + data["input"] if data['parameters']['qualityToggle'] == True: if "nai-diffusion-furry-3" not in data['model']: data['input'] += ', best quality, amazing quality, very aesthetic, absurdres' elif 'nai-diffusion-4' in data["model"]: data['input'] +=", best quality, amazing quality, very aesthetic, absurdres, year 2024" else: data['input'] += ', {best quality}, {amazing quality}' 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"]: data['parameters']['sampler'] = "k_euler_ancestral" if data['parameters']['cfg_rescale'] > 1: data['parameters']['cfg_rescale'] = 0 if 'nai-diffusion-4' in data["model"]: data['parameters']['params_version'] = 3 data['parameters']['add_original_image'] = True data['parameters']['characterPrompts'] = [] data['parameters']['legacy'] = False data['parameters']['legacy_v3_extend'] = False data['parameters']['prefer_brownian'] = True data['parameters']['ucPreset'] = 0 data['parameters']['use_coords'] = False del data['parameters']['sm'] del data['parameters']['sm_dyn'] del data['parameters']['enable_hr'] del data['parameters']['enable_AD'] if "skip_cfg_above_sigma" in data['parameters']: del data['parameters']['skip_cfg_above_sigma'] data['parameters']['v4_negative_prompt'] = { 'caption' : { 'base_caption' : data["parameters"]['negative_prompt'], 'char_captions' : [] }, } data['parameters']['v4_prompt'] = { 'caption' : { 'base_caption' : data['input'], 'char_captions' : [] }, 'use_coords' : False, 'use_order' : True } main_pr = [k.strip() for k in data['input'].split(', ')] sub_pr = [] if 'characters' in data['parameters'] and data['parameters']['characters'][0] != '': for i, k in enumerate(data['parameters']['characters']): _dict = { "center": { "x": 0.5, "y": 0.5 }, "prompt": k, "uc": data['parameters']['characters_uc'][i] } for _k in k.split(','): sub_pr.append(_k.strip()) data['parameters']['characterPrompts'].append(_dict.copy()) del _dict["uc"] _dict["centers"] = [ _dict["center"].copy() ] del _dict["center"] _dict['char_caption'] = _dict['prompt'] del _dict['prompt'] data['parameters']['v4_prompt']['caption']['char_captions'].append(_dict.copy()) _dict['char_caption'] = data['parameters']['characters_uc'][i] data['parameters']['v4_negative_prompt']['caption']['char_captions'].append(_dict.copy()) del data['parameters']['characters'] del data['parameters']['characters_uc'] main_pr = [prompt for prompt in main_pr if prompt not in sub_pr] main_pr = ', '.join(main_pr) data['input'] = main_pr data['parameters']['v4_prompt']['caption']['base_caption'] = main_pr response = requests.post(f"{BASE_URL}/ai/generate-image", json=data, headers={ "Authorization": f"Bearer {access_token}" }, timeout=180) # catch any errors return response.content def augment_image_NAI(gen_request): def resize_and_fill(image, max_size=None): if max_size is None: max_size = gen_request["user_screen_size"] original_width, original_height = image.size if original_width > max_size or original_height > max_size: # 비율을 유지하면서 크기 조정 image.thumbnail((max_size, max_size)) # 새 이미지 크기 계산 width, height = image.size new_image = Image.new("RGB", (max_size, max_size), "black") new_image.paste(image, ((max_size - width) // 2, (max_size - height) // 2)) return new_image else: return image def log_error(e, output_file_path="output_file_path"): # 현재 시간을 얻습니다 current_time = datetime.now().strftime("%m/%d %H:%M:%S") # 에러 로그 메시지 error_message = f"#### Error occured at {current_time} ####\nError: {e}\n############################################\n" # 지정된 출력 폴더의 error_log.txt 파일에 쓰기 with open(f"error_log.txt", "a") as file: file.write(error_message) access_token = gen_request["access_token"] mode = gen_request["mode"] defry = gen_request["defry"] prompt = gen_request["prompt"] iw = gen_request["width"] ih = gen_request["height"] image = gen_request["image"] if mode in ["declutter", "lineart", "sketch", "colorize"]: _mode = mode else: _mode = "emotion" mode = mode.lower() data = { "req_type": _mode, "width": iw, "height": ih, "image" : image_to_base64(image) } dfval = { "Normal" : 0, "Slightly Weak" : 1, "Weak" : 2, "Even Weaker" : 3, "Very Weak" : 4, "Weakest" : 5 } if _mode == "emotion": try: data["prompt"] = mode+";;"+prompt+";" except: data["prompt"] = mode+";" try: data["defry"] = dfval[defry] except: data["defry"] = 0 prompt = data["prompt"] elif _mode == "colorize": data["prompt"] = prompt try: data["defry"] = dfval[defry] except: data["defry"] = 0 else: prompt = "" start_time = time.time() print(access_token[:10]) aug_response = requests.post(f"{BASE_URL}/ai/augment-image", json=data, headers={ "Authorization": f"Bearer {access_token}" }, timeout=180) save_folder = gen_request["save_folder"] additional_folder = '' if gen_request["png_rule"] == "count": additional_folder = "/" + gen_request["start_time"] if "additional_save_folder" in gen_request: if gen_request["additional_save_folder"]["command1"] != "": additional_folder += "/" + gen_request["additional_save_folder"]["command1"].replace('/',"_") if gen_request["additional_save_folder"]["command2"] != "": additional_folder += "/" + gen_request["additional_save_folder"]["command2"].replace('/',"_") forbidden_chars = r'[\\:*?"<>|]' additional_folder = re.sub(forbidden_chars, '_', additional_folder).strip() additional_folder += "/director" if len(additional_folder) > 200: additional_folder = additional_folder[:200] d = Path(save_folder + additional_folder) d.mkdir(parents=True, exist_ok=True) try: zipped = zipfile.ZipFile(io.BytesIO(aug_response.content)) end_time = time.time() response_time = round(end_time - start_time, 2) result_list = [] for idx, file_info in enumerate(zipped.infolist()): image_bytes = zipped.read(file_info) if gen_request["png_rule"] == "count": _count = gen_request["count"] if "batch_size" in gen_request: filename = (d / f"{_count:05}_{idx}.png" ) else: filename = (d / f"{_count:05}.png" ) else: if "batch_size" in gen_request: filename = (d / f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{idx}.png" ) else: filename = (d / f"{datetime.now().strftime('%Y%m%d_%H%M%S')}.png" ) filename.write_bytes(image_bytes) i = Image.open(io.BytesIO(image_bytes)) i = ImageOps.exif_transpose(i).convert("RGB") if "artist" not in gen_request: i_resized = resize_and_fill(i) next = i_resized, prompt, 0, i.info, str(filename) result_list.append(next) return i_resized, prompt, 0, i.info, str(filename) if len(result_list) ==1 else result_list, response_time except Exception as e: end_time = time.time() response_time = round(end_time - start_time, 2) try: if aug_response.content is None: raise ValueError("Connection broken (Protocol Error)") error_message = aug_response.decode('utf-8')[2:-2] except Exception as inner_exception: error_message = str(inner_exception) log_error(error_message, "path_to_output_folder") return None, error_message, 0, None, None, response_time def upscale_NAI(access_token, image, gen_request): def resize_and_fill(image, max_size=None): if max_size is None: max_size = gen_request["user_screen_size"] original_width, original_height = image.size if original_width > max_size or original_height > max_size: # 비율을 유지하면서 크기 조정 image.thumbnail((max_size, max_size)) # 새 이미지 크기 계산 width, height = image.size new_image = Image.new("RGB", (max_size, max_size), "black") new_image.paste(image, ((max_size - width) // 2, (max_size - height) // 2)) return new_image else: return image def log_error(e, output_file_path="output_file_path"): # 현재 시간을 얻습니다 current_time = datetime.now().strftime("%m/%d %H:%M:%S") # 에러 로그 메시지 error_message = f"#### Error occured at {current_time} ####\nError: {e}\n############################################\n" # 지정된 출력 폴더의 error_log.txt 파일에 쓰기 with open(f"error_log.txt", "a") as file: file.write(error_message) img_str = image_to_base64(image) _width, _height = image.size data = { "image": img_str, "width": _width, "height": _height, "scale": 2, } response = requests.post(f"https://api.novelai.net/ai/upscale", json=data, headers={ "Authorization": f"Bearer {access_token}" }) save_folder = gen_request["save_folder"] additional_folder = '' if gen_request["png_rule"] == "count": additional_folder = "/" + gen_request["start_time"] additional_folder += "/upscale" if len(additional_folder) > 200: additional_folder = additional_folder[:200] d = Path(save_folder + additional_folder) forbidden_chars = r'[\\:*?"<>|]' additional_folder = re.sub(forbidden_chars, '_', additional_folder).strip() d.mkdir(parents=True, exist_ok=True) try: zipped = zipfile.ZipFile(io.BytesIO(response.content)) result_list = [] for idx, file_info in enumerate(zipped.infolist()): image_bytes = zipped.read(file_info) if gen_request["png_rule"] == "count": _count = gen_request["count"] filename = (d / f"{_count:05}.png" ) else: filename = (d / f"{datetime.now().strftime('%Y%m%d_%H%M%S')}.png" ) filename.write_bytes(image_bytes) i = Image.open(io.BytesIO(image_bytes)) i = ImageOps.exif_transpose(i).convert("RGB") if "artist" not in gen_request: i_resized = resize_and_fill(i) next = i_resized, gen_request["prompt"], 0, i.info, str(filename) result_list.append(next) return i_resized, gen_request["prompt"], 0, i.info, str(filename) if len(result_list) ==1 else result_list except Exception as e: try: if response.content is None: raise ValueError("Connection broken (Protocol Error)") error_message = response.decode('utf-8')[2:-2] except Exception as inner_exception: error_message = str(inner_exception) log_error(error_message, "path_to_output_folder") return None, error_message, 0, None, None def convert_prompt(prompt): keywords = [keyword.strip() for keyword in prompt.split(',')] converted_keywords = [] for keyword in keywords: #if 'artist:' in keyword: # keyword = keyword.replace('artist:', '') #if '(' in keyword: # keyword = keyword.replace('(', '\(').replace(')', '\)') if '{' in keyword: keyword = keyword.replace('{', '(').replace('}', ')') converted_keywords.append(keyword) return ', '.join(converted_keywords) def convert_prompt_ad(prompt, data): keywords = [keyword.strip() for keyword in prompt.split(',')] converted_keywords = [] closed_eyes_check = False if "closed eyes" not in keywords else True for idx, keyword in enumerate(keywords): if idx > 4 and (keyword in data.bag_of_tags): if ("eyes" in keyword or "pupils" in keyword) and closed_eyes_check: continue keywords[idx] = "((" + keyword + "))" if 'artist:' in keyword: keyword = keyword.replace('artist:', '').replace('[','').replace(']','') keyword = keyword if '(' in keyword: keyword = keyword.replace('(', '\(').replace(')', '\)') if '{' in keyword: keyword = keyword.replace('{', '(').replace('}', ')') if '}' in keyword: keyword = keyword.replace('}', ')') converted_keywords.append(keyword) return ', '.join(converted_keywords) def interrogate_webui(img, access_token): img_str = image_to_base64(img) req_img = { "image": img_str, "threshold": 0.35, "model": "wd-v1-4-moat-tagger.v2", "queue": "" } try: res = requests.post(f"{access_token}/tagger/v1/interrogate", json=req_img) if res.status_code != 200: raise Exception(f"Error code: {res.status_code}") text = res.json().get('caption', {}).get('tag', {}) text_list = [txt.replace("_", " ") for txt in text.keys()] _rating = res.json().get('caption', {}).get('rating', {}) rating = max(_rating, key=_rating.get) if rating == "sensitive": rv = "s" elif rating == "questionable": rv = "q" elif rating == "explicit": rv = "e" else: rv = "g" return ", ".join(text_list), rv except Exception as e: return f"{e} : WEBUI의 Extension에 stable-diffusion-webui-wd14-tagger 가 정상적으로 설치되어있는지 확인하세요. (WD14 기준버전 : e72d984b, 2023-11-04)", None def generate_image_webui(access_token, prompt, model, action, parameters): try: samplers = { "k_euler": "Euler", "k_euler_ancestral": "Euler a", "k_dpmpp_2s_ancestral": "DPM++ 2S a", "k_dpmpp_sde": "DPM++ SDE", "k_dpm_3m_sde": "DPM++ 3M SDE" } # matching data format data = { "input": prompt, "model": model, "action": action, "parameters": parameters, } if data['parameters']['sampler'] in samplers: data['parameters']['sampler'] = samplers[data['parameters']['sampler']] params = { "prompt": convert_prompt(data['input']) if "nai_enable_AD" not in parameters else convert_prompt_ad(data['input'], parameters["ad_data"]), "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"]), "steps": math.floor(data['parameters']['cfg_rescale'] / 0.5) * 0.5 if "nai-diffusion" not in model else data['parameters']['steps'], "width": data['parameters']['width'], "height": data['parameters']['height'], "cfg_scale": data['parameters']['scale'], "sampler_index": data['parameters']['sampler'], "sampler_name": data['parameters']['sampler'], "seed": data['parameters']['seed'], "seed_resize_from_h": -1, "seed_resize_from_w": -1, "denoising_strength": None if "denoising_strength" not in parameters else parameters["denoising_strength"], "n_iter": "1", "batch_size": data['parameters']['n_samples'] } if "sampler_awai" in parameters: data['parameters']['sampler'] = parameters["sampler_awai"] params["sampler_index"] = parameters["sampler_awai"] if 'scheduler' in parameters: params["scheduler"] = data['parameters']['scheduler'] if "scheduler_awai" in parameters: params["scheduler"] = parameters["scheduler_awai"] if data['parameters']['enable_hr'] == True: params['enable_hr'] = True params["hr_upscaler"] = data['parameters']["hr_upscaler"] params["hr_scale"] = data['parameters']["hr_scale"] params["hr_second_pass_steps"] = data['parameters']["hr_second_pass_steps"] params["denoising_strength"] = data['parameters']["denoising_strength"] #params["hr_checkpoint_name"] = "bijutsubu_v10.safetensors" #bijutsubu_v10.safetensors [0a35ea21f1] #params["hr_sampler_name"] = "Use same scheduler" params["hr_scheduler"] = "Automatic" params["hr_prompt"] = params["prompt"] params["hr_negative_prompt"] = params["negative_prompt"] if params["height"] * params["width"] > 1048576: params["width"], params["height"] = find_max_resolution(params["width"], params["height"]) if data['parameters']['enable_AD'] == True: params["alwayson_scripts"] = {"ADetailer": { "args": [ True, False if ("nai_enable_AD" not in parameters or "denoising_strength" in parameters) else True, { "ad_cfg_scale": data['parameters']['scale'], "ad_checkpoint": "Use same checkpoint", "ad_clip_skip": 1, "ad_confidence": 0.3, "ad_controlnet_guidance_end": 1, "ad_controlnet_guidance_start": 0, "ad_controlnet_model": "None", "ad_controlnet_module": "None", "ad_controlnet_weight": 1, "ad_denoising_strength": 0.4 if "nai_enable_AD" not in parameters else parameters["ad_data_str"], "ad_dilate_erode": 4, "ad_inpaint_height": 768, "ad_inpaint_only_masked": True, "ad_inpaint_only_masked_padding": 32, "ad_inpaint_width": 768, "ad_mask_blur": 4, "ad_mask_filter_method": "Area", "ad_mask_k": 0, "ad_mask_max_ratio": 1, "ad_mask_merge_invert": "None", "ad_mask_min_ratio": 0, "ad_model": "face_yolov8n.pt", "ad_model_classes": "", "ad_negative_prompt": params["negative_prompt"], "ad_noise_multiplier": 1, "ad_prompt": params["prompt"], "ad_restore_face": False, "ad_sampler": data['parameters']['sampler'], "ad_scheduler": "Use same scheduler", "ad_steps": params["steps"], "ad_tab_enable": True, "ad_use_cfg_scale": False, "ad_use_checkpoint": False, "ad_use_clip_skip": False, "ad_use_inpaint_width_height": True, "ad_use_noise_multiplier": False, "ad_use_sampler": False, "ad_use_steps": False, "ad_use_vae": False, "ad_vae": "Use same VAE", "ad_x_offset": 0, "ad_y_offset": 0, "is_api": True }, ] } } if "nai_enable_AD" in parameters: params["steps"] = 28 else: params["alwayson_scripts"] = {} if 'enable_TV' in data['parameters'] and data['parameters']['enable_TV'] == True: params["alwayson_scripts"]["Tiled VAE"] = {} params["alwayson_scripts"]["Tiled VAE"]["args"] = data['parameters']["tiled_vae_args"] if 'pag' in parameters: params["alwayson_scripts"]["Incantations"] = {} pag_pre = math.floor(float(data['parameters']['pag']) / 0.5) * 0.5 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] if "image" in data['parameters']: params['init_images'] = [data['parameters']['image']] params['include_init_images'] = True if "strength" in data['parameters']: params['denoising_strength'] = data['parameters']['strength'] if 'denoising_strength' in data['parameters']: params['denoising_strength'] = data['parameters']['denoising_strength'] params['inpainting_fill'] = 1 params["initial_noise_multiplier"] = 1 if "steps" not in params: params["steps"] = 28 if "mask" in data['parameters']: params['mask'] = data['parameters']['mask'] params['inpainting_fill'] = 1 params['inpaint_full_res'] = data['parameters']['add_original_image'] params['inpaint_full_res_padding'] = 32 if 'denoising_strength' not in params: params['denoising_strength'] = 0.7 res = requests.post(f"{access_token}/sdapi/v1/img2img", json=params) else: res = requests.post(f"{access_token}/sdapi/v1/txt2img", json=params) res.raise_for_status() try: response_json = res.json() except ValueError: error_message = f"Invalid JSON response from server. Status code: {res.status_code}, Response: {res.text}" raise RuntimeError(error_message) if 'images' not in response_json: error_message = f"No 'images' in response. Status code: {res.status_code}, Response: {res.text}" raise RuntimeError(error_message) imageb64s = res.json()['images'] if not imageb64s: raise RuntimeError("No images generated") content = None for b64 in imageb64s: img = b64.encode() content = base64.b64decode(img) s = io.BytesIO() zf = zipfile.ZipFile(s, "w") for idx, b64 in enumerate(imageb64s): content = base64.b64decode(b64) zf.writestr(f"generated_{idx}.png", content) zf.close() return s.getvalue() except requests.exceptions.RequestException as e: error_message = f"Network error: {str(e)}" if hasattr(e.response, 'status_code') and hasattr(e.response, 'text'): error_message = f"Server error {e.response.status_code}: {e.response.text}" raise RuntimeError(error_message) except Exception as e: raise RuntimeError(f"Error during image generation: {str(e)}") def generate_guide_image_webui(access_token, parameters): samplers = { "k_euler": "Euler", "k_euler_ancestral": "Euler a", "k_dpmpp_2s_ancestral": "DPM++ 2S a", "k_dpmpp_sde": "DPM++ SDE", "k_dpm_3m_sde": "DPM++ 3M SDE" } params = { "prompt": convert_prompt(parameters['prompt']), "negative_prompt": convert_prompt(parameters['negative prompt']), "steps": parameters['steps'], "width": parameters["width"], "height": parameters["height"], "cfg_scale": parameters['cfg_scale'], "sampler_index": samplers[parameters['sampler']] if parameters['sampler'] in samplers else parameters['sampler'], "seed": parameters['seed'], "seed_resize_from_h": -1, "seed_resize_from_w": -1, "denoising_strength": None, "n_iter": "1", "batch_size": 1 } additional_folder = "/guide" if "start_time" in parameters and parameters["png_rule"] == "count": additional_folder = "/" + parameters["start_time"] + "/guide" if "save_folder" in parameters: save_folder = parameters["save_folder"] try: res = requests.post(f"{access_token}/sdapi/v1/txt2img", json=params) imageb64s = res.json()['images'] image_data = base64.b64decode(imageb64s[0]) image_file = io.BytesIO(image_data) i = Image.open(image_file) if "save_folder" in parameters: s = io.BytesIO() zf = zipfile.ZipFile(s, "w") for idx, b64 in enumerate(imageb64s): img_data = base64.b64decode(b64) img = Image.open(io.BytesIO(img_data)) output = io.BytesIO() if 'parameters' in img.info or 'exif' in img.info: exif_dict = {"Exif": {}} if 'parameters' in img.info: exif_dict["Exif"][piexif.ExifIFD.UserComment] = img.info['parameters'].encode('utf-8') else: temp = extract_prompt_info(img.info) exif_dict["Exif"][piexif.ExifIFD.UserComment] = temp.encode('utf-8') exif_bytes = piexif.dump(exif_dict) img.save(output, format="JPEG", exif=exif_bytes) else: img.save(output, format="JPEG") jpeg_content = output.getvalue() zf.writestr(f"generated_{idx}.jpg", jpeg_content) zf.close() if len(additional_folder) > 200: additional_folder = additional_folder[:200] forbidden_chars = r'[\\:*?"<>|]' additional_folder = re.sub(forbidden_chars, '_', additional_folder).strip() d = Path(save_folder + additional_folder) d.mkdir(parents=True, exist_ok=True) zipped = zipfile.ZipFile(io.BytesIO(s.getvalue())) result_list = [] for idx, file_info in enumerate(zipped.infolist()): image_bytes = zipped.read(file_info) filename = (d / f"{datetime.now().strftime('%Y%m%d_%H%M%S')}.png" ) filename.write_bytes(image_bytes) except: i = Image.new('RGB', (768, 768), 'white') filename = "" return i, params["prompt"], params["seed"], filename def generate_typer_image_webui(access_token, parameters): params = { "prompt": convert_prompt(parameters['prompt']), "negative_prompt": convert_prompt(parameters['negative prompt']), "steps": parameters['steps'], "width": parameters["width"], "height": parameters["height"], "cfg_scale": parameters['cfg_scale'], "sampler_index": parameters['sampler'], "seed": parameters['seed'], "seed_resize_from_h": -1, "seed_resize_from_w": -1, "denoising_strength": None, "n_iter": "1", "batch_size": 1 } if "scheduler" in parameters: params["scheduler"] = parameters["scheduler"] try: res = requests.post(f"{access_token}/sdapi/v1/txt2img", json=params) imageb64s = res.json()['images'] image_data = base64.b64decode(imageb64s[0]) image_file = io.BytesIO(image_data) i = Image.open(image_file) except: i = Image.new('RGB', (768, 768), 'white') return i def generate_dtg(access_token, parameters): params = { "prompt": parameters['prompt'], "negative_prompt": parameters['negative prompt'], "steps": 1, "width": 64, "height": 64, "cfg_scale": 1, "sampler_index": "Euler a", "seed": 1, "seed_resize_from_h": -1, "seed_resize_from_w": -1, "denoising_strength": None, "n_iter": "1", "batch_size": 1, } if parameters["format"] == "Animagine": format = "<|special|>, <|characters|>, <|copyrights|>, <|artist|>, <|general|>, <|quality|>, <|meta|>, <|rating|>" elif parameters["format"] == "Pony": format = "<|quality|>, <|special|>, <|characters|>, <|copyrights|>, <|artist|>, <|general|>, <|meta|>, <|rating|>" params["alwayson_scripts"] = {} params["alwayson_scripts"]["DanTagGen"] = { "args" : [ True, "After applying other prompt processings", -1, #random.randint(0, 4294967295 - params["seed"]) parameters["length"], params["negative_prompt"], format, parameters["temp"], 0.95, 100, "KBlueLeaf/DanTagGen-delta-rev2", False, False ] } try: res = requests.post(f"{access_token}/sdapi/v1/txt2img", json=params) imageb64s = res.json()['images'] image_data = base64.b64decode(imageb64s[0]) image_file = io.BytesIO(image_data) i = Image.open(image_file) except: return "" return i.info def generate(gen_request, _naia): def parse_and_execute_commands(_prompt, negative, user_input, rating): negative = negative.split(',') negative = [neg.strip() for neg in negative] prompt = _prompt.split(',') prompt = [key.strip() for key in prompt] commands = [cmd.strip() for cmd in user_input.split(',') if not cmd.strip().startswith('#')] for command in commands: condition, cmd = parse_conditional_command(command) if check_condition(prompt, condition, rating): negative = execute_command(negative, cmd) return ', '.join(negative) def parse_conditional_command(command): match = re.match(r"\((.*?)\)\:(.*)", command) if match: return match.groups() return '', command def check_condition(prompt, condition, rating): if not condition: return True sub_conditions = re.split(r'\)\s*&\s*\(', condition) sub_conditions = [re.sub(r'^\(|\)$', '', cond) for cond in sub_conditions] results = [] for sub_cond in sub_conditions: if '&' in sub_cond: results.append(all(check_condition(prompt, cond, rating) for cond in sub_cond.split('&'))) elif '|' in sub_cond: results.append(any(check_condition(prompt, cond, rating) for cond in sub_cond.split('|'))) else: if sub_cond in ['e', 'q', 's', 'g']: results.append(sub_cond == rating) elif sub_cond in ['~e', '~q', '~s', '~g']: results.append(sub_cond != rating) # PM elif sub_cond.startswith('*'): results.append(sub_cond[1:] in prompt) # NOT IN elif sub_cond.startswith('~!'): results.append(sub_cond[2:] not in prompt) elif sub_cond.startswith('~'): results.append(all(sub_cond[1:] not in element for element in prompt)) # CONTAIN else: results.append(any(sub_cond in element for element in prompt)) return all(results) def execute_command(negative, command): if '+=' in command: keyword, addition = command.split('+=', 1) addition = addition.replace('^', ', ') return insert_text_after_keyword(negative, keyword, addition) elif command.startswith('add '): keyword = command[4:] keyword = keyword.replace('^', ', ') keys = keyword.split(',') keys = [key.strip() for key in keys] for key in keys: if key not in negative: negative.append(key) return negative elif command.startswith('rem '): keyword = command[4:] keyword = keyword.replace('^', ', ') keys = keyword.split(',') keys = [key.strip() for key in keys] for key in keys: if key in negative: negative.remove(key) return negative elif '=' in command: keyword, replacement = command.split('=', 1) if keyword in negative: replacement = replacement.replace('^', ', ') index = negative.index(keyword) negative[index] = replacement return negative def insert_text_after_keyword(negative, user_keyword, user_additional_keyword): if user_keyword in negative: index = negative.index(user_keyword) + 1 negative.insert(index, user_additional_keyword) return negative def open_random_png(folderpath): files = os.listdir(folderpath) png_files = [f for f in files if f.endswith('.png')] if not png_files: return None random_png = random.choice(png_files) img = Image.open(os.path.join(folderpath, random_png)) return img try: seed = int(gen_request["seed"]) except: seed = random.randint(0,9999999999) try: width = int(gen_request["width"]) height = int(gen_request["height"]) except: width, height= 1024, 1024 params = { "legacy": False, "qualityToggle": True if gen_request["quality_toggle"] == 1 else False, "width": width, "height": height, "n_samples": 1 if "batch_size" not in gen_request else gen_request["batch_size"], "seed": seed, "extra_noise_seed": random.randint(0,9999999999), "sampler": gen_request["sampler"], "steps": 28 if "steps" not in gen_request and gen_request["type"]!="upper" else gen_request["steps"], "scale": gen_request["scale"], "negative_prompt": ', '.join([keyword.strip() for keyword in gen_request["negative"].split(',') if not keyword.strip().startswith('#')]), "sm" : True if gen_request["sema"] == 1 else False, "sm_dyn" : True if gen_request["sema_dyn"] == 1 else False, "dynamic_thresholding": True if ("dynamic_thresholding" in gen_request and gen_request["dynamic_thresholding"] == True) else False, "controlnet_strength": 1.0, "add_original_image": False, "cfg_rescale": gen_request["cfg_rescale"], "noise_schedule": gen_request["noise_schedule"], "enable_hr" : gen_request["enable_hr"], "enable_AD" : gen_request["enable_AD"] } if "fts1" in gen_request: params["fts1"] = gen_request["fts1"] + " :" + gen_request["fts1_val"].strip() +" | " if "skip_cfg_above_sigma" in gen_request and gen_request["skip_cfg_above_sigma"] == True: if float(gen_request["uncond_scale"]) == 19.19: params["skip_cfg_above_sigma"] = 19.191344202730882 else: try: _scale = float(gen_request["uncond_scale"]) params["skip_cfg_above_sigma"] = _scale except: params["skip_cfg_above_sigma"] = 19.191344202730882 if params["enable_hr"] == True: params["hr_upscaler"] = gen_request["hr_upscaler"] params["hr_scale"] = gen_request["hr_scale"] params["hr_second_pass_steps"] = gen_request["hr_second_pass_steps"] params["denoising_strength"] = gen_request["denoising_strength"] if "enable_TV" in gen_request and gen_request["enable_TV"] == True: params["enable_TV"] = True params["tiled_vae_args"] = gen_request["tiled_vae_args"] if "reference_image" in gen_request: params["reference_image_multiple"] = gen_request["reference_image"] if isinstance(gen_request["reference_image"], list) else [gen_request["reference_image"]] params["reference_information_extracted_multiple"] = gen_request["reference_information_extracted"] if isinstance(gen_request["reference_information_extracted"], list) else [gen_request["reference_information_extracted"]] params["reference_strength_multiple"] = gen_request["reference_strength"] if isinstance(gen_request["reference_strength"], list) else [gen_request["reference_strength"]] if gen_request["type"]=="inpaint": if "mask" in gen_request: params["mask"] = gen_request["mask"] params['add_original_image'] = gen_request['add_original_image'] positive = gen_request["prompt"] positive = re.sub(r'#.*?(\n|,|$)', '', positive) keywords = [key.strip() for key in positive.split(',')] if "cond_negative" in gen_request and gen_request["cond_negative"]: user_input = gen_request["cond_negative"] rating = gen_request["rating"] params["negative_prompt"] = parse_and_execute_commands(positive, params["negative_prompt"], user_input, rating) if "repeat" in gen_request: max = gen_request["repeat_max"] for i, key in enumerate(keywords): if "->" in key: instant_keyword = [k for k in key.split('->')] if len(instant_keyword) > gen_request["repeat"]: current_key = instant_keyword[gen_request["repeat"]] else: current_key = instant_keyword[gen_request["repeat"] % len(instant_keyword)] keywords[i] = ', '.join(current_key.split('^')) filename_rule = gen_request["png_rule"] save_folder = gen_request["save_folder"] access_token = gen_request["access_token"] additional_folder = "" request_type = "generate" if "*i2i:(" in gen_request["prompt"] and "image" not in gen_request: try: start_index = gen_request["prompt"].index('*i2i:(') + len('*i2i:(') end_index = gen_request["prompt"].index(')', start_index) i2i_content = gen_request["prompt"][start_index:end_index] _img, _str = [item.strip() for item in i2i_content.split(':')] _str = float(_str) _img = open_random_png(_img) if _img: gen_request["image"] = image_to_base64(_img) gen_request["strength"] = _str gen_request["noise"] = 0 except: pass if "nai_enable_AD" in gen_request: params["nai_enable_AD"] = True params["ad_data"] = gen_request["ad_data"] params["ad_data_str"] = gen_request["ad_data_str"] if "denoising_strength" in gen_request: params["denoising_strength"] = gen_request["denoising_strength"] if "scheduler" in gen_request: params["scheduler"] = gen_request["scheduler"] if "pag" in gen_request: params["pag"] = gen_request["pag"] if "image" in gen_request: params["image"] = gen_request["image"] if "strength" in gen_request: params["strength"] = gen_request["strength"] params["noise"] = gen_request["noise"] params["sm"] = False params["sm_dyn"] = False request_type = "img2img" if "mask" not in gen_request else "infill" temp_del = [] for key in keywords: if key.startswith('*'): temp_del.append(key) for key in temp_del: if key in keywords: keywords.remove(key) positive = ', '.join(k for k in keywords if k.strip()) def resize_and_fill(image, max_size=None): if max_size is None: max_size = gen_request["user_screen_size"] original_width, original_height = image.size if original_width > max_size or original_height > max_size: # 비율을 유지하면서 크기 조정 image.thumbnail((max_size, max_size)) # 새 이미지 크기 계산 width, height = image.size new_image = Image.new("RGB", (max_size, max_size), "black") new_image.paste(image, ((max_size - width) // 2, (max_size - height) // 2)) return new_image else: return image def resize_and_fill_artist(image, max_size=None, _text=""): if _text != "": if "_" in _text: _text0 = _text.split("_")[0] _text1 = _text.split("_")[1] _text = _text1+ " : "+_text0 else: _text1 = _text if max_size is None: max_size = gen_request["user_screen_size"] original_width, original_height = image.size max_size = (max_size, max_size) if original_width > max_size[0] or original_height > max_size[1]: # 비율을 유지하면서 크기 조정 image.thumbnail(max_size) # 새 이미지 크기 계산 width, height = image.size new_image = Image.new("RGB", (max_size[0], max_size[1]), "black") # 하얀 상자를 위해 높이에 100 추가 new_image.paste(image, ((max_size[0] - width) // 2, (max_size[1] - height) // 2)) # 텍스트 추가 draw = ImageDraw.Draw(new_image) font_size = 24 font = ImageFont.truetype("arial.ttf", font_size) # 폰트 파일 경로 확인 필요 # textbbox 사용하여 텍스트 바운딩 박스 계산 text_x = (max_size[0]) // 2 text_y = (max_size[1] - 40) + (font_size) // 2 bbox = draw.textbbox((text_x, text_y), _text, font=font, anchor="mm") text_width = bbox[2] - bbox[0] text_height = bbox[3] - bbox[1] # 텍스트 배경 그리기 (하얀 상자) 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") # 텍스트 그리기 draw.text((text_x, text_y), _text, fill="black", font=font, anchor="mm") return new_image, _text1 def log_error(e, output_file_path="output_file_path"): # 현재 시간을 얻습니다 current_time = datetime.now().strftime("%m/%d %H:%M:%S") # 에러 로그 메시지 error_message = f"#### Error occured at {current_time} ####\nError: {e}\n############################################\n" # 지정된 출력 폴더의 error_log.txt 파일에 쓰기 with open(f"error_log.txt", "a") as file: file.write(error_message) if _naia == "NAID3-Furry": _m1 = "nai-diffusion-furry-3" _m2 = "nai-diffusion-furry-3-inpainting" elif _naia == "NAID4": _m1 = "nai-diffusion-4-curated-preview" _m2 = "nai-diffusion-3-inpainting" if "characters" in gen_request: params['characters'] = gen_request['characters'] params['characters_uc'] = gen_request['characters_uc'] else: _m1 = "nai-diffusion-3" _m2 = "nai-diffusion-3-inpainting" if "sampler_awai" in gen_request: params["sampler_awai"] = gen_request["sampler_awai"] if "scheduler_awai" in gen_request: params["scheduler_awai"] = gen_request["scheduler_awai"] start_time = time.time() try: zipped_bytes = None error_message = None try: zipped_bytes = generate_image(access_token, positive, _m1 if "mask" not in params else _m2, request_type, params) end_time = time.time() response_time = round(end_time - start_time, 2) if gen_request["png_rule"] == "count": additional_folder = "/" + gen_request["start_time"] if "additional_save_folder" in gen_request: if gen_request["additional_save_folder"]["command1"] != "": additional_folder += "/" + gen_request["additional_save_folder"]["command1"].replace('/',"_") if gen_request["additional_save_folder"]["command2"] != "": additional_folder += "/" + gen_request["additional_save_folder"]["command2"].replace('/',"_") if gen_request["type"] == "turbo": additional_folder += "/turbo" 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): additional_folder += "/hires" if "auto_i2i" in gen_request and gen_request["auto_i2i"]: if gen_request["png_rule"] == "count": additional_folder = "/" + gen_request["start_time"] + "/autoi2i" else: additional_folder = "/autoi2i" if "nai_enable_AD" in gen_request and gen_request["nai_enable_AD"]: additional_folder += "/Adetailer" if "webui_nai_i2i" in gen_request and gen_request["webui_nai_i2i"] and "http" not in access_token: additional_folder += "/webui_nai_i2i" if len(additional_folder) > 200: additional_folder = additional_folder[:200] forbidden_chars = r'[\\:*?"<>|]' additional_folder = re.sub(forbidden_chars, '_', additional_folder).strip() d = Path(save_folder + additional_folder) d.mkdir(parents=True, exist_ok=True) zipped = zipfile.ZipFile(io.BytesIO(zipped_bytes)) result_list = [] for idx, file_info in enumerate(zipped.infolist()): image_bytes = zipped.read(file_info) image_format = what(None, h=image_bytes) if image_format is None: if image_bytes.startswith(b'RIFF') and b'WEBP' in image_bytes[:12]: image_format = 'webp' else: image_format = 'png' if gen_request["png_rule"] == "count": _count = gen_request["count"] if "batch_size" in gen_request: filename = (d / f"{_count:05}_{idx}.{image_format}" ) else: filename = (d / f"{_count:05}.{image_format}" ) else: if "batch_size" in gen_request: filename = (d / f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{idx}.{image_format}" ) else: filename = (d / f"{datetime.now().strftime('%Y%m%d_%H%M%S')}.{image_format}" ) filename.write_bytes(image_bytes) i = Image.open(io.BytesIO(image_bytes)) i = ImageOps.exif_transpose(i).convert("RGB") if "artist" not in gen_request: i_resized = resize_and_fill(i) else: i_resized, _i_temp = resize_and_fill_artist(i, None, gen_request["artist"]) if not os.path.exists(f"artist_thumb"): os.makedirs(f"artist_thumb") model_name = gen_request["artist_model"] if not os.path.exists(f"artist_thumb/{model_name}"): os.makedirs(f"artist_thumb/{model_name}") #_counts = gen_request["artist"].split("_")[0] forbidden_chars = r'[\\/:*?"<>|]' _i_temp = re.sub(forbidden_chars, '_', _i_temp) i_resized.save(f"artist_thumb/{model_name}/{_i_temp}.jpg") next = i_resized, positive, params['seed'], i.info, str(filename) result_list.append(next) return i_resized, positive, params['seed'], i.info, str(filename) if len(result_list) ==1 else result_list, response_time except Exception as e: error_message = str(e) raise except Exception as e: end_time = time.time() response_time = round(end_time - start_time, 2) if not error_message: if zipped_bytes is not None: try: error_message = zipped_bytes.decode('utf-8')[2:-2] except: error_message = str(e) else: error_message = str(e) log_error(error_message, "path_to_output_folder") return None, error_message, params['seed'], None, None, response_time