├── app ├── config │ ├── __init__.py │ ├── strings.py │ ├── paths.py │ ├── consts.py │ └── themes.py ├── core │ ├── __init__.py │ ├── thread │ │ ├── __init__.py │ │ ├── anlas_thread.py │ │ ├── token_thread.py │ │ ├── loading_thread.py │ │ ├── login_thread.py │ │ ├── completiontagload_thread.py │ │ └── generate_thread.py │ └── worker │ │ ├── __init__.py │ │ ├── wildcard_applier.py │ │ ├── naiinfo_getter.py │ │ ├── stealth_pnginfo.py │ │ ├── danbooru_tagger.py │ │ ├── completer.py │ │ └── nai_generator.py ├── gui │ ├── __init__.py │ ├── dialog │ │ ├── __init__.py │ │ ├── fileio_dialog.py │ │ ├── generate_dialog.py │ │ ├── etc_dialog.py │ │ ├── login_dialog.py │ │ ├── option_dialog.py │ │ ├── miniutil_dialog.py │ │ └── inpaint_dialog.py │ ├── layout │ │ ├── __init__.py │ │ ├── model_options_layout.py │ │ ├── expand_layout.py │ │ ├── prompt_layout.py │ │ ├── generate_buttons_layout.py │ │ ├── main_layout.py │ │ ├── parameter_options_layout.py │ │ ├── resolution_options_layout.py │ │ ├── image_options_layout.py │ │ └── character_prompt_layout.py │ └── widget │ │ ├── __init__.py │ │ ├── status_bar.py │ │ ├── result_image_view.py │ │ ├── menu_bar.py │ │ └── custom_slider_widget.py ├── util │ ├── __init__.py │ ├── common_util.py │ ├── ui_util.py │ ├── file_util.py │ ├── tagger_util.py │ ├── image_util.py │ └── string_util.py ├── getter_window.py ├── tagger_window.py └── main_window.py ├── .gitattributes ├── assets ├── icon.ico ├── icon.png ├── getter.png ├── tagger.png ├── no_image.png ├── icon_getter.ico ├── icon_getter.png ├── icon_tagger.ico ├── icon_tagger.png ├── image_clear.png └── open_image.png ├── credit.txt ├── .gitignore └── README.md /app/config/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /app/core/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /app/gui/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /app/util/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /app/core/thread/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /app/core/worker/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /app/gui/dialog/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /app/gui/layout/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /app/gui/widget/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /.gitattributes: -------------------------------------------------------------------------------- 1 | # Auto detect text files and perform LF normalization 2 | * text=auto 3 | -------------------------------------------------------------------------------- /assets/icon.ico: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCP-arca/NAI-Auto-Generator/HEAD/assets/icon.ico -------------------------------------------------------------------------------- /assets/icon.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCP-arca/NAI-Auto-Generator/HEAD/assets/icon.png -------------------------------------------------------------------------------- /assets/getter.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCP-arca/NAI-Auto-Generator/HEAD/assets/getter.png -------------------------------------------------------------------------------- /assets/tagger.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCP-arca/NAI-Auto-Generator/HEAD/assets/tagger.png -------------------------------------------------------------------------------- /assets/no_image.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCP-arca/NAI-Auto-Generator/HEAD/assets/no_image.png -------------------------------------------------------------------------------- /assets/icon_getter.ico: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCP-arca/NAI-Auto-Generator/HEAD/assets/icon_getter.ico -------------------------------------------------------------------------------- /assets/icon_getter.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCP-arca/NAI-Auto-Generator/HEAD/assets/icon_getter.png -------------------------------------------------------------------------------- /assets/icon_tagger.ico: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCP-arca/NAI-Auto-Generator/HEAD/assets/icon_tagger.ico -------------------------------------------------------------------------------- /assets/icon_tagger.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCP-arca/NAI-Auto-Generator/HEAD/assets/icon_tagger.png -------------------------------------------------------------------------------- /assets/image_clear.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCP-arca/NAI-Auto-Generator/HEAD/assets/image_clear.png -------------------------------------------------------------------------------- /assets/open_image.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/DCP-arca/NAI-Auto-Generator/HEAD/assets/open_image.png -------------------------------------------------------------------------------- /credit.txt: -------------------------------------------------------------------------------- 1 | tags.json from 2 | https://gist.githubusercontent.com/bem13/0bc5091819f0594c53f0d96972c8b6ff/raw/b0aacd5ea4634ed4a9f320d344cc1fe81a60db5a/danbooru_tags_post_count.csv -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | build/dist/ 2 | build/dist_onefile/ 3 | build/dist_release_zip/ 4 | build/_pyinstaller_build/ 5 | mine/ 6 | __pycache__ 7 | *.spec 8 | assets/ART_WORK/ 9 | .venv/ 10 | .venv_test/ 11 | *.ini 12 | .vscode -------------------------------------------------------------------------------- /app/getter_window.py: -------------------------------------------------------------------------------- 1 | import sys 2 | 3 | from PyQt5.QtWidgets import QApplication, QMainWindow 4 | 5 | from gui.dialog.miniutil_dialog import MiniUtilDialog 6 | 7 | if __name__ == "__main__": 8 | app = QApplication(sys.argv) 9 | qw = QMainWindow() 10 | qw.move(200, 200) 11 | MiniUtilDialog(qw, "getter").exec_() 12 | -------------------------------------------------------------------------------- /app/core/thread/anlas_thread.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtCore import QThread, pyqtSignal 2 | 3 | class AnlasThread(QThread): 4 | anlas_result = pyqtSignal(int) 5 | 6 | def __init__(self, parent): 7 | super(AnlasThread, self).__init__(parent) 8 | 9 | def run(self): 10 | anlas = self.parent().nai.get_anlas() or -1 11 | 12 | self.anlas_result.emit(anlas) -------------------------------------------------------------------------------- /app/core/thread/token_thread.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtCore import QThread, pyqtSignal 2 | 3 | class TokenValidateThread(QThread): 4 | validation_result = pyqtSignal(int) 5 | 6 | def __init__(self, parent): 7 | super(TokenValidateThread, self).__init__(parent) 8 | 9 | def run(self): 10 | is_login_success = self.parent().nai.check_logged_in() 11 | 12 | self.validation_result.emit(0 if is_login_success else 1) -------------------------------------------------------------------------------- /app/core/thread/loading_thread.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtCore import QThread, pyqtSignal 2 | 3 | class LoadingThread(QThread): 4 | finished = pyqtSignal(str) 5 | 6 | def __init__(self, func): 7 | super().__init__() 8 | self.func = func 9 | 10 | def run(self): 11 | try: 12 | self.finished.emit(self.func()) 13 | except Exception as e: 14 | print(e) 15 | self.finished.emit("") -------------------------------------------------------------------------------- /app/core/thread/login_thread.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtCore import QThread, pyqtSignal 2 | 3 | class LoginThread(QThread): 4 | login_result = pyqtSignal(int) 5 | 6 | def __init__(self, parent, nai, username, password): 7 | super(LoginThread, self).__init__(parent) 8 | self.nai = nai 9 | self.username = username 10 | self.password = password 11 | 12 | def run(self): 13 | if not self.username or not self.password: 14 | self.login_result.emit(1) 15 | return 16 | 17 | is_login_success = self.nai.try_login( 18 | self.username, self.password) 19 | 20 | self.login_result.emit(0 if is_login_success else 2) 21 | -------------------------------------------------------------------------------- /app/tagger_window.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import os 3 | 4 | from PIL import Image 5 | 6 | from PyQt5.QtWidgets import QMessageBox, QApplication, QMainWindow 7 | from PyQt5.QtCore import QSettings 8 | 9 | from gui.dialog.miniutil_dialog import MiniUtilDialog 10 | from core.worker.danbooru_tagger import DanbooruTagger 11 | from config.paths import DEFAULT_PATH 12 | 13 | if __name__ == "__main__": 14 | app = QApplication(sys.argv) 15 | qw = QMainWindow() 16 | qw.move(200, 200) 17 | TOP_NAME = "dcp_arca" 18 | APP_NAME = "nag_gui" 19 | qw.settings = QSettings(TOP_NAME, APP_NAME) 20 | qw.dtagger = DanbooruTagger(qw.settings.value( 21 | "path_models", os.path.abspath(DEFAULT_PATH["path_models"]))) 22 | MiniUtilDialog(qw, "tagger").exec_() 23 | -------------------------------------------------------------------------------- /app/gui/layout/model_options_layout.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtWidgets import QGroupBox, QLabel, QVBoxLayout, QHBoxLayout, QComboBox 2 | 3 | from core.worker.nai_generator import MODEL_INFO_DICT 4 | 5 | 6 | def init_model_options_layout(self): 7 | layout = QVBoxLayout() 8 | 9 | model_options_group = QGroupBox("Image Settings") 10 | layout.addWidget(model_options_group) 11 | 12 | model_options_layout = QHBoxLayout() 13 | model_options_group.setLayout(model_options_layout) 14 | 15 | model_options_layout.addWidget(QLabel("모델: ")) 16 | 17 | model_combo = QComboBox() 18 | model_combo.addItems(MODEL_INFO_DICT.keys()) 19 | model_combo.currentTextChanged.connect(self.on_model_changed) 20 | model_options_layout.addWidget(model_combo, stretch=2) 21 | 22 | self.dict_ui_settings["model"] = model_combo 23 | 24 | return layout -------------------------------------------------------------------------------- /app/gui/widget/status_bar.py: -------------------------------------------------------------------------------- 1 | # status_bar.py 2 | from config.strings import STRING 3 | 4 | class StatusBar: 5 | def __init__(self, status_bar): 6 | self.status_bar = status_bar 7 | self.status_state = None 8 | self.status_list_format = [] 9 | self.status_bar.messageChanged.connect( 10 | self.on_statusbar_message_changed) 11 | 12 | def set_statusbar_text(self, status_key="", list_format=[]): 13 | if status_key: 14 | self.status_state = status_key 15 | self.status_list_format = list_format 16 | else: 17 | status_key = self.status_state 18 | list_format = self.status_list_format 19 | 20 | message = STRING.LIST_STATSUBAR_STATE[status_key].format( 21 | *list_format) 22 | self.status_bar.showMessage(message) 23 | 24 | def on_statusbar_message_changed(self, t): 25 | if not t: 26 | self.set_statusbar_text() 27 | -------------------------------------------------------------------------------- /app/gui/layout/expand_layout.py: -------------------------------------------------------------------------------- 1 | 2 | from PyQt5.QtWidgets import QVBoxLayout, QLabel, QTextEdit 3 | 4 | from gui.widget.result_image_view import ResultImageView 5 | 6 | from config.paths import PATH_IMG_NO_IMAGE 7 | 8 | def init_expand_layout(self): 9 | vbox_expand = QVBoxLayout() 10 | vbox_expand.setContentsMargins(30, 30, 30, 30) 11 | 12 | image_result = ResultImageView(PATH_IMG_NO_IMAGE) 13 | image_result.setStyleSheet(""" 14 | background-color: white; 15 | background-position: center 16 | """) 17 | self.installEventFilter(image_result) 18 | self.image_result = image_result 19 | vbox_expand.addWidget(image_result, stretch=9) 20 | 21 | vbox_expand.addWidget(QLabel("결과창")) 22 | prompt_result = QTextEdit("") 23 | prompt_result.setPlaceholderText("이곳에 결과가 출력됩니다.") 24 | prompt_result.setReadOnly(True) 25 | prompt_result.setAcceptRichText(False) 26 | prompt_result.setAcceptDrops(False) 27 | vbox_expand.addWidget(prompt_result, stretch=1) 28 | self.prompt_result = prompt_result 29 | 30 | return vbox_expand -------------------------------------------------------------------------------- /app/config/strings.py: -------------------------------------------------------------------------------- 1 | class STRING: 2 | LIST_STATSUBAR_STATE = { 3 | "BEFORE_LOGIN": "로그인이 필요합니다", 4 | "LOGGINGIN": "로그인 시도중...", 5 | "LOGINED": "로그인 성공.", 6 | "IDLE": "대기 중", 7 | "GENEARTING": "생성 요청 중...", 8 | "LOADING": "불러오는 중...", 9 | "LOAD_COMPLETE": "불러오기 완료", 10 | "AUTO_GENERATING_COUNT": "자동생성 중... 총 {}장 중 {}번째", 11 | "AUTO_GENERATING_INF": "자동생성 중...", 12 | "AUTO_ERROR_WAIT": "생성 중 에러가 발생. {}초 뒤 다시 시작.", 13 | "AUTO_WAIT": "자동생성 딜레이를 기다리는 중... {}초" 14 | } 15 | 16 | ABOUT = """ 17 | 본진 : 18 | 아카라이브 AI그림 채널 https://arca.live/b/aiart 19 | 만든이 : 20 | https://arca.live/b/aiart @DCP 21 | 크레딧 : 22 | https://huggingface.co/baqu2213 23 | https://github.com/neggles/sd-webui-stealth-pnginfo/ 24 | """ 25 | 26 | LABEL_PROMPT = "프롬프트(Prompt)" 27 | LABEL_PROMPT_HINT = "이곳에 원하는 특징을 입력하세요.\n(예 - 1girl, Tendou Aris (Blue archive), happy)" 28 | LABEL_NPROMPT = "네거티브 프롬프트(Undesired Content)" 29 | LABEL_NPROMPT_HINT = "이곳에 원하지 않는 특징을 입력하세요.\n(예 - bad quality, low quality, lowres, displeasing)" 30 | LABEL_AISETTING = "생성 옵션(AI Settings)" -------------------------------------------------------------------------------- /app/gui/dialog/fileio_dialog.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtWidgets import QLabel, QVBoxLayout, QDialog 2 | from PyQt5.QtCore import Qt, QTimer 3 | 4 | from core.thread.loading_thread import LoadingThread 5 | 6 | class FileIODialog(QDialog): 7 | def __init__(self, text, func): 8 | super().__init__() 9 | self.text = text 10 | self.func = func 11 | self.init_ui() 12 | 13 | def init_ui(self): 14 | self.setWindowTitle("작업 중") 15 | 16 | layout = QVBoxLayout() 17 | self.progress_label = QLabel(self.text) 18 | layout.addWidget(self.progress_label) 19 | 20 | self.setLayout(layout) 21 | 22 | self.resize(200, 100) 23 | self.setWindowFlag(Qt.WindowCloseButtonHint, False) 24 | 25 | def showEvent(self, event): 26 | QTimer.singleShot(100, self.start_work) 27 | super().showEvent(event) 28 | 29 | def start_work(self): 30 | self.worker_thread = LoadingThread(self.func) 31 | self.worker_thread.finished.connect(self.on_finished) 32 | self.worker_thread.start() 33 | 34 | def on_finished(self, result): 35 | self.result = result 36 | self.accept() -------------------------------------------------------------------------------- /app/util/common_util.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | 4 | def strtobool(val): 5 | """Convert a string representation of truth to true (1) or false (0). 6 | True values are 'y', 'yes', 't', 'true', 'on', and '1'; false values 7 | are 'n', 'no', 'f', 'false', 'off', and '0'. Raises ValueError if 8 | 'val' is anything else. 9 | """ 10 | if isinstance(val, bool): 11 | return val 12 | 13 | val = val.lower() 14 | if val in ('y', 'yes', 't', 'true', 'on', '1'): 15 | return True 16 | elif val in ('n', 'no', 'f', 'false', 'off', '0'): 17 | return False 18 | else: 19 | raise ValueError("invalid truth value %r" % (val,)) 20 | 21 | 22 | def get_key_from_dict(dictionary, value): 23 | return next((key for key, val in dictionary.items() if val == value), None) 24 | 25 | 26 | def try_loads(json_str): 27 | try: 28 | json_str = json.loads(json_str) 29 | except Exception: 30 | json_str = None 31 | 32 | return json_str 33 | 34 | 35 | def try_dumps(target): 36 | try: 37 | result = json.dumps(target) 38 | except Exception: 39 | result = None 40 | 41 | return result 42 | -------------------------------------------------------------------------------- /app/util/ui_util.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtWidgets import QWidget, QPushButton, QGraphicsOpacityEffect 2 | 3 | def create_empty(minimum_width=1, minimum_height=1, fixed_height=0, maximum_height=0): 4 | w = QWidget() 5 | w.setMinimumWidth(minimum_width) 6 | w.setMinimumHeight(minimum_height) 7 | w.setStyleSheet("background-color:#00000000") 8 | if fixed_height != 0: 9 | w.setFixedHeight(fixed_height) 10 | if maximum_height != 0: 11 | w.setMaximumHeight(maximum_height) 12 | return w 13 | 14 | 15 | def add_button(hbox, text, callback, minimum_width=-1, maximum_width=-1, maximum_height=-1): 16 | button = QPushButton(text) 17 | button.pressed.connect(callback) 18 | if minimum_width != -1: 19 | button.setMinimumWidth(minimum_width) 20 | if maximum_width != -1: 21 | button.setMaximumWidth(maximum_width) 22 | if maximum_height != -1: 23 | button.setMaximumHeight(maximum_height) 24 | hbox.addWidget(button) 25 | return button 26 | 27 | 28 | def set_opacity(widget, opacity): 29 | opacity_effect = QGraphicsOpacityEffect(widget) 30 | opacity_effect.setOpacity(opacity) # 0은 완전히 투명함을 의미함 31 | widget.setGraphicsEffect(opacity_effect) -------------------------------------------------------------------------------- /app/config/paths.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | 5 | def resource_path(relative_path): 6 | try: 7 | base_path = sys._MEIPASS 8 | except Exception: 9 | base_path = os.path.abspath(".") 10 | 11 | return os.path.join(base_path, relative_path).replace("\\", "/") 12 | 13 | 14 | PATH_IMG_NO_IMAGE = resource_path("assets/no_image.png") 15 | PATH_IMG_OPEN_IMAGE = resource_path("assets/open_image.png") 16 | PATH_IMG_OPEN_FOLDER = resource_path("assets/open_folder.png") 17 | PATH_IMG_IMAGE_CLEAR = resource_path("assets/image_clear.png") 18 | PATH_IMG_TAGGER = resource_path("assets/tagger.png") 19 | PATH_IMG_GETTER = resource_path("assets/getter.png") 20 | PATH_IMG_ICON = resource_path("assets/icon.png") 21 | PATH_IMG_ICON_GETTER = resource_path("assets/icon_getter.png") 22 | PATH_IMG_ICON_TAGGER = resource_path("assets/icon_tagger.png") 23 | 24 | try: 25 | base_path = sys._MEIPASS 26 | PATH_CSV_TAG_COMPLETION = "danbooru_tags_post_count.csv" 27 | except Exception: 28 | PATH_CSV_TAG_COMPLETION = "assets/danbooru_tags_post_count.csv" 29 | 30 | DEFAULT_PATH = { 31 | "path_results": "results/", 32 | "path_wildcards": "wildcards/", 33 | "path_settings": "settings/", 34 | "path_models": "models/" 35 | } 36 | -------------------------------------------------------------------------------- /app/util/file_util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | def create_folder_if_not_exists(foldersrc): 5 | if not os.path.exists(foldersrc): 6 | os.makedirs(foldersrc) 7 | 8 | 9 | def get_imgcount_from_foldersrc(foldersrc): 10 | return len([file for file in os.listdir(foldersrc) if file.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))]) 11 | 12 | 13 | def get_filename_only(path): 14 | filename, _ = os.path.splitext(os.path.basename(path)) 15 | return filename 16 | 17 | 18 | def create_windows_filepath(base_path, filename, extension, max_length=150): 19 | # 파일 이름으로 사용할 수 없는 문자 제거 20 | cleaned_filename = filename.replace("\n", "") 21 | cleaned_filename = cleaned_filename.replace("\\", "") 22 | 23 | invalid_chars = r'<>:"/\|?*' 24 | cleaned_filename = ''.join( 25 | char for char in cleaned_filename if char not in invalid_chars) 26 | 27 | # 파일 이름의 최대 길이 제한 (확장자 길이 고려) 28 | max_filename_length = max_length - len(base_path) - len(extension) - 1 29 | if max_filename_length < 5: 30 | return None 31 | cleaned_filename = cleaned_filename[:max_filename_length] 32 | 33 | # 경로, 파일 이름, 확장자 합치기 34 | filepath = os.path.join(base_path, cleaned_filename + extension) 35 | 36 | return filepath -------------------------------------------------------------------------------- /app/util/tagger_util.py: -------------------------------------------------------------------------------- 1 | 2 | from PyQt5.QtWidgets import QMessageBox, QDialog 3 | from PIL import Image 4 | 5 | from gui.dialog.fileio_dialog import FileIODialog 6 | 7 | def predict_tag_from(window, filemode, target, with_dialog): 8 | result = "" 9 | 10 | target_model_name = window.settings.value("selected_tagger_model", '') 11 | if not target_model_name: 12 | QMessageBox.information( 13 | window, '경고', "먼저 본체 어플의 옵션에서 태깅 모델을 다운/선택 해주세요.") 14 | return "" 15 | else: 16 | window.dtagger.options["model_name"] = target_model_name 17 | 18 | if filemode == "src": 19 | target = Image.open(target) 20 | 21 | if with_dialog: 22 | loading_dialog = FileIODialog( 23 | "태그하는 중...", lambda: window.dtagger.tag(target)) 24 | if loading_dialog.exec_() == QDialog.Accepted: 25 | result = loading_dialog.result 26 | if not result: 27 | list_installed_model = window.dtagger.get_installed_models() 28 | if not (target_model_name in list_installed_model): 29 | window.settings.setValue("selected_tagger_model", '') 30 | else: 31 | try: 32 | result = window.dtagger.tag(target) 33 | except Exception as e: 34 | print(e) 35 | 36 | return result 37 | -------------------------------------------------------------------------------- /app/core/thread/completiontagload_thread.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtWidgets import QMainWindow 2 | from PyQt5.QtCore import QThread, pyqtSignal 3 | 4 | from config.paths import PATH_CSV_TAG_COMPLETION 5 | 6 | def _on_load_completiontag_success(window, tag_list): 7 | if tag_list: 8 | window.completiontag_list = tag_list 9 | 10 | target_code = ["prompt", "negative_prompt"] 11 | for code in target_code: 12 | window.dict_ui_settings[code].start_complete_mode(tag_list) 13 | 14 | window.dict_ui_settings["characterPrompts"].refresh_completiontag_list() 15 | 16 | class CompletionTagLoadThread(QThread): 17 | on_load_completiontag_success = pyqtSignal(QMainWindow, list) 18 | 19 | def __init__(self, parent): 20 | super(CompletionTagLoadThread, self).__init__(parent) 21 | self.parent = parent 22 | self.on_load_completiontag_success.connect( 23 | _on_load_completiontag_success) 24 | 25 | def run(self): 26 | try: 27 | with open(PATH_CSV_TAG_COMPLETION, "r", encoding='utf8') as f: 28 | tag_list = f.readlines() 29 | if tag_list: 30 | self.on_load_completiontag_success.emit(self.parent, tag_list) 31 | except Exception: 32 | pass 33 | 34 | def stop(self): 35 | self.is_dead = True 36 | self.quit() 37 | -------------------------------------------------------------------------------- /app/gui/widget/result_image_view.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtWidgets import QLabel 2 | from PyQt5.QtGui import QPixmap 3 | from PyQt5.QtCore import Qt, pyqtSignal, QEvent, pyqtSignal, QTimer 4 | 5 | from util.image_util import pil2pixmap 6 | 7 | 8 | class ResultImageView(QLabel): 9 | clicked = pyqtSignal() 10 | 11 | def __init__(self, first_src): 12 | super(ResultImageView, self).__init__() 13 | self.set_custom_pixmap(first_src) 14 | self.setAlignment(Qt.AlignCenter) 15 | 16 | def set_custom_pixmap(self, img_obj): 17 | if isinstance(img_obj, str): 18 | self.pixmap = QPixmap(img_obj) 19 | else: 20 | self.pixmap = pil2pixmap(img_obj) 21 | self.refresh_size() 22 | 23 | def refresh_size(self): 24 | self.setPixmap(self.pixmap.scaled( 25 | self.width(), self.height(), 26 | aspectRatioMode=Qt.KeepAspectRatio, 27 | transformMode=Qt.SmoothTransformation)) 28 | self.setMinimumWidth(100) 29 | 30 | def setFixedSize(self, qsize): 31 | super(ResultImageView, self).setFixedSize(qsize) 32 | QTimer.singleShot(20, self.refresh_size) 33 | 34 | def eventFilter(self, obj, event): 35 | if event.type() == QEvent.Resize: 36 | self.refresh_size() 37 | return True 38 | return super(ResultImageView, self).eventFilter(obj, event) 39 | 40 | def mousePressEvent(self, ev): 41 | self.clicked.emit() -------------------------------------------------------------------------------- /app/config/consts.py: -------------------------------------------------------------------------------- 1 | TITLE_NAME = "NAI Auto Generator" 2 | TOP_NAME = "dcp_arca" 3 | APP_NAME = "nag_gui" 4 | 5 | MAX_COUNT_FOR_WHILELOOP = 10 6 | 7 | DEFAULT_RESOLUTION = "Square (640x640)" 8 | 9 | RESOLUTION_ITEMS = [ 10 | "NORMAL", 11 | "Portrait (832x1216)", 12 | "Landscape (1216x832)", 13 | "Square (1024x1024)", 14 | "LARGE", 15 | "Portrait (1024x1536)", 16 | "Landscape (1536x1024)", 17 | "Square (1472x1472)", 18 | "WALLPAPER", 19 | "Portrait (1088x1920)", 20 | "Landscape (1920x1088)", 21 | "SMALL", 22 | "Portrait (512x768)", 23 | "Landscape (768x512)", 24 | "Square (640x640)", 25 | "CUSTOM", 26 | "Custom", 27 | ] 28 | RESOLUTION_ITEMS_NOT_SELECTABLES = [0, 4, 8, 11, 15] 29 | 30 | RESOLUTION_FAMILIY = { 31 | 0: ["Portrait (832x1216)", "Landscape (1216x832)", "Square (1024x1024)"], 32 | 1: ["Portrait (1024x1536)", "Landscape (1536x1024)", "Square (1472x1472)"], 33 | 2: ["Portrait (1088x1920)", "Landscape (1920x1088)"], 34 | 3: ["Portrait (512x768)", "Landscape (768x512)", "Square (640x640)"], 35 | 4: [] 36 | } 37 | RESOLUTION_FAMILIY_MASK = [ 38 | -1, 39 | 0, 40 | 0, 41 | 0, 42 | -1, 43 | 1, 44 | 1, 45 | 1, 46 | -1, 47 | 2, 48 | 2, 49 | -1, 50 | 3, 51 | 3, 52 | 3, 53 | -1, 54 | 4 55 | ] 56 | 57 | DEFAULT_TAGGER_MODEL = "wd-v1-4-moat-tagger-v2" 58 | LIST_TAGGER_MODEL = ("wd-v1-4-moat-tagger-v2", 59 | "wd-v1-4-convnext-tagger-v2", "wd-v1-4-convnext-tagger", 60 | "wd-v1-4-convnextv2-tagger-v2", "wd-v1-4-vit-tagger-v2") 61 | 62 | -------------------------------------------------------------------------------- /app/config/themes.py: -------------------------------------------------------------------------------- 1 | class COLOR: 2 | WHITE = "#898a96" 3 | GRAY = "#7A7A7A" 4 | BRIGHT = "#212335" 5 | MIDIUM = "#1A1C2E" 6 | DARK = "#101224" 7 | BUTTON = "#F5F3C2" 8 | BUTTON_DSIABLED = "#999682" 9 | BUTTON_SELECTED = "#F5B5B5" 10 | 11 | 12 | MAIN_STYLESHEET = """ 13 | QWidget { 14 | color: white; 15 | background-color: """ + COLOR.BRIGHT + """; 16 | } 17 | QTextEdit { 18 | background-color: """ + COLOR.DARK + """; 19 | } 20 | QLineEdit { 21 | background-color: """ + COLOR.DARK + """; 22 | border: 1px solid """ + COLOR.GRAY + """; 23 | } 24 | QComboBox { 25 | background-color: """ + COLOR.DARK + """; 26 | border: 1px solid """ + COLOR.GRAY + """; 27 | } 28 | QComboBox QAbstractItemView { 29 | border: 2px solid """ + COLOR.GRAY + """; 30 | selection-background-color: black; 31 | } 32 | QPushButton { 33 | color:black; 34 | background-color: """ + COLOR.BUTTON + """; 35 | } 36 | QPushButton:disabled { 37 | background-color: """ + COLOR.BUTTON_DSIABLED + """; 38 | } 39 | QSplitter::handle { 40 | background-color: qlineargradient(spread:pad, x1:0, y1:0, x2:1, y2:0, 41 | stop:0 rgba(0, 0, 0, 0), 42 | stop:0.1 rgba(0, 0, 0, 0), 43 | stop:0.1001 rgba(255, 255, 255, 255), 44 | stop:0.2 rgba(255, 255, 255, 255), 45 | stop:0.2001 rgba(0, 0, 0, 0), 46 | stop:0.8 rgba(0, 0, 0, 0), 47 | stop:0.8001 rgba(255, 255, 255, 255), 48 | stop:0.9 rgba(255, 255, 255, 255), 49 | stop:0.9001 rgba(0, 0, 0, 0)); 50 | image: url(:/images/splitter.png); 51 | } 52 | """ 53 | -------------------------------------------------------------------------------- /app/gui/layout/prompt_layout.py: -------------------------------------------------------------------------------- 1 | 2 | from PyQt5.QtWidgets import QTextEdit, QLabel, QVBoxLayout, QHBoxLayout, QSizePolicy 3 | 4 | from core.worker.completer import CompletionTextEdit 5 | 6 | from util.ui_util import create_empty 7 | from util.string_util import prettify_naidict 8 | 9 | from gui.layout.character_prompt_layout import CharacterPromptsContainer 10 | 11 | from config.strings import STRING 12 | 13 | def create_prompt_layout(self, title_text): 14 | hbox_prompt_title = QHBoxLayout() 15 | 16 | label = QLabel(title_text) 17 | hbox_prompt_title.addWidget(label) 18 | 19 | return hbox_prompt_title 20 | 21 | def create_prompt_edit(self, placeholder_text, code, minimum_height): 22 | textedit = CompletionTextEdit() 23 | textedit.setPlaceholderText(placeholder_text) 24 | textedit.setSizePolicy(QSizePolicy.Expanding, QSizePolicy.Expanding) 25 | textedit.setMinimumHeight(minimum_height) 26 | self.dict_ui_settings[code] = textedit 27 | 28 | return textedit 29 | 30 | class PromptLayout(): 31 | def __init__(self, parent): 32 | self.parent = parent 33 | self.parent.prompt_layout = self 34 | 35 | def init(self): 36 | parent = self.parent 37 | 38 | vbox = QVBoxLayout() 39 | 40 | vbox.addLayout(create_prompt_layout(self, STRING.LABEL_PROMPT)) 41 | 42 | vbox.addWidget(create_prompt_edit(parent, STRING.LABEL_PROMPT_HINT, "prompt", 43 | minimum_height=200), stretch=3) 44 | 45 | vbox.addWidget(create_empty(minimum_height=4)) 46 | 47 | vbox.addLayout(create_prompt_layout( 48 | parent, STRING.LABEL_NPROMPT)) 49 | 50 | vbox.addWidget(create_prompt_edit( 51 | parent, STRING.LABEL_NPROMPT_HINT, "negative_prompt", 52 | minimum_height=100), stretch=2) 53 | 54 | vbox.addWidget(create_empty(minimum_height=4)) 55 | 56 | parent.character_prompts_container = CharacterPromptsContainer(parent) 57 | vbox.addWidget(parent.character_prompts_container, stretch=1) 58 | 59 | vbox.addWidget(create_empty(minimum_height=4)) 60 | 61 | return vbox -------------------------------------------------------------------------------- /app/gui/widget/menu_bar.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtWidgets import QAction 2 | 3 | from gui.dialog.etc_dialog import show_file_dialog 4 | 5 | class MenuBar: 6 | def __init__(self, parent): 7 | self.parent = parent 8 | self.menubar = parent.menuBar() 9 | self.menubar.setNativeMenuBar(False) 10 | self.init_menu_bar() 11 | 12 | def init_menu_bar(self): 13 | openAction = QAction('파일 열기(Open file)', self.parent) 14 | openAction.setShortcut('Ctrl+O') 15 | openAction.triggered.connect(lambda: show_file_dialog(self.parent, "file")) 16 | 17 | loginAction = QAction('로그인(Log in)', self.parent) 18 | loginAction.setShortcut('Ctrl+L') 19 | loginAction.triggered.connect(self.parent.show_login_dialog) 20 | 21 | optionAction = QAction('옵션(Option)', self.parent) 22 | optionAction.setShortcut('Ctrl+U') 23 | optionAction.triggered.connect(self.parent.show_option_dialog) 24 | 25 | exitAction = QAction('종료(Exit)', self.parent) 26 | exitAction.setShortcut('Ctrl+W') 27 | exitAction.triggered.connect(self.parent.quit_app) 28 | 29 | aboutAction = QAction('만든 이(About)', self.parent) 30 | aboutAction.triggered.connect(self.parent.show_about_dialog) 31 | 32 | getterAction = QAction('이미지 정보 확인기(Info Getter)', self.parent) 33 | getterAction.setShortcut('Ctrl+I') 34 | getterAction.triggered.connect(self.parent.on_click_getter) 35 | 36 | taggerAction = QAction('태그 확인기(Danbooru Tagger)', self.parent) 37 | taggerAction.setShortcut('Ctrl+T') 38 | taggerAction.triggered.connect(self.parent.on_click_tagger) 39 | 40 | # 파일 메뉴 생성 41 | file_menu = self.menubar.addMenu('&파일(Files)') 42 | file_menu.addAction(openAction) 43 | file_menu.addAction(loginAction) 44 | file_menu.addAction(optionAction) 45 | file_menu.addAction(exitAction) 46 | 47 | # 도구 메뉴 생성 48 | tool_menu = self.menubar.addMenu('&도구(Tools)') 49 | tool_menu.addAction(getterAction) 50 | tool_menu.addAction(taggerAction) 51 | 52 | # 기타 메뉴 생성 53 | etc_menu = self.menubar.addMenu('&기타(Etc)') 54 | etc_menu.addAction(aboutAction) 55 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # NAI-Auto-Generator 2 | 제목 없음 3 | 4 | # 특징 5 | 1. NAI UI를 유사하게 구현한 윈도우 어플리케이션. 6 | 2. 세팅 값을 텍스트 파일로 쉽게 내보내고 들여올 수 있음. 7 | 3. img2img, 멀티 vibe 기능 사용 가능. 8 | * 폴더를 드래그 드랍하거나, 열기 버튼을 눌러 폴더 안의 이미지를 순차적/랜덤으로 불러오기 가능. 9 | 4. 와일드카드 기능 사용 가능. (ex- \_\_와일드카드 이름\_\_, \_\_폴더이름1/폴더이름2/와일드카드 이름\_\_) 10 | 5. 프롬프트 내 인스턴스 랜덤 구문 사용 가능. (ex- <1|2|3>) 11 | * 1.1.0 부터 와일드카드와 인스턴스 구문은 서로 교차해서 사용 가능. 12 | 6. 웹페이지 이미지나 파일을 드래그 드랍하여 세팅 정보를 바로 가져오기 가능. 13 | 7. 자동 생성 기능 지원. 14 | 8. 정보를 가져오는 이미지 정보 확인기 / 단부루 태그를 읽어오는 태그 확인기 지원 15 | 9. v4 지원, v4 캐릭터 프롬프트 지원 16 | 17 | 18 | # 활용 19 | 20 | ![img-ezgif com-video-to-gif-converter](https://github.com/DCP-arca/NAI-Auto-Generator/assets/127241088/4e246046-b5a5-41d5-abd8-4f0e0622715f) 21 | 22 | img2img와 vibe transfer를 지원합니다. 23 | 폴더모드를 눌러 폴더 내부의 이미지들을 자동으로 선택하게 할 수도 있습니다. 24 | 25 | 26 | 27 | ![drag_drop-ezgif com-video-to-gif-converter](https://github.com/DCP-arca/NAI-Auto-Generator/assets/127241088/ae88a709-05fc-4387-9cbb-c61be980622d) 28 | 29 | 웹(주로 arca.live)/이미지파일/세팅텍스트파일을 드래그 드랍해서 세팅 가져올 수 있습니다. 30 | 31 | 32 | 33 | ![splitter-ezgif com-video-to-gif-converter](https://github.com/DCP-arca/NAI-Auto-Generator/assets/127241088/851aa7c4-c756-4e5e-835a-936be201716c) 34 | 35 | 이미지 뷰어는 하단의 >>를 눌러 키고 끌 수 있으며, 이동가능합니다. 36 | 37 | 38 | 39 | ![tagger-ezgif com-video-to-gif-converter](https://github.com/DCP-arca/NAI-Auto-Generator/assets/127241088/0f4ef2e2-121d-4435-b92a-279b2406ef4c) 40 | ![infogetter-ezgif com-video-to-gif-converter](https://github.com/DCP-arca/NAI-Auto-Generator/assets/127241088/18f837f4-f2df-42ef-99e5-c92cf51c1283) 41 | 42 | 하단의 Info/Tag 버튼이나, 상단의 메뉴바, 또는 직접 .exe를 실행하여, 정보를 가져오는 이미지 정보 확인기 / 단부루 태그를 읽어오는 태그 확인기를 사용할 수 있습니다. 43 | 44 | 45 | # 주의사항 46 | 자동 생성의 과도한 사용시 밴을 당할 수 있습니다. 주의하세요. 47 | 가능한 자동 생성버튼을 눌렀을 때 나오는 '지연 시간'을 설정하길 추천합니다. 48 | 해당 프로그램을 사용함으로써 발생하는 책임은 모두 사용자에게 있습니다. 49 | 50 | 51 | # 크레딧 52 | https://github.com/neggles/sd-webui-stealth-pnginfo/ 53 | 54 | https://huggingface.co/baqu2213/ 55 | 56 | https://github.com/pythongosssss/ComfyUI-WD14-Tagger/ 57 | 58 | https://huggingface.co/SmilingWolf 59 | 60 | 61 | # 릴리즈 62 | https://github.com/DCP-arca/NAI-Auto-Generator/releases 63 | 64 | onefile버전을 받으면 하나의 .exe파일로 지원합니다. 65 | 반대쪽 버전을 받으면 용량이 크지만, 어플리케이션 시작이 조금 더 빠르며, 이미지 정보확인기와 태그 확인기를 exe로 단일 실행할 수 있습니다. 66 | 67 | -------------------------------------------------------------------------------- /app/util/image_util.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | import io 4 | import base64 5 | 6 | from PIL import Image 7 | 8 | from PyQt5.QtGui import QImage, QPixmap 9 | from PyQt5.QtCore import QBuffer 10 | 11 | 12 | def convert_qimage_to_imagedata(qimage): 13 | try: 14 | buf = QBuffer() 15 | buf.open(QBuffer.ReadWrite) 16 | qimage.save(buf, "PNG") 17 | pil_im = Image.open(io.BytesIO(buf.data())) 18 | 19 | buf = io.BytesIO() 20 | pil_im.save(buf, format='png', quality=100) 21 | return base64.b64encode(buf.getvalue()).decode("utf-8") 22 | except Exception: 23 | return "" 24 | 25 | 26 | def pick_imgsrc_from_foldersrc(foldersrc, index, sort_order): 27 | files = [file for file in os.listdir(foldersrc) if file.lower( 28 | ).endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp'))] 29 | 30 | is_reset = False 31 | if index != 0 and index % len(files) == 0: 32 | is_reset = True 33 | 34 | # 파일들을 정렬 35 | if sort_order == '오름차순': 36 | files.sort() 37 | elif sort_order == '내림차순': 38 | files.sort(reverse=True) 39 | elif sort_order == '랜덤': 40 | random.seed(random.randint(0, 1000000)) 41 | random.shuffle(files) 42 | is_reset = False 43 | 44 | # 인덱스가 파일 개수를 초과하는 경우 45 | while index >= len(files): 46 | index -= len(files) 47 | 48 | # 정렬된 파일 리스트에서 인덱스에 해당하는 파일의 주소 반환 49 | return os.path.join(foldersrc, files[index]), is_reset 50 | 51 | 52 | def pil2pixmap(im): 53 | if im.mode == "RGB": 54 | r, g, b = im.split() 55 | im = Image.merge("RGB", (b, g, r)) 56 | elif im.mode == "RGBA": 57 | r, g, b, a = im.split() 58 | im = Image.merge("RGBA", (b, g, r, a)) 59 | elif im.mode == "L": 60 | im = im.convert("RGBA") 61 | # Bild in RGBA konvertieren, falls nicht bereits passiert 62 | im2 = im.convert("RGBA") 63 | data = im2.tobytes("raw", "RGBA") 64 | qim = QImage( 65 | data, im.size[0], im.size[1], QImage.Format_ARGB32) 66 | pixmap = QPixmap.fromImage(qim) 67 | return pixmap 68 | 69 | 70 | def convert_src_to_imagedata(img_path, quality=100): 71 | try: 72 | img = Image.open(img_path) 73 | buf = io.BytesIO() 74 | img.save(buf, format='png', quality=quality) 75 | return base64.b64encode(buf.getvalue()).decode("utf-8") 76 | except Exception as e: 77 | return "" 78 | -------------------------------------------------------------------------------- /app/gui/dialog/generate_dialog.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtWidgets import QLabel, QLineEdit, QCheckBox, QVBoxLayout, QHBoxLayout, QPushButton, QDialog 2 | 3 | class GenerateDialog(QDialog): 4 | def __init__(self, parent): 5 | super().__init__() 6 | self.parent = parent 7 | self.initUI() 8 | 9 | def initUI(self): 10 | parent_pos = self.parent.pos() 11 | 12 | self.setWindowTitle('자동 생성') 13 | self.move(parent_pos.x() + 50, parent_pos.y() + 50) 14 | self.setFixedSize(400, 200) 15 | 16 | layout = QVBoxLayout() 17 | self.setLayout(layout) 18 | 19 | layout_count = QHBoxLayout() 20 | layout.addLayout(layout_count) 21 | 22 | label_count = QLabel("생성 횟수(빈칸시 무한) : ", self) 23 | layout_count.addWidget(label_count, 1) 24 | lineedit_count = QLineEdit("") 25 | lineedit_count.setMaximumWidth(40) 26 | self.lineedit_count = lineedit_count 27 | layout_count.addWidget(lineedit_count, 1) 28 | 29 | layout_delay = QHBoxLayout() 30 | layout.addLayout(layout_delay) 31 | 32 | label_delay = QLabel("지연 시간(매 생성시, 에러시 대기시간) : ", self) 33 | layout_delay.addWidget(label_delay, 1) 34 | lineedit_delay = QLineEdit("3") 35 | lineedit_delay.setMaximumWidth(40) 36 | self.lineedit_delay = lineedit_delay 37 | layout_delay.addWidget(lineedit_delay, 1) 38 | 39 | checkbox_ignoreerror = QCheckBox("에러 발생 시에도 계속 하기") 40 | checkbox_ignoreerror.setChecked(True) 41 | self.checkbox_ignoreerror = checkbox_ignoreerror 42 | layout.addWidget(checkbox_ignoreerror, 1) 43 | 44 | layout_buttons = QHBoxLayout() 45 | layout.addLayout(layout_buttons) 46 | 47 | start_button = QPushButton("시작") 48 | start_button.clicked.connect(self.on_start_button_clicked) 49 | layout_buttons.addWidget(start_button) 50 | self.start_button = start_button 51 | 52 | close_button = QPushButton("닫기") 53 | close_button.clicked.connect(self.on_close_button_clicked) 54 | layout_buttons.addWidget(close_button) 55 | self.close_button = close_button 56 | 57 | def on_start_button_clicked(self): 58 | self.count = self.lineedit_count.text() or None 59 | self.delay = self.lineedit_delay.text() or None 60 | self.ignore_error = self.checkbox_ignoreerror.isChecked() 61 | self.accept() 62 | 63 | def on_close_button_clicked(self): 64 | self.reject() 65 | -------------------------------------------------------------------------------- /app/gui/dialog/etc_dialog.py: -------------------------------------------------------------------------------- 1 | import os 2 | import json 3 | from PyQt5.QtWidgets import QFileDialog, QMessageBox 4 | 5 | from config.paths import DEFAULT_PATH 6 | 7 | def show_prompt_dialog(self, title, prompt, nprompt): 8 | QMessageBox.about(self, title, 9 | "프롬프트:\n" + 10 | prompt + 11 | "\n\n" + 12 | "네거티브 프롬프트:\n" + 13 | nprompt) 14 | 15 | def show_file_dialog(self): 16 | select_dialog = QFileDialog() 17 | select_dialog.setFileMode(QFileDialog.ExistingFile) 18 | target_type = '이미지, 텍스트 파일(*.txt *.png *.webp)' 19 | fname = select_dialog.getOpenFileName( 20 | self, '불러올 파일을 선택해 주세요.', '', target_type) 21 | 22 | if fname[0]: 23 | fname = fname[0] 24 | 25 | if fname.endswith(".png") or fname.endswith(".webp"): 26 | self.get_image_info_bysrc(fname) 27 | elif fname.endswith(".txt"): 28 | self.get_image_info_bytxt(fname) 29 | else: 30 | QMessageBox.information( 31 | self, '경고', "png, webp, txt 파일만 가능합니다.") 32 | return 33 | 34 | def show_openfolder_dialog(self, mode): 35 | select_dialog = QFileDialog() 36 | select_dialog.setFileMode(QFileDialog.Directory) 37 | select_dialog.setOption(QFileDialog.Option.ShowDirsOnly) 38 | fname = select_dialog.getExistingDirectory( 39 | self, mode + '모드로 열 폴더를 선택해주세요.', '') 40 | 41 | if fname: 42 | if os.path.isdir(fname): 43 | self.set_imagefolder_as_param(mode, fname) 44 | else: 45 | QMessageBox.information( 46 | self, '경고', "폴더만 선택 가능합니다.") 47 | 48 | def show_setting_load_dialog(self): 49 | path = self.settings.value( 50 | "path_settings", DEFAULT_PATH["path_settings"]) 51 | 52 | select_dialog = QFileDialog() 53 | select_dialog.setFileMode(QFileDialog.ExistingFile) 54 | path, _ = select_dialog.getOpenFileName( 55 | self, "불러올 세팅 파일을 선택해주세요", path, "Txt File (*.txt)") 56 | if path: 57 | is_success = self._load_settings(path) 58 | 59 | if not is_success: 60 | QMessageBox.information( 61 | self, '경고', "세팅을 불러오는데 실패했습니다.\n\n") 62 | 63 | def show_setting_save_dialog(self): 64 | path = self.settings.value( 65 | "path_settings", DEFAULT_PATH["path_settings"]) 66 | path, _ = QFileDialog.getSaveFileName( 67 | self, "세팅 파일을 저장할 곳을 선택해주세요", path, "Txt File (*.txt)") 68 | if path: 69 | try: 70 | json_str = json.dumps(self.get_data_from_savetarget_ui()) 71 | with open(path, "w", encoding="utf8") as f: 72 | f.write(json_str) 73 | except Exception as e: 74 | print(e) 75 | QMessageBox.information( 76 | self, '경고', "세팅 저장에 실패했습니다.\n\n" + str(e)) 77 | -------------------------------------------------------------------------------- /app/gui/widget/custom_slider_widget.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtWidgets import QSlider, QLabel, QLineEdit, QHBoxLayout 2 | from PyQt5.QtCore import Qt 3 | from PyQt5.QtGui import QIntValidator 4 | 5 | 6 | class CustomSliderWidget(QHBoxLayout): 7 | def __init__(self, **option_dict): 8 | assert all(key in option_dict for key in [ 9 | "title", "min_value", "max_value", "default_value", "edit_width", "mag", "slider_text_lambda"]) 10 | super(CustomSliderWidget, self).__init__() 11 | 12 | label = QLabel(option_dict["title"]) 13 | slider = QSlider(Qt.Horizontal) 14 | slider.setMinimum(option_dict["min_value"]) 15 | slider.setMaximum(option_dict["max_value"]) 16 | 17 | edit = CustomSliderWidget.CustomLineEdit(option_dict, slider) 18 | 19 | if "gui_nobackground" in option_dict and option_dict["gui_nobackground"]: 20 | label.setStyleSheet("QLabel{background-color:#00000000}") 21 | slider.setStyleSheet("QSlider{background-color:#00000000}") 22 | 23 | self.addWidget(label) 24 | self.addWidget(edit) 25 | if "enable_percent_label" in option_dict and option_dict["enable_percent_label"]: 26 | self.addWidget(QLabel("%")) 27 | self.addWidget(slider) 28 | 29 | self.slider = slider 30 | self.edit = edit 31 | 32 | slider.setValue(int(float(option_dict["default_value"]))) 33 | 34 | class CustomLineEdit(QLineEdit): 35 | def __init__(self, option_dict, target_slider): 36 | super(QLineEdit, self).__init__(str(option_dict["default_value"])) 37 | self.min_value = option_dict["min_value"] 38 | self.max_value = option_dict["max_value"] 39 | self.mag = option_dict["mag"] 40 | self.target_slider = target_slider 41 | self.slider_text_lambda = option_dict["slider_text_lambda"] 42 | 43 | self.setMinimumWidth(option_dict["edit_width"]) 44 | self.setMaximumWidth(option_dict["edit_width"]) 45 | self.setAlignment(Qt.AlignCenter) 46 | self.setValidator(QIntValidator(0, 100)) 47 | 48 | target_slider.valueChanged.connect( 49 | lambda value: self.setText(self.slider_text_lambda(value))) 50 | self.returnPressed.connect( 51 | self.on_enter_or_focusout) 52 | 53 | def on_enter_or_focusout(self): 54 | value = self.text() 55 | if not value: 56 | value = self.min_value 57 | value = max(self.min_value, min( 58 | self.max_value, float(value))) 59 | value *= self.mag 60 | self.setText(self.slider_text_lambda(value)) 61 | self.target_slider.setValue(int(value)) 62 | 63 | def focusOutEvent(self, event): 64 | super(CustomSliderWidget.CustomLineEdit, self).focusOutEvent(event) 65 | self.on_enter_or_focusout() 66 | 67 | def setText(self, text): 68 | super(CustomSliderWidget.CustomLineEdit, self).setText(text) 69 | self.target_slider.setValue(int(float(text) * self.mag)) -------------------------------------------------------------------------------- /app/core/worker/wildcard_applier.py: -------------------------------------------------------------------------------- 1 | import os 2 | import random 3 | 4 | MAX_TRY_AMOUNT = 10 5 | 6 | 7 | class WildcardApplier(): 8 | def __init__(self, src_wildcards_folder): 9 | self.src_wildcards_folder = src_wildcards_folder 10 | self._wildcards_dict = {} 11 | 12 | def set_src(self, src): 13 | self.src_wildcards_folder = src 14 | 15 | def load_wildcards(self): 16 | self._wildcards_dict.clear() 17 | 18 | for dirpath, dname_list, fname_list in os.walk(self.src_wildcards_folder): 19 | path = "" # path for wildcards 20 | path = dirpath.replace(self.src_wildcards_folder, "") 21 | path = path.replace("\\", "/") + "/" 22 | path = path[1:] 23 | 24 | for filename in fname_list: 25 | if filename.endswith(".txt"): 26 | src = os.path.join(dirpath, filename) 27 | with open(src, "r", encoding="utf8") as f: 28 | lines = f.readlines() 29 | if lines: 30 | onlyname = os.path.splitext( 31 | os.path.basename(filename))[0] 32 | key = path + onlyname 33 | self._wildcards_dict[key.lower()] = lines 34 | 35 | def _apply_wildcard_once(self, target_str, except_list=[]): 36 | result = target_str 37 | 38 | applied_wildcard_list = [] 39 | prev_point = 0 40 | while "__" in result: 41 | p_left = result.find("__", prev_point) 42 | if p_left == -1: 43 | break 44 | 45 | p_right = result.find("__", p_left + 1) 46 | if p_right == -1: 47 | print("Warning : A single __ exists") 48 | break 49 | 50 | str_left = result[0:p_left] 51 | str_center = result[p_left + 2:p_right].lower() 52 | str_right = result[p_right + 2:len(result)] 53 | 54 | if str_center in self._wildcards_dict and not (str_center in except_list): 55 | wc_list = self._wildcards_dict[str_center] 56 | str_center = wc_list[random.randrange(0, len(wc_list))].strip() 57 | 58 | applied_wildcard_list.append(str_center) 59 | else: 60 | print("Warning : Unknown wildcard", str_center) 61 | str_center = "__" + str_center + "__" 62 | 63 | result_left = str_left + str_center 64 | prev_point = len(result_left) + 1 65 | 66 | result = result_left + str_right 67 | 68 | return result, applied_wildcard_list 69 | 70 | def apply_wildcards(self, target_str): 71 | self.load_wildcards() 72 | 73 | result = target_str 74 | 75 | index = 0 76 | except_list = [] 77 | while True: 78 | result, applied_wildcard_list = self._apply_wildcard_once( 79 | result, except_list) 80 | 81 | except_list.extend(applied_wildcard_list) 82 | if len(applied_wildcard_list) == 0: 83 | break 84 | 85 | index += 1 86 | if index > MAX_TRY_AMOUNT: 87 | print("Warning : Too much recursion") 88 | break 89 | 90 | return result 91 | 92 | 93 | if __name__ == "__main__": 94 | wa = WildcardApplier("wildcards") 95 | 96 | # result = wa.apply_wildcards("") 97 | 98 | # print(result) 99 | -------------------------------------------------------------------------------- /app/core/worker/naiinfo_getter.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | import json 3 | 4 | from core.worker.stealth_pnginfo import read_info_from_image_stealth 5 | 6 | TARGETKEY_NAIDICT_OPTION = ("steps", "height", "width", 7 | "scale", "seed", "sampler", "n_samples", "sm", "sm_dyn") 8 | 9 | 10 | def _get_infostr_from_img(img): 11 | exif = None 12 | pnginfo = None 13 | 14 | # exif 15 | if img.info: 16 | try: 17 | exif = json.dumps(img.info) 18 | except Exception as e: 19 | print("[_get_infostr_from_img]", e) 20 | 21 | # stealth pnginfo 22 | try: 23 | pnginfo = read_info_from_image_stealth(img) 24 | except Exception as e: 25 | print("[_get_infostr_from_img]", e) 26 | 27 | return exif, pnginfo 28 | 29 | 30 | def _get_exifdict_from_infostr(info_str): 31 | try: 32 | infoDict = json.loads(info_str) 33 | if ('Comment' in infoDict): 34 | exif_dict = json.loads(infoDict['Comment']) 35 | return exif_dict 36 | except Exception as e: 37 | print("[_get_exifdict_from_infostr]", e) 38 | 39 | return None 40 | 41 | 42 | def _get_naidict_from_exifdict(exif_dict): 43 | try: 44 | nai_dict = {} 45 | nai_dict["prompt"] = exif_dict["prompt"].strip() 46 | nai_dict["negative_prompt"] = exif_dict["uc"].strip() if "uc" in exif_dict else exif_dict["negative_prompt"].strip() 47 | option_dict = {} 48 | for key in TARGETKEY_NAIDICT_OPTION: 49 | if key in exif_dict.keys(): 50 | option_dict[key] = exif_dict[key] 51 | nai_dict["option"] = option_dict 52 | 53 | etc_dict = {} 54 | for key in exif_dict.keys(): 55 | if key in TARGETKEY_NAIDICT_OPTION + ("uc", "prompt"): 56 | continue 57 | etc_dict[key] = exif_dict[key] 58 | nai_dict["etc"] = etc_dict 59 | return nai_dict 60 | except Exception as e: 61 | print("[_get_naidict_from_exifdict]", e) 62 | 63 | return None 64 | 65 | 66 | def get_naidict_from_file(src): 67 | try: 68 | img = Image.open(src) 69 | img.load() 70 | except Exception as e: 71 | print("[get_naidict_from_file]", e) 72 | return None, 1 73 | 74 | return get_naidict_from_img(img) 75 | 76 | 77 | def get_naidict_from_txt(src): 78 | try: 79 | with open(src, "r", encoding="utf8") as f: 80 | info_str = f.read() 81 | ed = json.loads(info_str) 82 | except Exception as e: 83 | print("[get_naidict_from_txt]", e) 84 | return info_str or "", 1 85 | 86 | nd = _get_naidict_from_exifdict(ed) 87 | if not nd: 88 | return ed, 2 89 | else: 90 | return nd, 0 91 | 92 | 93 | def get_naidict_from_img(img): 94 | exif, pnginfo = _get_infostr_from_img(img) 95 | if not exif and not pnginfo: 96 | return None, 1 97 | 98 | ed1 = _get_exifdict_from_infostr(exif) 99 | ed2 = _get_exifdict_from_infostr(pnginfo) 100 | if not ed1 and not ed2: 101 | return exif or pnginfo, 2 102 | 103 | nd1 = _get_naidict_from_exifdict(ed1) 104 | nd2 = _get_naidict_from_exifdict(ed2) 105 | if not nd1 and not nd2: 106 | return exif or pnginfo, 3 107 | 108 | if nd1: 109 | return nd1, 0 110 | elif nd2: 111 | return nd2, 0 112 | 113 | 114 | if __name__ == "__main__": 115 | src = "settings/aris_noimage.txt" 116 | 117 | nd = get_naidict_from_txt(src) 118 | print(nd) -------------------------------------------------------------------------------- /app/util/string_util.py: -------------------------------------------------------------------------------- 1 | import random 2 | import json 3 | from collections import OrderedDict 4 | 5 | from config.consts import MAX_COUNT_FOR_WHILELOOP 6 | 7 | def pickedit_lessthan_str(original_str): 8 | try_count = 0 9 | 10 | edited_str = original_str 11 | while try_count < MAX_COUNT_FOR_WHILELOOP: 12 | try_count += 1 13 | 14 | before_edit_str = edited_str 15 | pos_prev = 0 16 | while True: 17 | pos_r = edited_str.find(">", pos_prev + 1) 18 | if pos_r == -1: 19 | break 20 | 21 | pos_l = edited_str.rfind("<", pos_prev, pos_r) 22 | if pos_l != -1: 23 | left = edited_str[0:pos_l] 24 | center = edited_str[pos_l + 1:pos_r] 25 | right = edited_str[pos_r + 1:len(edited_str)] 26 | 27 | center_splited = center.split("|") 28 | center_picked = center_splited[random.randrange( 29 | 0, len(center_splited))] 30 | 31 | result_left = left + center_picked 32 | pos_prev = len(result_left) 33 | edited_str = result_left + right 34 | else: 35 | pos_prev = pos_r 36 | 37 | if before_edit_str == edited_str: 38 | break 39 | 40 | return edited_str 41 | 42 | 43 | def apply_wc_and_lessthan(wcapplier, prompt): 44 | for x in range(MAX_COUNT_FOR_WHILELOOP): 45 | before_prompt = prompt 46 | prompt = pickedit_lessthan_str(prompt) # lessthan pick 47 | try: 48 | prompt = wcapplier.apply_wildcards(prompt) 49 | except Exception: 50 | pass 51 | 52 | if before_prompt == prompt: 53 | break 54 | 55 | return prompt 56 | 57 | 58 | def inject_imagetag(original_str, tagname, additional_str): 59 | result_str = original_str[:] 60 | 61 | tag_str_left = "@@" + tagname 62 | left_pos = original_str.find(tag_str_left) 63 | if left_pos != -1: 64 | right_pos = original_str.find("@@", left_pos + 1) 65 | except_tag_list = [x.strip() for x in original_str[left_pos + 66 | len(tag_str_left) + 1:right_pos].split(",")] 67 | original_tag_list = [x.strip() for x in additional_str.split(',')] 68 | target_tag_list = [ 69 | x for x in original_tag_list if x not in except_tag_list] 70 | 71 | result_str = original_str[0:left_pos] + ", ".join(target_tag_list) + \ 72 | original_str[right_pos + 2:len(original_str)] 73 | 74 | return result_str 75 | 76 | 77 | def prettify_naidict(nai_dict): 78 | nai_dict = nai_dict.copy() 79 | 80 | desired_order = [ 81 | "prompt", 82 | "negative_prompt", 83 | "width", 84 | "height", 85 | "scale", 86 | "steps", 87 | "sampler", 88 | "cfg_scale", 89 | "sm", 90 | "sm_dyn", 91 | "seed", 92 | "strength", 93 | "noise", 94 | "reference_information_extracted_multiple", 95 | "reference_strength_multiple", 96 | "v4_prompt", 97 | "v4_negative_prompt" 98 | ] 99 | 100 | ban_keys = [ 101 | "image", 102 | "reference_image", 103 | "reference_image_multiple" 104 | ] 105 | 106 | # 먼저 순서대로 정렬할 딕셔너리 만들기 107 | ordered_dict = OrderedDict() 108 | for key in desired_order: 109 | if key in nai_dict: 110 | ordered_dict[key] = nai_dict[key] 111 | 112 | # 남은 키들도 자동으로 추가해주기 113 | for key in nai_dict: 114 | if key not in ordered_dict and key: 115 | ordered_dict[key] = nai_dict[key] 116 | 117 | # ban_keys 작업 118 | for key in ordered_dict: 119 | if key in ban_keys: 120 | ordered_dict[key] = "True" 121 | 122 | 123 | content = json.dumps(ordered_dict, indent=4) 124 | content = content.replace("\\n", "\n") 125 | 126 | return content -------------------------------------------------------------------------------- /app/gui/layout/generate_buttons_layout.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtWidgets import QLabel, QVBoxLayout, QHBoxLayout, QPushButton 2 | 3 | from util.ui_util import add_button, create_empty 4 | 5 | from config.themes import COLOR 6 | 7 | class CustomQLabel(QLabel): 8 | def set_logged_in(self, is_loggedin): 9 | if is_loggedin: 10 | super(CustomQLabel, self).setText("") 11 | self.setStyleSheet(""" 12 | QLabel { 13 | color: white; 14 | font-size: 15px; 15 | } 16 | """) 17 | else: 18 | super(CustomQLabel, self).setText("로그인 안됨") 19 | self.setStyleSheet(""" 20 | QLabel { 21 | color: red; 22 | font-size: 15px; 23 | } 24 | """) 25 | 26 | 27 | class GenerateButtonsLayout(): 28 | def __init__(self, parent): 29 | self.parent = parent 30 | self.parent.generate_buttons_layout = self 31 | 32 | def init(self): 33 | parent = self.parent 34 | 35 | main_layout = QVBoxLayout() 36 | hbox_generate = QHBoxLayout() 37 | main_layout.addLayout(hbox_generate) 38 | 39 | parent.label_loginstate = CustomQLabel() 40 | parent.label_loginstate.set_logged_in(False) 41 | hbox_generate.addWidget(parent.label_loginstate) 42 | 43 | self.button_generate_once = add_button( 44 | hbox_generate, "생성", parent.on_click_generate_once, 200, 200) 45 | self.button_generate_auto = add_button( 46 | hbox_generate, "연속 생성", parent.on_click_generate_auto, 200, 200) 47 | self.button_generate_sett = add_button( 48 | hbox_generate, "연속 세팅 생성", parent.on_click_generate_sett, 200, 200) 49 | self.button_generate_once.setDisabled(True) 50 | self.button_generate_auto.setDisabled(True) 51 | self.button_generate_sett.setDisabled(True) 52 | 53 | main_layout.addWidget(create_empty(minimum_height=5)) 54 | 55 | hbox_expand = QHBoxLayout() 56 | main_layout.addLayout(hbox_expand) 57 | 58 | hbox_anlas = QHBoxLayout() 59 | label_anlas = QLabel("Anlas: ") 60 | label_anlas.setStyleSheet(""" 61 | font: bold; 62 | color: """ + COLOR.BUTTON + """; 63 | """) 64 | parent.label_anlas = label_anlas 65 | hbox_anlas.addWidget(label_anlas) 66 | hbox_expand.addLayout(hbox_anlas, stretch=1) 67 | 68 | hbox_expand.addStretch(9) 69 | 70 | stylesheet_button = """ 71 | color:white; 72 | text-align: center; 73 | font-weight: bold; 74 | background-color:#00000000; 75 | border: 1px solid white; 76 | padding: 8px; 77 | """ 78 | button_getter = QPushButton("Info") 79 | hbox_expand.addWidget(button_getter) 80 | button_getter.clicked.connect(parent.on_click_getter) 81 | button_getter.setStyleSheet(stylesheet_button) 82 | 83 | button_tagger = QPushButton("Tag") 84 | hbox_expand.addWidget(button_tagger) 85 | button_tagger.clicked.connect(parent.on_click_tagger) 86 | button_tagger.setStyleSheet(stylesheet_button) 87 | 88 | button_expand = QPushButton("<<") 89 | hbox_expand.addWidget(button_expand) 90 | button_expand.clicked.connect(parent.on_click_expand) 91 | button_expand.setStyleSheet(stylesheet_button) 92 | parent.button_expand = button_expand 93 | 94 | return main_layout 95 | 96 | def set_disable_button(self, will_disable): 97 | self.button_generate_once.setDisabled(will_disable) 98 | self.button_generate_sett.setDisabled(will_disable) 99 | self.button_generate_auto.setDisabled(will_disable) 100 | 101 | def set_autogenerate_mode(self, is_autogenrate): 102 | self.button_generate_once.setDisabled(is_autogenrate) 103 | self.button_generate_sett.setDisabled(is_autogenrate) 104 | 105 | stylesheet = """ 106 | color:black; 107 | background-color: """ + COLOR.BUTTON_SELECTED + """; 108 | """ if is_autogenrate else "" 109 | self.button_generate_auto.setStyleSheet(stylesheet) 110 | self.button_generate_auto.setText( 111 | "생성 중지" if is_autogenrate else "연속 생성") 112 | self.button_generate_auto.setDisabled(False) -------------------------------------------------------------------------------- /app/gui/layout/main_layout.py: -------------------------------------------------------------------------------- 1 | 2 | from PyQt5.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QGroupBox, QSplitter, QLabel 3 | 4 | from gui.layout.model_options_layout import init_model_options_layout 5 | from gui.layout.resolution_options_layout import init_resolution_options_layout 6 | from gui.layout.parameter_options_layout import init_parameter_options_layout 7 | from gui.layout.generate_buttons_layout import GenerateButtonsLayout 8 | from gui.layout.prompt_layout import PromptLayout 9 | from gui.layout.image_options_layout import init_image_options_layout 10 | from gui.layout.expand_layout import init_expand_layout 11 | 12 | from config.themes import MAIN_STYLESHEET 13 | 14 | from util.ui_util import create_empty, add_button 15 | 16 | 17 | def init_main_layout(self): 18 | widget = QWidget() 19 | widget.setStyleSheet(MAIN_STYLESHEET) 20 | 21 | hbox_headwidget = QHBoxLayout() 22 | widget.setLayout(hbox_headwidget) 23 | 24 | main_splitter = QSplitter() 25 | main_splitter.setHandleWidth(0) 26 | main_splitter.setSizes([9999, 1]) 27 | self.main_splitter = main_splitter 28 | hbox_headwidget.addWidget(main_splitter) 29 | 30 | # left - input 31 | self.dict_ui_settings = {} 32 | 33 | widget_left = QWidget() 34 | main_splitter.addWidget(widget_left) 35 | 36 | hbox_left = QVBoxLayout() 37 | hbox_left.setContentsMargins(30, 20, 30, 30) 38 | widget_left.setLayout(hbox_left) 39 | 40 | hbox_upper_buttons = QHBoxLayout() 41 | hbox_left.addLayout(hbox_upper_buttons) 42 | 43 | hbox_upper_buttons.addLayout(init_setting_buttonlayout(self), stretch=1) 44 | hbox_upper_buttons.addStretch(9999) 45 | hbox_upper_buttons.addWidget(init_openfolder_group(self), stretch=1) 46 | 47 | vbox_lower_settings = QHBoxLayout() 48 | hbox_left.addLayout(vbox_lower_settings) 49 | 50 | vbox_option = QVBoxLayout() 51 | vbox_lower_settings.addLayout(vbox_option, stretch=1) 52 | 53 | vbox_option.addLayout(init_model_options_layout(self), stretch=1) 54 | vbox_option.addLayout(init_resolution_options_layout(self), stretch=1) 55 | vbox_option.addLayout(init_parameter_options_layout(self), stretch=1) 56 | vbox_option.addLayout(init_image_options_layout(self), stretch=999) 57 | vbox_option.addWidget(QLabel("* 모델이 지원하지 않는 패러미터는 무시됩니다.")) 58 | 59 | vbox_lower_settings.addWidget(create_empty(minimum_width=5)) 60 | 61 | vbox_input = QVBoxLayout() 62 | vbox_lower_settings.addLayout(vbox_input, stretch=2000) 63 | 64 | vbox_input.addLayout(PromptLayout(self).init(), stretch=250) 65 | vbox_input.addWidget(create_empty(minimum_height=15), stretch=5) 66 | 67 | vbox_input.addLayout(GenerateButtonsLayout(self).init(), stretch=10) 68 | self.vbox_input = vbox_input 69 | 70 | ############################################# 71 | # right - expand 72 | 73 | widget_right = QWidget() 74 | main_splitter.addWidget(widget_right) 75 | vbox_expand = init_expand_layout(self) 76 | widget_right.setLayout(vbox_expand) 77 | 78 | # main_splitter.setCollapsible(0, False) 79 | main_splitter.widget(1).setMaximumSize(0, 0) 80 | 81 | def on_splitter_move(pos, index): 82 | if self.is_expand: 83 | self.image_result.refresh_size() 84 | if pos > self.size().width() * 0.9: 85 | self.on_click_expand() 86 | self.settings.setValue("splitterSizes", None) 87 | 88 | main_splitter.splitterMoved.connect(on_splitter_move) 89 | 90 | self.setCentralWidget(widget) 91 | 92 | 93 | def init_setting_buttonlayout(self,): 94 | hbox_upper_buttons = QHBoxLayout() 95 | 96 | add_button(hbox_upper_buttons, "세팅 파일로 저장", self.on_click_save_settings) 97 | add_button(hbox_upper_buttons, "세팅 파일 불러오기", self.on_click_load_settings) 98 | 99 | return hbox_upper_buttons 100 | 101 | 102 | def init_openfolder_group(self,): 103 | openfolder_group = QGroupBox("폴더 열기") 104 | 105 | buttons_layout = QHBoxLayout() 106 | openfolder_group.setLayout(buttons_layout) 107 | 108 | add_button(buttons_layout, "생성 결과", 109 | lambda: self.on_click_open_folder("path_results")) 110 | add_button(buttons_layout, "와일드 카드", 111 | lambda: self.on_click_open_folder("path_wildcards")) 112 | add_button(buttons_layout, "세팅 파일", 113 | lambda: self.on_click_open_folder("path_settings")) 114 | 115 | return openfolder_group 116 | -------------------------------------------------------------------------------- /app/gui/layout/parameter_options_layout.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtWidgets import QGroupBox, QLabel, QLineEdit, QCheckBox, QVBoxLayout, QHBoxLayout, QComboBox 2 | 3 | from gui.widget.custom_slider_widget import CustomSliderWidget 4 | 5 | from core.worker.nai_generator import SAMPLER_ITEMS_V4 6 | 7 | 8 | def init_parameter_options_layout(self): 9 | layout = QVBoxLayout() 10 | 11 | # AI Settings Group 12 | ai_settings_group = QGroupBox("AI Settings") 13 | layout.addWidget(ai_settings_group) 14 | 15 | ai_settings_layout = QVBoxLayout() 16 | ai_settings_group.setLayout(ai_settings_layout) 17 | 18 | # Steps Slider 19 | steps_layout = CustomSliderWidget( 20 | title="Steps: ", 21 | min_value=1, 22 | max_value=50, 23 | default_value=28, 24 | edit_width=35, 25 | mag=1, 26 | slider_text_lambda=lambda value: "%d" % value 27 | ) 28 | ai_settings_layout.addLayout(steps_layout, stretch=1) 29 | 30 | lower_layout = QHBoxLayout() 31 | ai_settings_layout.addLayout(lower_layout, stretch=1) 32 | 33 | lower_slider_layout = QVBoxLayout() 34 | lower_checkbox_layout = QVBoxLayout() 35 | lower_layout.addLayout(lower_slider_layout, stretch=2) 36 | lower_layout.addLayout(lower_checkbox_layout, stretch=1) 37 | 38 | # Seed and Sampler 39 | seed_layout = QVBoxLayout() 40 | seed_label_layout = QHBoxLayout() 41 | seed_label_layout.addWidget(QLabel("시드(Seed)")) 42 | seed_label_layout.addStretch(1) 43 | seed_layout.addLayout(seed_label_layout) 44 | seed_input = QLineEdit() 45 | seed_input.setPlaceholderText("여기에 시드 입력") 46 | seed_layout.addWidget(seed_input) 47 | sampler_layout = QVBoxLayout() 48 | sampler_layout.addWidget(QLabel("샘플러(Sampler)")) 49 | sampler_combo = QComboBox() 50 | sampler_combo.addItems(SAMPLER_ITEMS_V4) 51 | self.dict_ui_settings["sampler"] = sampler_combo 52 | sampler_layout.addWidget(sampler_combo) 53 | lower_slider_layout.addLayout(seed_layout) 54 | lower_slider_layout.addLayout(sampler_layout) 55 | 56 | # SMEA Checkbox 57 | seed_opt_layout = QHBoxLayout() 58 | seed_opt_layout.setContentsMargins(10, 25, 0, 0) 59 | seed_fix_checkbox = QCheckBox("고정") 60 | seed_opt_layout.addWidget(seed_fix_checkbox) 61 | variety_plus_checkbox = QCheckBox("variety+") 62 | seed_opt_layout.addWidget(variety_plus_checkbox) 63 | checkbox_layout = QHBoxLayout() 64 | checkbox_layout.setContentsMargins(10, 30, 0, 0) 65 | smea_checkbox = QCheckBox("SMEA") 66 | dyn_checkbox = QCheckBox("+DYN") 67 | checkbox_layout.addWidget(smea_checkbox, stretch=1) 68 | checkbox_layout.addWidget(dyn_checkbox, stretch=1) 69 | lower_checkbox_layout.addLayout(seed_opt_layout) 70 | lower_checkbox_layout.addLayout(checkbox_layout) 71 | 72 | ######## 73 | 74 | # Advanced Settings 75 | advanced_settings_group = QGroupBox("Advanced Settings") 76 | advanced_settings_layout = QVBoxLayout() 77 | 78 | # Prompt Guidance Slider 79 | 80 | prompt_guidance_layout = CustomSliderWidget( 81 | title="Prompt Guidance(CFG):", 82 | min_value=0, 83 | max_value=100, 84 | default_value=5.0, 85 | edit_width=40, 86 | mag=10, 87 | slider_text_lambda=lambda value: "%.1f" % (value / 10) 88 | ) 89 | advanced_settings_layout.addLayout(prompt_guidance_layout) 90 | 91 | # Prompt Guidance Rescale 92 | prompt_rescale_layout = CustomSliderWidget( 93 | title="Prompt Guidance Rescale: ", 94 | min_value=0, 95 | max_value=100, 96 | default_value="0.00", 97 | edit_width=50, 98 | mag=100, 99 | slider_text_lambda=lambda value: "%.2f" % (value / 100) 100 | ) 101 | advanced_settings_layout.addLayout(prompt_rescale_layout) 102 | 103 | advanced_settings_group.setLayout(advanced_settings_layout) 104 | layout.addWidget(advanced_settings_group) 105 | 106 | self.dict_ui_settings["sampler"] = sampler_combo 107 | self.dict_ui_settings["steps"] = steps_layout.edit 108 | self.dict_ui_settings["seed"] = seed_input 109 | self.dict_ui_settings["seed_fix_checkbox"] = seed_fix_checkbox 110 | self.dict_ui_settings["scale"] = prompt_guidance_layout.edit 111 | self.dict_ui_settings["cfg_rescale"] = prompt_rescale_layout.edit 112 | self.dict_ui_settings["variety_plus"] = variety_plus_checkbox 113 | self.dict_ui_settings["sm"] = smea_checkbox 114 | self.dict_ui_settings["sm_dyn"] = dyn_checkbox 115 | 116 | return layout 117 | -------------------------------------------------------------------------------- /app/gui/dialog/login_dialog.py: -------------------------------------------------------------------------------- 1 | 2 | from PyQt5.QtWidgets import QLabel, QLineEdit, QCheckBox, QGridLayout, QVBoxLayout, QHBoxLayout, QPushButton, QDialog 3 | from PyQt5.QtCore import Qt 4 | 5 | from core.thread.login_thread import LoginThread 6 | 7 | class LoginDialog(QDialog): 8 | def __init__(self, parent): 9 | super().__init__() 10 | self.parent = parent 11 | self.initUI() 12 | self.check_already_login() 13 | 14 | def initUI(self): 15 | parent_pos = self.parent.pos() 16 | 17 | self.setWindowTitle('로그인') 18 | self.move(parent_pos.x() + 50, parent_pos.y() + 50) 19 | self.resize(400, 150) 20 | 21 | layout = QVBoxLayout() 22 | 23 | hbox_username = QHBoxLayout() 24 | username_label = QLabel('Username:') 25 | username_edit = QLineEdit(self) 26 | self.username_edit = username_edit 27 | hbox_username.addWidget(username_label) 28 | hbox_username.addWidget(username_edit) 29 | 30 | hbox_password = QHBoxLayout() 31 | password_label = QLabel('Password:') 32 | password_edit = QLineEdit(self) 33 | password_edit.setEchoMode(QLineEdit.Password) # 비밀번호 입력시 마스킹 처리 34 | self.password_edit = password_edit 35 | hbox_password.addWidget(password_label) 36 | hbox_password.addWidget(password_edit) 37 | 38 | vbox_button = QGridLayout() 39 | autologin_check = QCheckBox("자동 로그인") 40 | self.autologin_check = autologin_check 41 | autologin_check.setChecked(True) 42 | login_button = QPushButton('Login') 43 | login_button.clicked.connect(self.try_login) 44 | self.login_button = login_button 45 | logout_button = QPushButton('Logout') 46 | logout_button.setDisabled(True) 47 | logout_button.clicked.connect(self.logout) 48 | self.logout_button = logout_button 49 | vbox_button.addWidget(autologin_check, 0, 0) 50 | vbox_button.addWidget(login_button, 0, 1) 51 | vbox_button.addWidget(logout_button, 1, 1) 52 | 53 | instruct_label = QLabel('Novel AI 계정을 입력해주세요.') 54 | instruct_label.setAlignment(Qt.AlignRight) 55 | self.instruct_label = instruct_label 56 | 57 | layout.addLayout(hbox_username) 58 | layout.addLayout(hbox_password) 59 | layout.addLayout(vbox_button) 60 | layout.addWidget(instruct_label) 61 | 62 | self.setLayout(layout) 63 | 64 | def set_login_result_ui(self, is_login_success): 65 | self.username_edit.setDisabled(is_login_success) 66 | self.password_edit.setDisabled(is_login_success) 67 | 68 | self.login_button.setDisabled(is_login_success) 69 | self.logout_button.setDisabled(not is_login_success) 70 | 71 | self.instruct_label.setText( 72 | "로그인 성공! 창을 닫아도 됩니다." if is_login_success else 'Novel AI 계정을 입력해주세요.') 73 | 74 | def check_already_login(self): 75 | if self.parent.trying_auto_login: 76 | self.username_edit.setDisabled(True) 77 | self.password_edit.setDisabled(True) 78 | 79 | self.login_button.setDisabled(True) 80 | self.logout_button.setDisabled(True) 81 | 82 | self.instruct_label.setText("자동로그인 시도중입니다. 창을 꺼주세요.") 83 | else: # 오토 로그인 중이 아니라면 자동 로그인 84 | nai = self.parent.nai 85 | if nai.access_token: # 자동로그인 성공 86 | self.username_edit.setText(nai.username) 87 | self.password_edit.setText(nai.password) 88 | 89 | self.set_login_result_ui(True) 90 | else: # 자동로그인 실패 91 | self.set_login_result_ui(False) 92 | 93 | def try_login(self): 94 | username = self.username_edit.text() 95 | password = self.password_edit.text() 96 | 97 | self.instruct_label.setText("로그인 시도 중... 창을 닫지 마세요.") 98 | self.parent.statusbar.set_statusbar_text("LOGGINGIN") 99 | self.login_button.setDisabled(True) 100 | self.logout_button.setDisabled(True) 101 | 102 | login_thread = LoginThread(self, self.parent.nai, username, password) 103 | login_thread.login_result.connect(self.on_login_result) 104 | login_thread.login_result.connect(self.parent.on_login_result) 105 | login_thread.start() 106 | self.login_thread = login_thread 107 | 108 | def logout(self): 109 | self.parent.on_logout() 110 | self.set_login_result_ui(False) 111 | 112 | def on_login_result(self, error_code): 113 | if error_code == 0: 114 | self.set_login_result_ui(True) 115 | 116 | if self.autologin_check.isChecked(): 117 | self.parent.set_auto_login(True) 118 | elif error_code == 1: 119 | self.set_login_result_ui(False) 120 | self.instruct_label.setText("잘못된 아이디 또는 비번입니다.") 121 | elif error_code == 2: 122 | self.set_login_result_ui(False) 123 | self.instruct_label.setText("로그인에 실패했습니다.") 124 | -------------------------------------------------------------------------------- /app/core/worker/stealth_pnginfo.py: -------------------------------------------------------------------------------- 1 | from PIL import Image 2 | import gzip 3 | 4 | 5 | # from https://github.com/neggles/sd-webui-stealth-pnginfo/ 6 | def read_info_from_image_stealth(image): 7 | # trying to read stealth pnginfo 8 | width, height = image.size 9 | pixels = image.load() 10 | 11 | has_alpha = True if image.mode == 'RGBA' else False 12 | mode = None 13 | compressed = False 14 | binary_data = '' 15 | buffer_a = '' 16 | buffer_rgb = '' 17 | index_a = 0 18 | index_rgb = 0 19 | sig_confirmed = False 20 | confirming_signature = True 21 | reading_param_len = False 22 | reading_param = False 23 | read_end = False 24 | for x in range(width): 25 | for y in range(height): 26 | if has_alpha: 27 | r, g, b, a = pixels[x, y] 28 | buffer_a += str(a & 1) 29 | index_a += 1 30 | else: 31 | r, g, b = pixels[x, y] 32 | buffer_rgb += str(r & 1) 33 | buffer_rgb += str(g & 1) 34 | buffer_rgb += str(b & 1) 35 | index_rgb += 3 36 | if confirming_signature: 37 | if index_a == len('stealth_pnginfo') * 8: 38 | decoded_sig = bytearray(int(buffer_a[i:i + 8], 2) for i in 39 | range(0, len(buffer_a), 8)).decode('utf-8', errors='ignore') 40 | if decoded_sig in {'stealth_pnginfo', 'stealth_pngcomp'}: 41 | confirming_signature = False 42 | sig_confirmed = True 43 | reading_param_len = True 44 | mode = 'alpha' 45 | if decoded_sig == 'stealth_pngcomp': 46 | compressed = True 47 | buffer_a = '' 48 | index_a = 0 49 | else: 50 | read_end = True 51 | break 52 | elif index_rgb == len('stealth_pnginfo') * 8: 53 | decoded_sig = bytearray(int(buffer_rgb[i:i + 8], 2) for i in 54 | range(0, len(buffer_rgb), 8)).decode('utf-8', errors='ignore') 55 | if decoded_sig in {'stealth_rgbinfo', 'stealth_rgbcomp'}: 56 | confirming_signature = False 57 | sig_confirmed = True 58 | reading_param_len = True 59 | mode = 'rgb' 60 | if decoded_sig == 'stealth_rgbcomp': 61 | compressed = True 62 | buffer_rgb = '' 63 | index_rgb = 0 64 | elif reading_param_len: 65 | if mode == 'alpha': 66 | if index_a == 32: 67 | param_len = int(buffer_a, 2) 68 | reading_param_len = False 69 | reading_param = True 70 | buffer_a = '' 71 | index_a = 0 72 | else: 73 | if index_rgb == 33: 74 | pop = buffer_rgb[-1] 75 | buffer_rgb = buffer_rgb[:-1] 76 | param_len = int(buffer_rgb, 2) 77 | reading_param_len = False 78 | reading_param = True 79 | buffer_rgb = pop 80 | index_rgb = 1 81 | elif reading_param: 82 | if mode == 'alpha': 83 | if index_a == param_len: 84 | binary_data = buffer_a 85 | read_end = True 86 | break 87 | else: 88 | if index_rgb >= param_len: 89 | diff = param_len - index_rgb 90 | if diff < 0: 91 | buffer_rgb = buffer_rgb[:diff] 92 | binary_data = buffer_rgb 93 | read_end = True 94 | break 95 | else: 96 | # impossible 97 | read_end = True 98 | break 99 | if read_end: 100 | break 101 | if sig_confirmed and binary_data != '': 102 | # Convert binary string to UTF-8 encoded text 103 | byte_data = bytearray(int(binary_data[i:i + 8], 2) 104 | for i in range(0, len(binary_data), 8)) 105 | try: 106 | if compressed: 107 | decoded_data = gzip.decompress( 108 | bytes(byte_data)).decode('utf-8') 109 | else: 110 | decoded_data = byte_data.decode('utf-8', errors='ignore') 111 | return decoded_data 112 | except Exception as e: 113 | print(e) 114 | pass 115 | 116 | return None 117 | 118 | 119 | if __name__ == "__main__": 120 | im = Image.open("target.png") 121 | im.load() 122 | gi = read_info_from_image_stealth(im) 123 | print(type(gi)) 124 | -------------------------------------------------------------------------------- /app/core/worker/danbooru_tagger.py: -------------------------------------------------------------------------------- 1 | import os 2 | import io 3 | import base64 4 | import onnxruntime as ort 5 | import csv 6 | from PIL import Image 7 | import numpy as np 8 | import requests 9 | import shutil 10 | 11 | from util.file_util import create_folder_if_not_exists 12 | 13 | from config.paths import PATH_IMG_NO_IMAGE 14 | from config.consts import DEFAULT_TAGGER_MODEL, LIST_TAGGER_MODEL 15 | 16 | 17 | def download_file(url, dst): 18 | try: 19 | with requests.get(url, stream=True) as response: 20 | if response.status_code == 200: 21 | with open(dst, 'wb') as f: 22 | response.raw.decode_content = True 23 | shutil.copyfileobj(response.raw, f) 24 | return True 25 | else: 26 | print( 27 | f"Failed to download file from {url}. Status code: {response.status_code}") 28 | return False 29 | except Exception as e: 30 | print(f"An error occurred: {e}") 31 | return False 32 | 33 | 34 | class DanbooruTagger(): 35 | def __init__(self, models_dir): 36 | self.models_dir = models_dir 37 | self.options = { 38 | "model_name": DEFAULT_TAGGER_MODEL, 39 | "threshold": 0.35, 40 | "character_threshold": 0.85, 41 | "replace_underscore": True, 42 | "trailing_comma": False, 43 | "exclude_tags": "" 44 | } 45 | 46 | def get_installed_models(self): 47 | create_folder_if_not_exists(self.models_dir) 48 | return list(filter(lambda x: x.endswith(".onnx"), os.listdir(self.models_dir))) 49 | 50 | def tag(self, image): 51 | model_name = self.options['model_name'] 52 | threshold = self.options['threshold'] 53 | character_threshold = self.options['character_threshold'] 54 | replace_underscore = self.options['replace_underscore'] 55 | trailing_comma = self.options['trailing_comma'] 56 | exclude_tags = self.options['exclude_tags'] 57 | 58 | if model_name.endswith(".onnx"): 59 | model_name = model_name[0:-5] 60 | installed = self.get_installed_models() 61 | if not any(model_name + ".onnx" in s for s in installed): 62 | print("model not installed") 63 | return 64 | 65 | name = os.path.join(self.models_dir, model_name + ".onnx") 66 | model = ort.InferenceSession( 67 | name, providers=ort.get_available_providers()) 68 | 69 | input = model.get_inputs()[0] 70 | height = input.shape[1] 71 | 72 | # Reduce to max size and pad with white 73 | ratio = float(height) / max(image.size) 74 | new_size = tuple([int(x * ratio) for x in image.size]) 75 | image = image.resize(new_size, Image.LANCZOS) 76 | square = Image.new("RGB", (height, height), (255, 255, 255)) 77 | square.paste( 78 | image, ((height - new_size[0]) // 2, (height - new_size[1]) // 2)) 79 | 80 | image = np.array(square).astype(np.float32) 81 | image = image[:, :, ::-1] # RGB -> BGR 82 | image = np.expand_dims(image, 0) 83 | 84 | # Read all tags from csv and locate start of each category 85 | tags = [] 86 | general_index = None 87 | character_index = None 88 | with open(os.path.join(self.models_dir, model_name + ".csv")) as f: 89 | reader = csv.reader(f) 90 | next(reader) 91 | for row in reader: 92 | if general_index is None and row[2] == "0": 93 | general_index = reader.line_num - 2 94 | elif character_index is None and row[2] == "4": 95 | character_index = reader.line_num - 2 96 | if replace_underscore: 97 | tags.append(row[1].replace("_", " ")) 98 | else: 99 | tags.append(row[1]) 100 | 101 | label_name = model.get_outputs()[0].name 102 | probs = model.run([label_name], {input.name: image})[0] 103 | 104 | result = list(zip(tags, probs[0])) 105 | 106 | # rating = max(result[:general_index], key=lambda x: x[1]) 107 | general = [item for item in result[general_index:character_index] 108 | if item[1] > threshold] 109 | character = [item for item in result[character_index:] 110 | if item[1] > character_threshold] 111 | 112 | all = character + general 113 | remove = [s.strip() for s in exclude_tags.lower().split(",")] 114 | all = [tag for tag in all if tag[0] not in remove] 115 | 116 | res = ("" if trailing_comma else ", ").join((item[0].replace( 117 | "(", "\\(").replace(")", "\\)") + (", " if trailing_comma else "") for item in all)) 118 | 119 | return res 120 | 121 | def download_model(self, model): 122 | installed = self.get_installed_models() 123 | if any(model + ".onnx" in s for s in installed): 124 | print("model already installed") 125 | return True 126 | 127 | url = f"https://huggingface.co/SmilingWolf/{model}/resolve/main/" 128 | is_success = download_file( 129 | f"{url}model.onnx", 130 | os.path.join(self.models_dir, f"{model}.onnx")) 131 | is_success = is_success and download_file( 132 | f"{url}selected_tags.csv", 133 | os.path.join(self.models_dir, f"{model}.csv")) 134 | 135 | return is_success 136 | -------------------------------------------------------------------------------- /app/core/thread/generate_thread.py: -------------------------------------------------------------------------------- 1 | import os 2 | import time 3 | import zipfile 4 | import io 5 | import random 6 | import datetime 7 | from PIL import Image 8 | 9 | from PyQt5.QtCore import QThread, pyqtSignal 10 | 11 | from util.common_util import strtobool 12 | from util.file_util import create_folder_if_not_exists, get_filename_only, create_windows_filepath 13 | 14 | from config.paths import DEFAULT_PATH 15 | 16 | 17 | def _threadfunc_generate_image(thread_self, path): 18 | parent = thread_self.parent() 19 | nai = parent.nai 20 | data = nai.generate_image() 21 | if not data: 22 | return 1, "서버에서 정보를 가져오는데 실패했습니다." 23 | 24 | try: 25 | zipped = zipfile.ZipFile(io.BytesIO(data)) 26 | image_bytes = zipped.read(zipped.infolist()[0]) 27 | img = Image.open(io.BytesIO(image_bytes)) 28 | except Exception as e: 29 | return 2, str(e) + str(data) 30 | 31 | create_folder_if_not_exists(path) 32 | dst = "" 33 | timename = datetime.datetime.now().strftime("%y%m%d_%H%M%S%f")[:-4] 34 | filename = timename 35 | if strtobool(parent.settings.value("will_savename_prompt", True)): 36 | filename += "_" + nai.parameters["prompt"] 37 | dst = create_windows_filepath(path, filename, ".png") 38 | if not dst: 39 | dst = timename 40 | 41 | try: 42 | img.save(dst) 43 | except Exception as e: 44 | return 3, str(e) 45 | 46 | return 0, dst 47 | 48 | 49 | class GenerateThread(QThread): 50 | # Signals for 자동 생성 모드 51 | on_data_created = pyqtSignal() 52 | on_error = pyqtSignal(int, str) 53 | on_success = pyqtSignal(str) 54 | on_end = pyqtSignal() 55 | on_statusbar_change = pyqtSignal(str, list) 56 | 57 | # Signal for 단일 생성 모드 58 | generate_result = pyqtSignal(int, str) 59 | 60 | def __init__(self, parent, auto_mode=False, count=-1, delay=0.01, ignore_error=False): 61 | super(GenerateThread, self).__init__(parent) 62 | self.auto_mode = auto_mode 63 | self.count = int(count or -1) 64 | self.delay = float(delay or 0.01) 65 | self.ignore_error = ignore_error 66 | self.is_dead = False 67 | 68 | def run(self): 69 | parent = self.parent() 70 | if self.auto_mode: 71 | # 자동 생성 모드 (기존 AutoGenerateThread의 기능) 72 | count = self.count 73 | delay = self.delay 74 | temp_preserve_data_once = False 75 | while count != 0: 76 | # 1. 데이터 생성 77 | if not temp_preserve_data_once: 78 | data = parent._get_data_for_generate() 79 | parent.nai.set_param_dict(data) 80 | self.on_data_created.emit() 81 | temp_preserve_data_once = False 82 | 83 | # 상태 표시줄 업데이트 84 | if count <= -1: 85 | self.on_statusbar_change.emit("AUTO_GENERATING_INF", []) 86 | else: 87 | self.on_statusbar_change.emit("AUTO_GENERATING_COUNT", [ 88 | self.count, self.count - count + 1]) 89 | 90 | # 결과 저장 경로 설정 (배치 처리 고려) 91 | path = parent.settings.value( 92 | "path_results", DEFAULT_PATH["path_results"]) 93 | create_folder_if_not_exists(path) 94 | if parent.list_settings_batch_target: 95 | setting_path = parent.list_settings_batch_target[parent.index_settings_batch_target] 96 | setting_name = get_filename_only(setting_path) 97 | path = os.path.join(path, setting_name) 98 | create_folder_if_not_exists(path) 99 | 100 | # 이미지 생성 101 | error_code, result_str = _threadfunc_generate_image(self, path) 102 | if self.is_dead: 103 | return 104 | if error_code == 0: 105 | self.on_success.emit(result_str) 106 | else: 107 | if self.ignore_error: 108 | for t in range(int(delay), 0, -1): 109 | self.on_statusbar_change.emit( 110 | "AUTO_ERROR_WAIT", [t]) 111 | time.sleep(1+ random.uniform(0.01, 0.1)) 112 | if self.is_dead: 113 | return 114 | temp_preserve_data_once = True 115 | continue 116 | else: 117 | self.on_error.emit(error_code, result_str) 118 | return 119 | 120 | # 2. 대기 121 | count -= 1 122 | if count != 0: 123 | temp_delay = delay 124 | for _ in range(int(delay)): 125 | self.on_statusbar_change.emit( 126 | "AUTO_WAIT", [temp_delay]) 127 | time.sleep(1+ random.uniform(0.01, 0.1)) 128 | if self.is_dead: 129 | return 130 | temp_delay -= 1 131 | self.on_end.emit() 132 | else: 133 | # 단일 생성 모드 (기존 GenerateThread의 기능) 134 | path = parent.settings.value( 135 | "path_results", DEFAULT_PATH["path_results"]) 136 | error_code, result_str = _threadfunc_generate_image(self, path) 137 | self.generate_result.emit(error_code, result_str) 138 | 139 | def stop(self): 140 | self.is_dead = True 141 | self.quit() 142 | -------------------------------------------------------------------------------- /app/core/worker/completer.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtWidgets import QCompleter, QTextEdit 2 | from PyQt5.QtGui import QTextCursor, QTextCharFormat, QFont, QColor 3 | from PyQt5.QtCore import Qt, QStringListModel 4 | 5 | 6 | # complete_target_stringset = string.ascii_letters + string.digits + "~!#$%^&*_+?.-=" 7 | class CustomCompleter(QCompleter): 8 | def __init__(self, words, parent=None): 9 | super().__init__(words, parent) 10 | self.words = words 11 | self.setCaseSensitivity(Qt.CaseInsensitive) 12 | self.setFilterMode(Qt.MatchContains) 13 | self.model = QStringListModel(words, self) 14 | self.setModel(self.model) 15 | 16 | def setCompletionPrefix(self, prefix): 17 | self.prefix = prefix 18 | 19 | is_add_mode = len(self.prefix) > 3 20 | prefix_lower = self.prefix.lower() 21 | filtered_words = [] 22 | contains_matches = [] 23 | for word in self.words: 24 | word_lower = word.lower() 25 | if word_lower.startswith(prefix_lower): 26 | filtered_words.append(word) 27 | elif is_add_mode and prefix_lower in word_lower: 28 | contains_matches.append(word) 29 | 30 | if is_add_mode: 31 | filtered_words.extend(sorted(contains_matches)) 32 | 33 | self.model.setStringList(filtered_words) 34 | super().setCompletionPrefix(prefix) 35 | self.complete() 36 | 37 | 38 | class CompletionTextEdit(QTextEdit): 39 | def __init__(self): 40 | super().__init__() 41 | self.setAcceptRichText(False) 42 | self.setAcceptDrops(False) 43 | self.completer = None 44 | self.textChanged.connect(self.highlightBrackets) 45 | def highlightBrackets(self): 46 | text = self.toPlainText() 47 | 48 | # Stack to keep track of open brackets 49 | stack = [] 50 | bracket_pairs = {'(': ')', '{': '}', '[': ']', '<': '>'} 51 | open_brackets = bracket_pairs.keys() 52 | close_brackets = bracket_pairs.values() 53 | 54 | # Dictionary to keep track of bracket positions 55 | bracket_positions = {} 56 | for i, char in enumerate(text): 57 | if char in open_brackets: 58 | stack.append((char, i)) 59 | bracket_positions[i] = -1 # 기본 값으로 -1 설정 60 | elif char in close_brackets: 61 | if stack and bracket_pairs[stack[-1][0]] == char: 62 | open_bracket, open_pos = stack.pop() 63 | bracket_positions[open_pos] = i 64 | bracket_positions[i] = open_pos 65 | else: 66 | bracket_positions[i] = -1 67 | 68 | # Highlight unmatched brackets 69 | extraSelections = [] 70 | unmatched_format = QTextCharFormat() 71 | unmatched_format.setFontWeight(QFont.Bold) 72 | unmatched_format.setForeground(QColor("red")) 73 | 74 | for pos, matching_pos in bracket_positions.items(): 75 | if matching_pos == -1: 76 | selection = QTextEdit.ExtraSelection() 77 | selection.format = unmatched_format 78 | cursor = self.textCursor() 79 | cursor.setPosition(pos) 80 | cursor.movePosition(QTextCursor.NextCharacter, QTextCursor.KeepAnchor) 81 | selection.cursor = cursor 82 | extraSelections.append(selection) 83 | 84 | self.setExtraSelections(extraSelections) 85 | 86 | def start_complete_mode(self, tag_list): 87 | completer = CustomCompleter(tag_list) 88 | 89 | self.setCompleter(completer) 90 | 91 | def setCompleter(self, completer): 92 | if self.completer: 93 | self.disconnect(self.completer, self.insertCompletion) 94 | self.completer = completer 95 | if not self.completer: 96 | return 97 | self.completer.setWidget(self) 98 | self.completer.setCompletionMode(QCompleter.PopupCompletion) 99 | self.completer.setCaseSensitivity(False) 100 | self.completer.activated.connect(self.insertCompletion) 101 | 102 | def insertCompletion(self, completion): 103 | actual_text = completion.split('[')[0] # 'apple[30]' -> 'apple' 104 | actual_text = actual_text.replace("_", " ") 105 | tc = self.textCursor() 106 | extra = len(actual_text) - len(self.completer.completionPrefix()) 107 | tc.movePosition(QTextCursor.Left) 108 | tc.movePosition(QTextCursor.EndOfWord) 109 | tc.insertText(actual_text[-extra:]) 110 | self.setTextCursor(tc) 111 | 112 | def textUnderCursor(self): 113 | tc = self.textCursor() 114 | tc.select(QTextCursor.WordUnderCursor) 115 | return tc.selectedText() 116 | 117 | def keyPressEvent(self, event): 118 | if self.completer: 119 | if self.completer and self.completer.popup().isVisible(): 120 | if event.key() in (Qt.Key_Enter, Qt.Key_Return, Qt.Key_Escape, Qt.Key_Tab, Qt.Key_Backtab): 121 | event.ignore() 122 | return 123 | 124 | super().keyPressEvent(event) 125 | 126 | ctrlOrShift = event.modifiers() & (Qt.ControlModifier | Qt.ShiftModifier) 127 | if ctrlOrShift and event.text() == '': 128 | return 129 | 130 | if event.text(): 131 | # eow = "~!@#$%^&*()_+{}|:\"<>?,./;'[]\\-=" 132 | eow = "{},<>|@" 133 | hasModifier = (event.modifiers() != 134 | Qt.NoModifier) and not ctrlOrShift 135 | completionPrefix = self.textUnderCursor() 136 | 137 | if not self.completer or (hasModifier and event.text() == '') or len(completionPrefix) < 1 or event.text()[-1] in eow: 138 | self.completer.popup().hide() 139 | return 140 | 141 | if completionPrefix != self.completer.completionPrefix(): 142 | self.completer.setCompletionPrefix(completionPrefix) 143 | self.completer.popup().setCurrentIndex( 144 | self.completer.completionModel().index(0, 0)) 145 | 146 | cr = self.cursorRect() 147 | cr.setWidth(self.completer.popup().sizeHintForColumn( 148 | 0) + self.completer.popup().verticalScrollBar().sizeHint().width()) 149 | self.completer.complete(cr) 150 | else: 151 | super().keyPressEvent(event) 152 | -------------------------------------------------------------------------------- /app/gui/layout/resolution_options_layout.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtWidgets import QGroupBox, QLabel, QLineEdit, QCheckBox, QVBoxLayout, QHBoxLayout, QStyledItemDelegate, QComboBox 2 | from PyQt5.QtCore import Qt 3 | from PyQt5.QtGui import QPalette, QColor, QFont, QIntValidator 4 | 5 | from util.common_util import strtobool 6 | 7 | from config.themes import COLOR 8 | from config.consts import RESOLUTION_ITEMS, DEFAULT_RESOLUTION, RESOLUTION_ITEMS_NOT_SELECTABLES 9 | 10 | class CustomDelegate(QStyledItemDelegate): 11 | def __init__(self, parent=None, special_indexes=None): 12 | super(CustomDelegate, self).__init__(parent) 13 | self.special_indexes = special_indexes or set() 14 | self.parent = parent 15 | 16 | def paint(self, painter, option, index): 17 | if index.row() == self.parent.currentIndex(): 18 | painter.save() 19 | 20 | painter.fillRect(option.rect, QColor(COLOR.BUTTON)) 21 | palette = QPalette() 22 | palette.setColor( 23 | QPalette.Text, QColor('black')) # Text color 24 | option.palette = palette 25 | option.font.setBold(True) 26 | super(CustomDelegate, self).paint(painter, option, index) 27 | painter.restore() 28 | palette.setColor( 29 | QPalette.Text, QColor('white')) # Text color 30 | option.palette = palette 31 | option.font.setBold(False) 32 | elif index.row() in self.special_indexes: 33 | # Save the original state of the painter to restore it later 34 | painter.save() 35 | # Apply custom styles 36 | painter.setPen(QColor(COLOR.GRAY)) 37 | painter.fillRect(option.rect, QColor(COLOR.GRAY)) 38 | painter.setFont(QFont('Arial', 10, QFont.Bold)) 39 | painter.restore() 40 | super(CustomDelegate, self).paint(painter, option, index) 41 | else: 42 | super(CustomDelegate, self).paint(painter, option, index) 43 | 44 | 45 | class CustomWHLineEdit(QLineEdit): 46 | def __init__(self, size, custom_func): 47 | super(CustomWHLineEdit, self).__init__(size) 48 | self.custom_func = custom_func 49 | 50 | def focusOutEvent(self, event): 51 | self.custom_func() 52 | super(CustomWHLineEdit, self).focusOutEvent(event) 53 | 54 | def setText(self, text): 55 | super(CustomWHLineEdit, self).setText(text) 56 | self.custom_func() 57 | 58 | 59 | def create_custom_whlineedit(custom_func): 60 | whle = CustomWHLineEdit("640", custom_func) 61 | whle.setMinimumWidth(50) 62 | whle.setValidator(QIntValidator(0, 2000)) 63 | whle.setAlignment(Qt.AlignCenter) 64 | whle.returnPressed.connect(custom_func) 65 | 66 | return whle 67 | 68 | 69 | 70 | def init_resolution_options_layout(self): 71 | layout = QVBoxLayout() 72 | 73 | image_settings_group = QGroupBox("Image Settings") 74 | layout.addWidget(image_settings_group) 75 | 76 | image_settings_layout = QHBoxLayout() 77 | image_settings_group.setLayout(image_settings_layout) 78 | 79 | left_layout = QVBoxLayout() 80 | right_layout = QHBoxLayout() 81 | image_settings_layout.addStretch(2) 82 | image_settings_layout.addLayout(left_layout, stretch=10) 83 | image_settings_layout.addStretch(1) 84 | image_settings_layout.addLayout(right_layout, stretch=10) 85 | image_settings_layout.addStretch(2) 86 | 87 | combo_resolution = QComboBox() 88 | combo_resolution.setMinimumWidth(220) 89 | combo_resolution.setMaxVisibleItems(len(RESOLUTION_ITEMS)) 90 | combo_resolution.addItems(RESOLUTION_ITEMS) 91 | combo_resolution.setCurrentIndex( 92 | RESOLUTION_ITEMS.index(self.settings.value( 93 | "resolution", DEFAULT_RESOLUTION)) 94 | ) 95 | self.combo_resolution = combo_resolution 96 | 97 | combo_resolution.setItemDelegate(CustomDelegate( 98 | combo_resolution, special_indexes=RESOLUTION_ITEMS_NOT_SELECTABLES) 99 | ) 100 | 101 | for nselectable in RESOLUTION_ITEMS_NOT_SELECTABLES: 102 | combo_resolution.setItemData( 103 | nselectable, 0, Qt.UserRole - 1) 104 | 105 | def update_resolutions(text): 106 | # Custom이 아닌경우에만 적용 107 | if "(" in text and ")" in text: 108 | res_text = text.split("(")[1].split(")")[0] 109 | width, height = res_text.split("x") 110 | self.width_edit.setText(width) 111 | self.height_edit.setText(height) 112 | combo_resolution.currentTextChanged.connect(update_resolutions) 113 | 114 | left_layout.addWidget(combo_resolution) 115 | 116 | # Right side layout 117 | 118 | def on_enter_or_focusout(): 119 | now_width = self.width_edit.text() 120 | now_height = self.height_edit.text() 121 | 122 | # 일치하는 경우 확인 123 | found = False 124 | if now_width and now_height: 125 | for text in RESOLUTION_ITEMS: 126 | if "(" in text and ")" in text: 127 | res_text = text.split("(")[1].split(")")[0] 128 | width, height = res_text.split("x") 129 | if now_width == width.strip() and now_height == height.strip(): 130 | combo_resolution.setCurrentText(text) 131 | found = True 132 | 133 | # 그 외 모든 경우 맨마지막 - Custom으로 감 134 | if not found: 135 | combo_resolution.setCurrentIndex(combo_resolution.count() - 1) 136 | 137 | self.width_edit = create_custom_whlineedit(on_enter_or_focusout) 138 | self.height_edit = create_custom_whlineedit(on_enter_or_focusout) 139 | 140 | # Add widgets to the right layout 141 | right_layout.addWidget(QLabel("너비(Width)")) 142 | right_layout.addWidget(self.width_edit) 143 | right_layout.addWidget(QLabel("높이(Height)")) 144 | right_layout.addWidget(self.height_edit) 145 | 146 | # Check Box Layout 147 | checkbox_layout = QHBoxLayout() 148 | checkbox_layout.addStretch(2000) 149 | checkbox_random_resolution = QCheckBox("이미지 크기 랜덤") 150 | prev_value_checkbox = strtobool( 151 | self.settings.value("image_random_checkbox", False)) 152 | checkbox_random_resolution.setChecked(prev_value_checkbox) 153 | checkbox_layout.addWidget(checkbox_random_resolution) 154 | checkbox_random_resolution.stateChanged.connect( 155 | self.on_random_resolution_checked) 156 | self.checkbox_random_resolution = checkbox_random_resolution 157 | layout.addLayout(checkbox_layout) 158 | 159 | self.dict_ui_settings["resolution"] = combo_resolution 160 | self.dict_ui_settings["width"] = self.width_edit 161 | self.dict_ui_settings["height"] = self.height_edit 162 | 163 | return layout -------------------------------------------------------------------------------- /app/gui/dialog/option_dialog.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | 4 | from PyQt5.QtWidgets import QApplication, QRadioButton, QSlider, QGroupBox, QFileDialog, QLabel, QCheckBox, QVBoxLayout, QHBoxLayout, QPushButton, QDialog, QSizePolicy 5 | from PyQt5.QtCore import Qt, QSettings 6 | from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QDialog 7 | 8 | from core.worker.danbooru_tagger import DanbooruTagger 9 | 10 | from util.common_util import strtobool, get_key_from_dict 11 | from util.ui_util import create_empty 12 | 13 | from config.paths import DEFAULT_PATH 14 | from config.consts import LIST_TAGGER_MODEL 15 | 16 | class OptionDialog(QDialog): 17 | def __init__(self, parent): 18 | super().__init__(parent) 19 | self.initUI() 20 | 21 | def initUI(self): 22 | parent = self.parent() 23 | parent_pos = parent.pos() 24 | 25 | self.setWindowTitle('옵션') 26 | self.move(parent_pos.x() + 50, parent_pos.y() + 50) 27 | self.resize(600, 200) 28 | layout = QVBoxLayout() 29 | 30 | self.dict_label_loc = {} 31 | 32 | def add_item(layout, code, text): 33 | hbox_item = QHBoxLayout() 34 | layout.addLayout(hbox_item) 35 | 36 | label_title = QLabel(text) 37 | hbox_item.addWidget(label_title) 38 | 39 | path = parent.settings.value( 40 | code, DEFAULT_PATH[code]) 41 | label_loc = QLabel(os.path.abspath(path)) 42 | label_loc.setStyleSheet("font-size: 14px") 43 | self.dict_label_loc[code] = label_loc 44 | hbox_item.addWidget(label_loc, stretch=999) 45 | 46 | button_select_loc = QPushButton("위치 변경") 47 | button_select_loc.setSizePolicy( 48 | QSizePolicy.Minimum, QSizePolicy.Minimum) 49 | button_select_loc.clicked.connect( 50 | lambda: self.on_click_select_button(code)) 51 | hbox_item.addWidget(button_select_loc) 52 | 53 | button_reset_loc = QPushButton("리셋") 54 | button_reset_loc.setSizePolicy( 55 | QSizePolicy.Minimum, QSizePolicy.Minimum) 56 | button_reset_loc.clicked.connect( 57 | lambda: self.on_click_reset_button(code)) 58 | hbox_item.addWidget(button_reset_loc) 59 | 60 | folderloc_group = QGroupBox("폴더 위치") 61 | layout.addWidget(folderloc_group) 62 | 63 | folderloc_group_layout = QVBoxLayout() 64 | folderloc_group.setLayout(folderloc_group_layout) 65 | 66 | add_item(folderloc_group_layout, "path_results", "생성이미지 저장 위치 : ") 67 | add_item(folderloc_group_layout, "path_wildcards", "와일드카드 저장 위치 : ") 68 | add_item(folderloc_group_layout, "path_settings", "세팅 파일 저장 위치 : ") 69 | add_item(folderloc_group_layout, "path_models", "태거 모델 저장 위치 : ") 70 | 71 | layout.addWidget(create_empty(minimum_height=6)) 72 | 73 | groupBox_tagger = QGroupBox("태그 모델 설치 및 선택") 74 | self.groupBox_tagger = groupBox_tagger 75 | layout.addWidget(groupBox_tagger) 76 | 77 | groupBox_tagger_layout = QVBoxLayout() 78 | groupBox_tagger.setLayout(groupBox_tagger_layout) 79 | 80 | list_installed_model = parent.dtagger.get_installed_models() 81 | self.dict_tagger_model_radio = {} 82 | self.dict_tagger_model_button = {} 83 | for index, model_name in enumerate(LIST_TAGGER_MODEL): 84 | radio_layout = QHBoxLayout() 85 | groupBox_tagger_layout.addLayout(radio_layout) 86 | 87 | radio = QRadioButton(model_name + ("(추천)" if index == 0 else "")) 88 | radio.setEnabled(False) 89 | radio.clicked.connect(self.on_click_modelradio_button) 90 | self.dict_tagger_model_radio[model_name] = radio 91 | radio_layout.addWidget(radio) 92 | 93 | download_button = QPushButton("다운로드") 94 | download_button.clicked.connect(self.on_click_download_button) 95 | self.dict_tagger_model_button[model_name] = download_button 96 | radio_layout.addWidget(download_button) 97 | 98 | if model_name + ".onnx" in list_installed_model: 99 | radio.setEnabled(True) 100 | download_button.setEnabled(False) 101 | 102 | now_selected = parent.settings.value("selected_tagger_model", '') 103 | if now_selected and now_selected in list_installed_model: 104 | self.dict_tagger_model_radio[now_selected].setChecked(True) 105 | else: 106 | for k, v in self.dict_tagger_model_radio.items(): 107 | if v.isEnabled(): 108 | v.setChecked(True) 109 | parent.settings.setValue("selected_tagger_model", k) 110 | break 111 | 112 | layout.addWidget(create_empty(minimum_height=6)) 113 | 114 | font_layout = QHBoxLayout() 115 | layout.addLayout(font_layout) 116 | 117 | now_font_size = self.parent().settings.value("nag_font_size", 18) 118 | font_label = QLabel(f'글꼴 크기(종료시 적용, 기본 18): {now_font_size}') 119 | self.font_label = font_label 120 | font_layout.addWidget(font_label) 121 | 122 | font_progress_bar = QSlider(self) 123 | font_progress_bar.setMinimum(14) 124 | font_progress_bar.setMaximum(50) 125 | font_progress_bar.setValue(now_font_size) 126 | font_progress_bar.setOrientation(Qt.Horizontal) 127 | font_progress_bar.valueChanged.connect(self.on_fontlabel_updated) 128 | font_layout.addWidget(font_progress_bar) 129 | 130 | checkbox_completeoff = QCheckBox("태그 자동완성 키기") 131 | checkbox_completeoff.setChecked(strtobool( 132 | parent.settings.value("will_complete_tag", True))) 133 | self.checkbox_completeoff = checkbox_completeoff 134 | layout.addWidget(checkbox_completeoff) 135 | 136 | checkbox_savepname = QCheckBox("파일 생성시 이름에 프롬프트 넣기") 137 | checkbox_savepname.setChecked(strtobool( 138 | parent.settings.value("will_savename_prompt", True))) 139 | self.checkbox_savepname = checkbox_savepname 140 | layout.addWidget(checkbox_savepname) 141 | 142 | button_close = QPushButton("닫기") 143 | button_close.clicked.connect(self.on_click_close_button) 144 | self.button_close = button_close 145 | 146 | layout.addStretch(2) 147 | 148 | qhl_close = QHBoxLayout() 149 | qhl_close.addStretch(4) 150 | qhl_close.addWidget(self.button_close, 2) 151 | layout.addLayout(qhl_close) 152 | 153 | self.setLayout(layout) 154 | 155 | def on_fontlabel_updated(self, value): 156 | self.font_label.setText(f'폰트 사이즈(종료시 적용, 기본 18): {value}') 157 | self.parent().settings.setValue("nag_font_size", int(value)) 158 | 159 | def on_click_select_button(self, code): 160 | select_dialog = QFileDialog() 161 | save_loc = select_dialog.getExistingDirectory( 162 | self, '저장할 위치를 골라주세요.') 163 | 164 | if save_loc: 165 | self.parent().on_change_path(code, save_loc) 166 | 167 | self.refresh_label(code) 168 | 169 | def on_click_download_button(self): 170 | button = self.sender() 171 | model_name = get_key_from_dict(self.dict_tagger_model_button, button) 172 | 173 | self.parent().install_model(model_name) 174 | 175 | def on_click_modelradio_button(self): 176 | radio = self.sender() 177 | model_name = get_key_from_dict(self.dict_tagger_model_radio, radio) 178 | 179 | self.parent().settings.setValue("selected_tagger_model", model_name) 180 | 181 | def on_click_reset_button(self, code): 182 | self.parent().on_change_path(code, DEFAULT_PATH[code]) 183 | 184 | self.refresh_label(code) 185 | 186 | def on_model_downloaded(self, model_name): 187 | self.dict_tagger_model_radio[model_name].setEnabled(True) 188 | self.dict_tagger_model_button[model_name].setEnabled(False) 189 | 190 | def refresh_label(self, code): 191 | path = self.parent().settings.value(code, DEFAULT_PATH[code]) 192 | self.dict_label_loc[code].setText(path) 193 | 194 | def on_click_close_button(self): 195 | self.parent().settings.setValue( 196 | "will_complete_tag", self.checkbox_completeoff.isChecked()) 197 | self.parent().settings.setValue( 198 | "will_savename_prompt", self.checkbox_savepname.isChecked()) 199 | self.reject() 200 | 201 | 202 | if __name__ == '__main__': 203 | app = QApplication(sys.argv) 204 | from PyQt5.QtWidgets import QMainWindow 205 | from PyQt5.QtCore import QSettings 206 | TOP_NAME = "dcp_arca" 207 | APP_NAME = "nag_gui" 208 | qw = QMainWindow() 209 | qw.move(200, 200) 210 | qw.settings = QSettings(TOP_NAME, APP_NAME) 211 | qw.dtagger = DanbooruTagger(qw.settings.value( 212 | "path_models", os.path.abspath(DEFAULT_PATH["path_models"]))) 213 | OptionDialog(qw) -------------------------------------------------------------------------------- /app/gui/dialog/miniutil_dialog.py: -------------------------------------------------------------------------------- 1 | import os 2 | import sys 3 | from io import BytesIO 4 | from PIL import Image 5 | from urllib import request 6 | 7 | from PyQt5.QtWidgets import QApplication, QTextEdit, QFrame, QFileDialog, QLabel, QVBoxLayout, QDialog, QMessageBox 8 | from PyQt5.QtGui import QImage, QPainter 9 | from PyQt5.QtCore import Qt, pyqtSignal, QEvent, QRectF, QSize, pyqtSignal 10 | 11 | from core.worker.danbooru_tagger import DanbooruTagger 12 | 13 | from core.worker.naiinfo_getter import get_naidict_from_file, get_naidict_from_txt, get_naidict_from_img 14 | 15 | from util.image_util import pil2pixmap 16 | from util.tagger_util import predict_tag_from 17 | from util.string_util import prettify_naidict 18 | 19 | from config.paths import DEFAULT_PATH, PATH_IMG_TAGGER, PATH_IMG_GETTER 20 | 21 | class BackgroundFrame(QFrame): 22 | clicked = pyqtSignal() 23 | 24 | def __init__(self, parent, opacity): 25 | super(BackgroundFrame, self).__init__(parent) 26 | self.image = QImage() 27 | self.opacity = opacity 28 | 29 | def set_background_image_by_src(self, image_path): 30 | self.image.load(image_path) 31 | self.update() 32 | 33 | def set_background_image_by_img(self, image): 34 | self.image = pil2pixmap(image).toImage() 35 | self.update() 36 | 37 | def paintEvent(self, event): 38 | painter = QPainter(self) 39 | scaled_image = self.image.scaled( 40 | self.size(), Qt.KeepAspectRatioByExpanding, Qt.SmoothTransformation) 41 | target_rect = QRectF( 42 | (self.width() - scaled_image.width()) / 2, 43 | (self.height() - scaled_image.height()) / 2, 44 | scaled_image.width(), 45 | scaled_image.height()) 46 | painter.setOpacity(self.opacity) 47 | painter.drawImage(target_rect, scaled_image) 48 | 49 | def eventFilter(self, obj, event): 50 | if event.type() == QEvent.Resize: 51 | self.update() 52 | return True 53 | return super(BackgroundFrame, self).eventFilter(obj, event) 54 | 55 | def mousePressEvent(self, ev): 56 | self.clicked.emit() 57 | 58 | 59 | class DoubleClickableTextEdit(QTextEdit): 60 | doubleclicked = pyqtSignal() 61 | 62 | def mouseDoubleClickEvent(self, ev): 63 | self.doubleclicked.emit() 64 | 65 | 66 | class MiniUtilDialog(QDialog): 67 | def __init__(self, parent, mode): 68 | super(MiniUtilDialog, self).__init__(parent) 69 | self.mode = mode 70 | self.setAcceptDrops(True) 71 | self.setWindowTitle("태그 확인기" if self.mode == 72 | "tagger" else "이미지 정보 확인기") 73 | 74 | # 레이아웃 설정 75 | layout = QVBoxLayout() 76 | frame = BackgroundFrame(self, opacity=0.3) 77 | frame.set_background_image_by_src( 78 | PATH_IMG_TAGGER if self.mode == "tagger" else PATH_IMG_GETTER) 79 | self.frame = frame 80 | self.parent().installEventFilter(frame) 81 | frame.setFixedSize(QSize(512, 512)) 82 | frame.setStyleSheet("""QLabel{ 83 | font-size:30px; 84 | } 85 | QTextEdit{ 86 | font-size:20px; 87 | background-color:#00000000; 88 | }""") 89 | 90 | inner_layout = QVBoxLayout() 91 | label1 = QLabel("Double Click Me") 92 | label1.setAlignment(Qt.AlignCenter) 93 | inner_layout.addWidget(label1) 94 | label_content = DoubleClickableTextEdit("") 95 | label_content.setReadOnly(True) 96 | label_content.setAcceptRichText(False) 97 | label_content.setFrameStyle(QFrame.NoFrame) 98 | label_content.setAlignment(Qt.AlignCenter) 99 | inner_layout.addWidget(label_content, stretch=999) 100 | label2 = QLabel("Or Drag-Drop Here") 101 | label2.setAlignment(Qt.AlignCenter) 102 | inner_layout.addWidget(label2) 103 | frame.setLayout(inner_layout) 104 | 105 | label_content.doubleclicked.connect(self.show_file_minidialog) 106 | 107 | self.label_content = label_content 108 | self.label1 = label1 109 | self.label2 = label2 110 | 111 | layout.addWidget(frame) 112 | self.setLayout(layout) 113 | 114 | def set_content(self, src, nai_dict): 115 | if isinstance(src, str) and os.path.isfile(src): 116 | self.frame.set_background_image_by_src(src) 117 | else: 118 | self.frame.set_background_image_by_img(src) 119 | self.frame.opacity = 0.1 120 | self.label1.setVisible(False) 121 | self.label2.setVisible(False) 122 | self.label_content.setEnabled(True) 123 | self.label_content.setText(nai_dict) 124 | 125 | def execute(self, filemode, target): 126 | if self.mode == "getter": 127 | if filemode == "src": 128 | nai_dict, error_code = get_naidict_from_file( 129 | target) 130 | elif filemode == "img": 131 | nai_dict, error_code = get_naidict_from_img( 132 | target) 133 | elif filemode == "txt": 134 | nai_dict, error_code = get_naidict_from_txt( 135 | target) 136 | 137 | if nai_dict and error_code == 0: 138 | target_dict = { 139 | "prompt": nai_dict["prompt"], 140 | "negative_prompt": nai_dict["negative_prompt"] 141 | } 142 | target_dict.update(nai_dict["option"]) 143 | target_dict.update(nai_dict["etc"]) 144 | 145 | if 'reference_strength' in target_dict: 146 | rs_float = 0.0 147 | try: 148 | rs_float = float(target_dict['reference_strength']) 149 | except Exception as e: 150 | pass 151 | if rs_float > 0: 152 | target_dict["reference_image"] = "True" 153 | if "request_type" in target_dict and target_dict["request_type"] == "Img2ImgRequest": 154 | target_dict["image"] = "True" 155 | 156 | self.set_content(target, prettify_naidict(target_dict)) 157 | 158 | return 159 | elif self.mode == "tagger": 160 | window = self.parent() 161 | result = predict_tag_from(window, filemode, target, True) 162 | 163 | if result: 164 | self.set_content(target, result) 165 | return 166 | QMessageBox.information( 167 | self, '경고', "불러오는 중에 오류가 발생했습니다.") 168 | 169 | def show_file_minidialog(self): 170 | select_dialog = QFileDialog() 171 | select_dialog.setFileMode(QFileDialog.ExistingFile) 172 | fname, _ = select_dialog.getOpenFileName( 173 | self, '불러올 파일을 선택해 주세요.', '', 174 | '이미지 파일(*.png *.webp)' if self.mode == "tagger" else 175 | '이미지, 텍스트 파일(*.txt *.png *.webp)') 176 | 177 | if fname: 178 | if self.mode == "tagger": 179 | if fname.endswith(".png") or fname.endswith(".webp"): 180 | self.execute("src", fname) 181 | else: 182 | QMessageBox.information( 183 | self, '경고', "png, webp, txt 파일만 가능합니다.") 184 | return 185 | else: 186 | if fname.endswith(".png") or fname.endswith(".webp"): 187 | self.execute("src", fname) 188 | elif fname.endswith(".txt"): 189 | self.execute("txt", fname) 190 | QMessageBox.information( 191 | self, '경고', "png, webp 또는 폴더만 가능합니다.") 192 | return 193 | 194 | def dragEnterEvent(self, event): 195 | if event.mimeData().hasUrls(): 196 | event.accept() 197 | else: 198 | event.ignore() 199 | 200 | def dropEvent(self, event): 201 | files = [u for u in event.mimeData().urls()] 202 | 203 | if len(files) != 1: 204 | QMessageBox.information(self, '경고', "파일을 하나만 옮겨주세요.") 205 | return 206 | 207 | furl = files[0] 208 | if furl.isLocalFile(): 209 | fname = furl.toLocalFile() 210 | if fname.endswith(".png") or fname.endswith(".webp"): 211 | self.execute("src", fname) 212 | return 213 | elif fname.endswith(".txt"): 214 | if self.mode == "getter": 215 | self.execute("txt", fname) 216 | return 217 | elif fname.endswith('.jpg') and self.mode == 'tagger': 218 | self.execute("src", fname) 219 | return 220 | 221 | if self.mode == "getter": 222 | QMessageBox.information( 223 | self, '경고', "세팅 불러오기는 png, webp, txt 파일만 가능합니다.") 224 | else: 225 | QMessageBox.information( 226 | self, '경고', "태그 불러오기는 jpg, png, webp 파일만 가능합니다.") 227 | else: 228 | try: 229 | url = furl.url() 230 | res = request.urlopen(url).read() 231 | img = Image.open(BytesIO(res)) 232 | if img: 233 | self.execute("img", img) 234 | 235 | except Exception as e: 236 | print(e) 237 | QMessageBox.information(self, '경고', "이미지 파일 다운로드에 실패했습니다.") 238 | return 239 | 240 | 241 | if __name__ == '__main__': 242 | app = QApplication(sys.argv) 243 | from PyQt5.QtWidgets import QMainWindow 244 | from PyQt5.QtCore import QSettings 245 | qw = QMainWindow() 246 | qw.move(200, 200) 247 | TOP_NAME = "dcp_arca" 248 | APP_NAME = "nag_gui" 249 | qw.settings = QSettings(TOP_NAME, APP_NAME) 250 | qw.dtagger = DanbooruTagger(qw.settings.value( 251 | "path_models", os.path.abspath(DEFAULT_PATH["path_models"]))) 252 | loading_dialog = MiniUtilDialog(qw, "getter") 253 | if loading_dialog.exec_() == QDialog.Accepted: 254 | print(len(loading_dialog.result)) 255 | -------------------------------------------------------------------------------- /app/gui/dialog/inpaint_dialog.py: -------------------------------------------------------------------------------- 1 | import sys 2 | from collections import deque 3 | from PyQt5.QtWidgets import QMainWindow, QDialog, QApplication, QLabel, QSlider, QVBoxLayout, QWidget, QPushButton, QHBoxLayout, QCheckBox, QShortcut 4 | from PyQt5.QtGui import QPainter, QPen, QImage, qRgb, qRgba, QKeySequence 5 | from PyQt5.QtCore import Qt, QPoint, QSize, QEvent 6 | 7 | def convert_coord(point, original_image_size, now_image_rect): 8 | wseed = original_image_size.width() / now_image_rect.width() 9 | hseed = original_image_size.height() / now_image_rect.height() 10 | 11 | x = int((point.x() - now_image_rect.left()) * wseed) 12 | y = int((point.y() - now_image_rect.top()) * hseed) 13 | 14 | return QPoint(x, y) 15 | 16 | 17 | def convert_brush_coord(brush_size, original_image_size, now_image_rect): 18 | if original_image_size.width() > original_image_size.height(): 19 | seed = original_image_size.width() / now_image_rect.width() 20 | else: 21 | seed = original_image_size.height() / now_image_rect.height() 22 | return brush_size * seed * 2 23 | 24 | 25 | # 빈곳은 검게, 있는 곳은 하얗게 만든다. 투명도가 없어진다. 26 | # 하얀곳이 하나도 없으면 False를 내뱉는다. 27 | def convert_pixel_fit_to_mask(image): 28 | is_empty = True 29 | for y in range(image.height()): 30 | for x in range(image.width()): 31 | pixel = image.pixel(x, y) 32 | if pixel >> 24 == 0: 33 | image.setPixel(x, y, qRgb(0, 0, 0)) 34 | else: 35 | image.setPixel(x, y, qRgb(255, 255, 255)) 36 | is_empty = False 37 | 38 | return is_empty 39 | 40 | 41 | def convert_mask_to_pixel_fit(image): 42 | for y in range(image.height()): 43 | for x in range(image.width()): 44 | pixel = image.pixel(x, y) 45 | if pixel == 4278190080: 46 | image.setPixel(x, y, qRgba(0, 0, 0, 0)) 47 | else: 48 | image.setPixel(x, y, qRgba(0, 0, 0, 255)) 49 | 50 | 51 | """ 52 | img = QImage('target.png') 53 | # mask = QImage('mask.png') 54 | # d = InpaintDialog(img, mask) 55 | d = InpaintDialog(img) 56 | if d.exec_() == QDialog.Accepted: 57 | print(d.mask_add) 58 | print(d.mask_only) 59 | 60 | 이때, 변화가 없으면 d.mask_only=None 61 | """ 62 | 63 | 64 | class InpaintDialog(QDialog): 65 | def __init__(self, img: QImage, mask: QImage = None, undo_max_length=10): 66 | super().__init__() 67 | self.image_original_size = QSize() 68 | 69 | self.initUI(undo_max_length) 70 | self.load_image(img, mask) 71 | 72 | def initUI(self, undo_max_length): 73 | self.image = QImage() 74 | self.draw_image = QImage() 75 | self.drawing = False 76 | self.is_erase_mod = False 77 | self.last_image_deque = deque(maxlen=undo_max_length) 78 | self.last_point = QPoint() 79 | self.brush_size = 40 80 | self.now_mouse_pos = QPoint(0, 0) 81 | 82 | # Label for Image 83 | self.image_label = QLabel(self) 84 | self.image_label.setMinimumSize( 85 | QSize(400, 400)) 86 | self.image_label.installEventFilter(self) 87 | self.image_label.setMouseTracking(True) 88 | 89 | # Slider for brush size 90 | self.slider = QSlider(Qt.Horizontal, self) 91 | self.slider.setRange(10, 200) 92 | self.slider.setValue(40) 93 | self.slider.valueChanged.connect(self.change_brush_size) 94 | 95 | layout_buttons = QHBoxLayout() 96 | erase_box = QCheckBox('지우개') 97 | erase_box.clicked.connect(self.on_check_erase) 98 | undo_button = QPushButton('되돌리기') 99 | undo_button.clicked.connect(self.on_click_undo) 100 | clear_button = QPushButton('모두 지우기') 101 | clear_button.clicked.connect(self.on_click_clear) 102 | 103 | # Button to save mask 104 | save_button = QPushButton('저장') 105 | save_button.clicked.connect(self.save_mask) 106 | 107 | quit_button = QPushButton('취소') 108 | quit_button.clicked.connect(self.on_quit_button) 109 | 110 | layout_buttons.addWidget(erase_box) 111 | layout_buttons.addWidget(undo_button) 112 | layout_buttons.addWidget(clear_button) 113 | layout_buttons.addWidget(save_button) 114 | layout_buttons.addWidget(quit_button) 115 | 116 | layout = QVBoxLayout() 117 | layout.addWidget(self.image_label) 118 | layout.addWidget(self.slider) 119 | layout.addLayout(layout_buttons) 120 | self.setLayout(layout) 121 | 122 | shortcut = QShortcut(QKeySequence(Qt.CTRL + Qt.Key_Z), self) 123 | shortcut.activated.connect(self.on_press_ctrlz) 124 | 125 | def mousePressEvent(self, event): 126 | if event.button() == Qt.LeftButton and not self.draw_image.isNull(): 127 | self.drawing = True 128 | self.last_point = event.pos() 129 | self.last_image_deque.append(self.draw_image.copy()) 130 | 131 | def mouseMoveEvent(self, event): 132 | if event.buttons() & Qt.LeftButton and self.drawing: 133 | painter = QPainter(self.draw_image) 134 | painter.setBackgroundMode(Qt.BGMode.OpaqueMode) 135 | brush_size = convert_brush_coord( 136 | self.brush_size, self.image_original_size, self.now_image_rect) 137 | 138 | painter.setPen(QPen(Qt.black, brush_size, 139 | Qt.SolidLine, Qt.RoundCap, Qt.RoundJoin)) 140 | 141 | p1 = convert_coord( 142 | event.pos(), self.image_original_size, self.now_image_rect) 143 | p2 = convert_coord( 144 | self.last_point, self.image_original_size, self.now_image_rect) 145 | 146 | if not self.is_erase_mod: 147 | painter.drawLine(p1, p2) 148 | else: 149 | brush_pos = QPoint( 150 | int(p1.x() - brush_size / 4), 151 | int(p2.y() - brush_size / 4) 152 | ) 153 | 154 | painter.setCompositionMode(QPainter.CompositionMode_Clear) 155 | painter.eraseRect(brush_pos.x(), brush_pos.y(), int( 156 | brush_size/2), int(brush_size/2)) 157 | 158 | self.last_point = event.pos() 159 | self.update() 160 | 161 | def mouseReleaseEvent(self, event): 162 | if event.button() == Qt.LeftButton: 163 | self.drawing = False 164 | 165 | def eventFilter(self, obj, event): 166 | if event.type() == QEvent.MouseMove: 167 | self.onHovered(event) 168 | return super(QWidget, self).eventFilter(obj, event) 169 | 170 | def onHovered(self, event): 171 | self.now_mouse_pos = event.pos() 172 | self.update() 173 | 174 | def paintEvent(self, event): 175 | rect = self.image_label.rect() 176 | rect.setWidth(self.width()) 177 | 178 | w, h = self.image_original_size.width(), self.image_original_size.height() 179 | if rect.width() > rect.height(): 180 | new_width = int(rect.height() * w / h) 181 | rect.moveLeft(int((rect.width() - new_width) * 0.5)) 182 | rect.setWidth(new_width) 183 | else: 184 | new_height = int(rect.width() * h / w) 185 | rect.moveTop(int((rect.height() - new_height) * 0.5)) 186 | rect.setHeight(new_height) 187 | self.now_image_rect = rect 188 | 189 | canvas_painter = QPainter(self) 190 | canvas_painter.drawImage(rect, self.image) 191 | canvas_painter.setOpacity(0.5) 192 | canvas_painter.drawImage(rect, self.draw_image) 193 | canvas_painter.setOpacity(0.3) 194 | canvas_painter.setPen(QPen(Qt.black)) 195 | if not self.is_erase_mod: 196 | brush_size = self.brush_size 197 | # brush_pos = self.now_mouse_pos + 20 198 | brush_pos = QPoint( 199 | int(self.now_mouse_pos.x() + 15), 200 | int(self.now_mouse_pos.y() + 15) 201 | ) 202 | canvas_painter.drawEllipse(brush_pos, brush_size, brush_size) 203 | else: 204 | brush_size = self.brush_size 205 | # brush_pos = self.now_mouse_pos + 20 206 | brush_pos = QPoint( 207 | int(self.now_mouse_pos.x() + 15 - brush_size / 2), 208 | int(self.now_mouse_pos.y() + 15 - brush_size / 2) 209 | ) 210 | canvas_painter.drawRect( 211 | brush_pos.x(), brush_pos.y(), brush_size, brush_size) 212 | 213 | def resizeEvent(self, event): 214 | self.window_size = self.size() 215 | 216 | def load_image(self, img, mask=None): 217 | self.image = img 218 | w, h = self.image.width(), self.image.height() 219 | if mask: 220 | self.draw_image = mask 221 | convert_mask_to_pixel_fit(self.draw_image) 222 | else: 223 | self.draw_image = QImage(w, h, QImage.Format_ARGB32) 224 | # self.draw_image.invertPixels(QImage.InvertRgb); 225 | self.image_original_size = QSize(w, h) 226 | self.update() 227 | # self.resize(w, h) 228 | 229 | def save_mask(self): 230 | # Saving black mask 231 | mask_only = self.draw_image.copy() 232 | 233 | mask_add = self.image.copy() 234 | painter = QPainter(mask_add) 235 | painter.setOpacity(0.5) 236 | painter.drawImage(mask_only.rect(), mask_only) 237 | 238 | self.mask_add = mask_add 239 | is_empty = convert_pixel_fit_to_mask(mask_only) 240 | self.mask_only = None if is_empty else mask_only 241 | self.accept() 242 | 243 | def change_brush_size(self, value): 244 | self.brush_size = value 245 | self.update() 246 | 247 | def on_check_erase(self, is_erase_mod): 248 | self.is_erase_mod = is_erase_mod 249 | self.update() 250 | 251 | def on_click_undo(self): 252 | if len(self.last_image_deque) > 0: 253 | self.draw_image = self.last_image_deque.pop() 254 | self.update() 255 | 256 | def on_click_clear(self): 257 | self.last_image_deque.append(self.draw_image.copy()) 258 | self.draw_image.fill(0) 259 | self.update() 260 | 261 | def on_press_ctrlz(self): 262 | self.on_click_undo() 263 | 264 | def on_quit_button(self): 265 | self.reject() 266 | 267 | 268 | if __name__ == '__main__': 269 | app = QApplication(sys.argv) 270 | qw = QMainWindow() 271 | qw.move(200, 200) 272 | img = QImage('getter.png') 273 | # mask = QImage('mask.png') 274 | # d = InpaintDialog(img, mask) 275 | d = InpaintDialog(img) 276 | if d.exec_() == QDialog.Accepted: 277 | print(d.mask_add) 278 | print(d.mask_only) 279 | # d.mask_add.save('mask1.png') 280 | # d.mask_only.save('mask2.png') 281 | -------------------------------------------------------------------------------- /app/gui/layout/image_options_layout.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtWidgets import QWidget, QGroupBox, QFrame, QFileDialog,QScrollArea, QVBoxLayout, QHBoxLayout, QPushButton, QDialog, QSizePolicy 2 | from PyQt5.QtGui import QImage, QPainter 3 | from PyQt5.QtCore import Qt, QSize, QRectF, QEvent 4 | 5 | from gui.widget.custom_slider_widget import CustomSliderWidget 6 | from gui.widget.result_image_view import ResultImageView 7 | 8 | from gui.dialog.inpaint_dialog import InpaintDialog 9 | 10 | from util.image_util import convert_src_to_imagedata 11 | 12 | from config.paths import PATH_IMG_OPEN_IMAGE, PATH_IMG_IMAGE_CLEAR 13 | from config.themes import COLOR 14 | 15 | PADDING_IMAGE_ITEM = 15 16 | HEIGHT_IMAGE_ITEM = 150 17 | MINHEIGHT_IMAGE_ITEM = 180 18 | MAXHEIGHT_IMAGE_ITEM = 350 19 | 20 | VIBE_SLIDER_OPTION1 = { 21 | "title": "Information Extracted:", 22 | "min_value": 1, 23 | "max_value": 100, 24 | "edit_width": 50, 25 | "mag": 100 26 | } 27 | VIBE_SLIDER_OPTION2 = { 28 | "title": "Reference Strength: ", 29 | "min_value": 1, 30 | "max_value": 100, 31 | "edit_width": 50, 32 | "mag": 100 33 | } 34 | I2I_SLIDER_OPTION1 = { 35 | "title":"Strength:", 36 | "min_value":1, 37 | "max_value":99, 38 | "edit_width":50, 39 | "mag":100 40 | } 41 | I2I_SLIDER_OPTION2 = { 42 | "title": "Noise: ", 43 | "min_value": 0, 44 | "max_value": 99, 45 | "edit_width": 50, 46 | "mag": 100 47 | } 48 | 49 | 50 | class ImageFrame(QFrame): 51 | def __init__(self, parent=None): 52 | super(ImageFrame, self).__init__(parent) 53 | 54 | self.image = QImage() 55 | 56 | def set_image_by_src(self, image_path): 57 | self.image.load(image_path) 58 | self.update() 59 | 60 | def set_image_by_img(self, image): 61 | self.image = image 62 | self.update() 63 | 64 | def paintEvent(self, event): 65 | painter = QPainter(self) 66 | scaled_image = self.image.scaled( 67 | self.size(), Qt.KeepAspectRatioByExpanding, Qt.SmoothTransformation) 68 | target_rect = QRectF( 69 | (self.width() - scaled_image.width()) / 2, 70 | (self.height() - scaled_image.height()) / 2, 71 | scaled_image.width(), 72 | scaled_image.height()) 73 | painter.drawImage(target_rect, scaled_image) 74 | 75 | def eventFilter(self, obj, event): 76 | if event.type() == QEvent.Resize: 77 | self.update() 78 | return True 79 | return super(ImageFrame, self).eventFilter(obj, event) 80 | 81 | # CustomSliderWidget을 QWidget처럼 사용하기 위한 컨테이너 클래스 82 | class CustomSliderWidgetContainer(QWidget): 83 | def __init__(self, **option_dict): 84 | super().__init__() 85 | layout = CustomSliderWidget(**option_dict) 86 | self.setLayout(layout) 87 | # CustomSliderWidget 내부에서 slider를 속성으로 저장했으므로 그대로 노출 88 | self.slider = layout.slider 89 | self.edit = layout.edit 90 | 91 | class ImageItem(QWidget): 92 | def __init__(self, parent, is_i2i, image_path, parent_layout, settings, remove_callback, index, on_click_inpaint_button=None): 93 | super().__init__(parent) 94 | self.is_i2i = is_i2i 95 | self.image_path = image_path 96 | self.parent_layout = parent_layout 97 | self.settings = settings 98 | self.remove_callback = remove_callback 99 | self.index = index # 생성 시 할당된 인덱스 100 | 101 | # 초기 작업 102 | self.setFixedHeight(HEIGHT_IMAGE_ITEM) # 각 아이템 높이 HEIGHT_IMAGE_ITEM 103 | self.imagedata = convert_src_to_imagedata(image_path) 104 | 105 | # 전체 레이아웃 (HBox: 좌측 이미지, 우측 슬라이더 영역) 106 | main_layout = QHBoxLayout(self) 107 | main_layout.setContentsMargins(0, 0, 0, 0) 108 | main_layout.setSpacing(10) 109 | 110 | # 좌측: 이미지 프레임 (HEIGHT_IMAGE_ITEM) 111 | image_frame = ImageFrame() 112 | # image_frame.setStyleSheet("border: 1px solid gray;") 113 | image_frame.setFixedSize(HEIGHT_IMAGE_ITEM, HEIGHT_IMAGE_ITEM) 114 | image_frame.set_image_by_src(self.image_path) 115 | delete_button = QPushButton("X", image_frame) 116 | delete_button.setFixedSize(24, 24) 117 | delete_button.setStyleSheet("background-color: red; color: white; border-radius: 12px;") 118 | delete_button.move(15, 15) 119 | delete_button.clicked.connect(self.delete_self) 120 | main_layout.addWidget(image_frame) 121 | self.image_frame = image_frame 122 | 123 | # 우측: 슬라이더 영역 (VBox) 124 | slider_vbox = QVBoxLayout() 125 | slider_vbox.setSpacing(10) 126 | slider_vbox.setContentsMargins(0, 0, 0, 0) 127 | 128 | # 저장된 값 복원 (저장된 값이 없으면 기본값 50) 129 | mode_str = "i2i" if self.is_i2i else "vibe" 130 | val1 = int(self.settings.value(f"{mode_str}_slider_value_{self.index}_1", 50)) 131 | val2 = int(self.settings.value(f"{mode_str}_slider_value_{self.index}_2", 50)) 132 | 133 | # CustomSliderWidgetContainer로 슬라이더1 생성 134 | slider1_options=I2I_SLIDER_OPTION1 if self.is_i2i else VIBE_SLIDER_OPTION1 135 | slider1_options["default_value"] = val1 136 | slider1_options["slider_text_lambda"] = lambda value: "%.2f" % (value / 100) 137 | custom_slider1 = CustomSliderWidgetContainer(**slider1_options) 138 | # 슬라이더 값 변경 시 QSettings에 저장 139 | custom_slider1.slider.valueChanged.connect(lambda value: self.slider_value_changed(value, 1)) 140 | slider_vbox.addWidget(custom_slider1, stretch=9) 141 | self.custom_slider1 = custom_slider1 142 | 143 | # CustomSliderWidgetContainer로 슬라이더2 생성 144 | slider2_options=I2I_SLIDER_OPTION2 if self.is_i2i else VIBE_SLIDER_OPTION2 145 | slider2_options["default_value"] = val2 146 | slider2_options["slider_text_lambda"] = lambda value: "%.2f" % (value / 100) 147 | custom_slider2 = CustomSliderWidgetContainer(**slider2_options) 148 | # 슬라이더 값 변경 시 QSettings에 저장 149 | custom_slider2.slider.valueChanged.connect(lambda value: self.slider_value_changed(value, 2)) 150 | slider_vbox.addWidget(custom_slider2, stretch=9) 151 | self.custom_slider2 = custom_slider2 152 | 153 | main_layout.addLayout(slider_vbox) 154 | 155 | # i2i는 인페인트 버튼 추가 156 | if is_i2i: 157 | inpaint_button = QPushButton("인페인트") 158 | main_layout.addWidget(inpaint_button) 159 | self.inpaint_button = inpaint_button 160 | if on_click_inpaint_button: 161 | inpaint_button.clicked.connect(lambda: on_click_inpaint_button(self.image_path)) 162 | 163 | def slider_value_changed(self, value, slider_num): 164 | # 슬라이더 값 변경 시 QSettings에 저장 165 | mode_str = "i2i" if self.is_i2i else "vibe" 166 | key = f"{mode_str}_slider_value_{self.index}_{slider_num}" 167 | self.settings.setValue(key, value) 168 | 169 | def change_image(self, qimage, is_mask_applied): 170 | if is_mask_applied: 171 | self.inpaint_button.setStyleSheet(""" 172 | background-color: """ + COLOR.BUTTON_SELECTED + """; 173 | background-position: center; 174 | color: white; 175 | """) 176 | else: 177 | self.inpaint_button.setStyleSheet("") 178 | self.image_frame.set_image_by_img(qimage) 179 | 180 | def delete_self(self): 181 | self.parent_layout.removeWidget(self) 182 | self.deleteLater() 183 | if self.remove_callback: 184 | self.remove_callback() 185 | 186 | def value(self): 187 | return [self.imagedata, self.custom_slider1.edit.text(), self.custom_slider2.edit.text()] 188 | 189 | class ButtonsWidget(QWidget): 190 | def __init__(self, mode, add_callback, clear_callback=None, parent=None): 191 | """ 192 | mode: "initial"이면 add.png와 folder.png 버튼, 193 | "normal"이면 add.png와 clear 버튼. 194 | """ 195 | super().__init__(parent) 196 | self.mode = mode 197 | self.add_callback = add_callback 198 | self.clear_callback = clear_callback 199 | self.setFixedHeight(HEIGHT_IMAGE_ITEM) 200 | layout = QHBoxLayout(self) 201 | layout.setAlignment(Qt.AlignCenter) 202 | 203 | def create_open_button(img_src, func=None): 204 | open_button = ResultImageView( 205 | img_src) 206 | open_button.setStyleSheet(""" 207 | background-color: """ + COLOR.BRIGHT + """; 208 | background-position: center 209 | """) 210 | open_button.setFixedSize(QSize(80, 80)) 211 | if func: 212 | open_button.clicked.connect(func) 213 | return open_button 214 | 215 | # add 버튼 (add.png 아이콘) 216 | layout.addStretch(1) 217 | self.add_button = create_open_button(PATH_IMG_OPEN_IMAGE, self.add_callback) 218 | layout.addWidget(self.add_button) 219 | if mode == "initial": 220 | pass 221 | elif mode == "normal": 222 | layout.addStretch(1) 223 | # clear 버튼 (Clear 텍스트) 224 | self.clear_button = create_open_button(PATH_IMG_IMAGE_CLEAR, self.clear_callback) 225 | layout.addWidget(self.clear_button) 226 | layout.addStretch(1) 227 | 228 | class ImageLayout(QWidget): 229 | def __init__(self, parent, title, is_i2i): 230 | super().__init__(parent) 231 | 232 | self.mask = None 233 | 234 | self.settings = parent.settings 235 | self.is_i2i = is_i2i 236 | 237 | main_layout = QVBoxLayout(self) 238 | groupbox = QGroupBox(title) 239 | main_layout.addWidget(groupbox) 240 | grouplayout = QVBoxLayout(groupbox) 241 | main_layout.setContentsMargins(0,0,0,0) 242 | 243 | self.scroll_area = QScrollArea() 244 | self.scroll_area.setFrameShape(QFrame.NoFrame) 245 | self.scroll_area.setWidgetResizable(True) 246 | grouplayout.addWidget(self.scroll_area) 247 | 248 | self.container = QWidget() 249 | self.scroll_layout = QVBoxLayout(self.container) 250 | self.scroll_layout.setSpacing(PADDING_IMAGE_ITEM) 251 | self.container.setLayout(self.scroll_layout) 252 | self.scroll_area.setWidget(self.container) 253 | 254 | # 처음엔 아무 아이템도 없으므로 초기 모드 버튼 위젯 추가 (높이 HEIGHT_IMAGE_ITEM) 255 | self.buttons_widget = self.create_buttons_widget("initial") 256 | self.scroll_layout.addWidget(self.buttons_widget) 257 | self.update_container_height() 258 | 259 | def create_buttons_widget(self, mode): 260 | """mode에 따라 ButtonsWidget 생성""" 261 | if mode == "initial": 262 | return ButtonsWidget(mode, add_callback=self.on_click_add_item) 263 | else: 264 | return ButtonsWidget(mode, add_callback=self.on_click_add_item, clear_callback=self.clear_all_items) 265 | 266 | def on_click_add_item(self): 267 | """파일 다이얼로그를 통해 이미지 선택 후, ImageItem 추가 (버튼 위젯 교체)""" 268 | file_dialog = QFileDialog() 269 | file_dialog.setFileMode(QFileDialog.ExistingFile) 270 | file_dialog.setNameFilter("Images (*.png *.jpg *.jpeg *.webp)") 271 | if file_dialog.exec_(): 272 | image_path = file_dialog.selectedFiles()[0] 273 | 274 | self._add_item(image_path) 275 | 276 | def _add_item(self, image_path): 277 | # 마지막 위젯(버튼 위젯)이 있으면 제거 278 | last_index = self.scroll_layout.count() - 1 279 | last_widget = self.scroll_layout.itemAt(last_index).widget() 280 | if isinstance(last_widget, ButtonsWidget): 281 | self.scroll_layout.removeWidget(last_widget) 282 | last_widget.deleteLater() 283 | # 새 ImageItem 생성 시 현재 count()를 인덱스로 사용 284 | new_item = ImageItem(self, self.is_i2i, image_path, self.scroll_layout, 285 | settings=self.settings, 286 | remove_callback=self.on_click_remove_button, 287 | index=self.scroll_layout.count(), 288 | on_click_inpaint_button=self.on_click_inpaint_button) 289 | self.scroll_layout.addWidget(new_item) 290 | # 새 버튼 위젯 (normal 모드: add와 clear 버튼) 추가 291 | if not self.is_i2i: 292 | new_buttons = self.create_buttons_widget("normal") 293 | self.scroll_layout.addWidget(new_buttons) 294 | self.update_container_height() 295 | 296 | def update_container_height(self): 297 | # 스크롤 레이아웃에 추가된 위젯 개수를 구함 (ImageItem과 ButtonsWidget 포함) 298 | count = self.scroll_layout.count() 299 | # 계산: 총 높이 = (개수 * (이미지 아이템 높이 + 패딩)) + 상단 패딩 300 | calc_height = count * (HEIGHT_IMAGE_ITEM + PADDING_IMAGE_ITEM) + PADDING_IMAGE_ITEM 301 | # 최소 최대 높이 지정 302 | calc_height = max(calc_height, MINHEIGHT_IMAGE_ITEM) 303 | calc_height = min(calc_height, MAXHEIGHT_IMAGE_ITEM) 304 | # 컨테이너의 최소 높이와 스크롤 영역의 고정 높이를 설정함 305 | self.scroll_area.setFixedHeight(calc_height) 306 | 307 | def on_click_remove_button(self): 308 | if self.is_i2i: 309 | self.mask = None 310 | self.buttons_widget = self.create_buttons_widget("initial") 311 | self.scroll_layout.addWidget(self.buttons_widget) 312 | self.update_container_height() 313 | 314 | def clear_all_items(self): 315 | """Clear 버튼 클릭 시 모든 아이템과 버튼 위젯 제거, QSettings 초기화""" 316 | while self.scroll_layout.count() > 0: 317 | item = self.scroll_layout.takeAt(0) 318 | widget = item.widget() 319 | if widget: 320 | widget.deleteLater() 321 | # 초기 상태: 초기 모드 버튼 위젯 추가 322 | init_buttons = self.create_buttons_widget("initial") 323 | self.scroll_layout.addWidget(init_buttons) 324 | self.update_container_height() 325 | 326 | def on_click_inpaint_button(self, src): 327 | if self.is_i2i: 328 | last_widget = self.scroll_layout.itemAt(0).widget() 329 | if isinstance(last_widget, ImageItem): 330 | img = QImage(src) 331 | mask = self.mask if self.mask else None 332 | 333 | d = InpaintDialog(img, mask) 334 | if d.exec_() == QDialog.Accepted: 335 | is_mask_applied = d.mask_only != None 336 | 337 | last_widget.change_image(d.mask_add, is_mask_applied) 338 | self.mask = d.mask_only 339 | 340 | # i2i는 [img, slider1, slider2, mask_only] 341 | # vibe는 [[img, slider1, slider2]] 342 | def get_nai_param(self): 343 | result = [] 344 | if self.is_i2i: 345 | last_widget = self.scroll_layout.itemAt(0).widget() 346 | if isinstance(last_widget, ImageItem): 347 | result = last_widget.value() 348 | 349 | if self.mask: 350 | result.append(self.mask) 351 | else: 352 | for i in range(self.scroll_layout.count()): 353 | item = self.scroll_layout.itemAt(i) 354 | widget = item.widget() 355 | if isinstance(widget, ImageItem): 356 | values = widget.value() 357 | if values[0]: 358 | result.append(values) 359 | return result 360 | 361 | def init_image_options_layout(self): 362 | image_options_layout = QVBoxLayout() 363 | 364 | # I2I Settings Group 365 | i2i_settings_group = ImageLayout(self, title="I2I Settings", is_i2i=True) 366 | image_options_layout.addWidget(i2i_settings_group ) 367 | 368 | vibe_settings_group = ImageLayout(self, title="Vibe Settings", is_i2i=False) 369 | image_options_layout.addWidget(vibe_settings_group) 370 | 371 | image_options_layout.addStretch() 372 | 373 | # Assign 374 | self.i2i_settings_group = i2i_settings_group 375 | self.vibe_settings_group = vibe_settings_group 376 | # self.dict_ui_settings["strength"] = i2i_settings_group.slider_1.edit 377 | # self.dict_ui_settings["noise"] = i2i_settings_group.slider_2.edit 378 | # self.dict_ui_settings["reference_information_extracted"] = vibe_settings_group.slider_1.edit 379 | # self.dict_ui_settings["reference_strength"] = vibe_settings_group.slider_2.edit 380 | 381 | return image_options_layout 382 | -------------------------------------------------------------------------------- /app/core/worker/nai_generator.py: -------------------------------------------------------------------------------- 1 | from hashlib import blake2b 2 | import argon2 3 | from base64 import urlsafe_b64encode 4 | import requests 5 | import json 6 | import copy 7 | 8 | 9 | # 10 | # 11 | # 12 | # 13 | # 14 | 15 | BASE_URL_MAIN = "https://api.novelai.net" 16 | BASE_URL_IMAGE = "https://image.novelai.net/ai/generate-image" 17 | 18 | 19 | DEFAULT_MODEL_V4 = "NAI Diffusion V4.5 Full" 20 | 21 | SAMPLER_ITEMS_V3 = ['k_euler', 'k_euler_ancestral', 'k_dpmpp_2s_ancestral', "k_dpmpp_2m_sde", 22 | "k_dpmpp_2m", 'k_dpmpp_sde', "ddim_v3"] 23 | 24 | SAMPLER_ITEMS_V4 = ['k_euler_ancestral', 'k_dpmpp_2s_ancestral', "k_dpmpp_2m_sde", 25 | 'k_euler', "k_dpmpp_2m",'k_dpmpp_sde'] 26 | 27 | # 이 dict에 있는 값만 설정 가능. 28 | TARGET_PARAMETERS = { 29 | "model": "k_euler_ancestral", 30 | "prompt": "1girl", 31 | "negative_prompt": "worst quality", 32 | "width": 640, 33 | "height": 640, 34 | "noise_schedule": "karras", 35 | "sampler": "k_euler_ancestral", 36 | "steps": 25, 37 | "seed": 9999999999, 38 | "scale": 6, 39 | "cfg_rescale": 0.3, 40 | "sm": True, 41 | "sm_dyn": True, 42 | "variety_plus": True, 43 | "image": "", # i2i image 44 | "mask": "", # inpant mask 45 | "strength": 0.7, # i2i 세팅값1 46 | "noise": 0.0, # i2i 세팅값2 47 | "reference_image_multiple": [], # 바이브 이미지 48 | "reference_information_extracted_multiple": [0], # 바이브 세팅값1 49 | "reference_strength_multiple": [0], # 바이스 세팅값2 50 | "legacy_uc": False, 51 | "use_coords": False, 52 | "characterPrompts": [] # DEFAULT_PARAMETER_CHARPROMPTS를 생성해서 넣어야함. 53 | } 54 | 55 | DEFAULT_PARAMETER_CHARPROMPTS = { 56 | "prompt": "girl, ", 57 | "uc": "lowres, aliasing, ", 58 | "center": { 59 | "x": 0.5, 60 | "y": 0.5 61 | }, 62 | "enabled": True 63 | } 64 | 65 | 66 | DEFAULT_PARAMETER_CHARCAPTIONS = { 67 | "char_caption": "boy, ", 68 | "centers": [ 69 | { 70 | "x": 0.1, 71 | "y": 0.5 72 | } 73 | ] 74 | } 75 | 76 | 77 | V3_PARAMETERS = { 78 | "params_version": 3, 79 | "width": 1024, 80 | "height": 1024, 81 | "scale": 6, 82 | "sampler": "k_euler_ancestral", 83 | "steps": 28, 84 | "n_samples": 1, 85 | "ucPreset": 0, 86 | "qualityToggle": True, 87 | "sm": False, 88 | "sm_dyn": False, 89 | "dynamic_thresholding": False, 90 | "controlnet_strength": 1, 91 | "legacy": False, 92 | "add_original_image": False, 93 | "cfg_rescale": 0, 94 | "noise_schedule": "karras", 95 | "legacy_v3_extend": False, 96 | "skip_cfg_above_sigma": None, 97 | "use_coords": False, 98 | "seed": 9999999999, 99 | "extra_noise_seed":9999999999, 100 | "characterPrompts": [], 101 | "prompt": "1girl", 102 | "negative_prompt": "nsfw, lowres, {bad}, error, fewer, extra, missing, worst quality, jpeg artifacts, bad quality, watermark, unfinished, displeasing, chromatic aberration, signature, extra digits, artistic error, username, scan, [abstract], 3d, blender, pixel art, realistic, blurry, lowres, error, film grain, scan artifacts, worst quality, bad quality, jpeg artifacts, very displeasing, chromatic aberration, multiple views, logo, too many watermarks, white blank page, blank page, {{{worst quality, bad quality}}}, normal quality, very displeasing, censored, displeasing, gundam, furry, simple background, logo, sign, glitch, sketch, simple coloring", 103 | "reference_image_multiple": [], 104 | "reference_information_extracted_multiple": [], 105 | "reference_strength_multiple": [], 106 | "deliberate_euler_ancestral_bug": False, 107 | "prefer_brownian": True 108 | } 109 | 110 | V4_PARAMETERS = { 111 | "params_version": 3, 112 | "width": 832, 113 | "height": 1216, 114 | "scale": 4.8, 115 | "sampler": "k_euler_ancestral", 116 | "steps": 28, 117 | "n_samples": 1, 118 | "ucPreset": 2, 119 | "qualityToggle": False, 120 | "autoSmea": True, 121 | "dynamic_thresholding": False, 122 | "controlnet_strength": 1, 123 | "legacy": False, 124 | "add_original_image": True, 125 | "cfg_rescale": 0.3, 126 | "noise_schedule": "karras", 127 | "legacy_v3_extend": False, 128 | "skip_cfg_above_sigma": None, 129 | "use_coords": True, 130 | "v4_prompt": { 131 | "caption": { 132 | "base_caption": "1girl", 133 | "char_captions": [], 134 | }, 135 | "use_coords": True, 136 | "use_order": True 137 | }, 138 | "v4_negative_prompt": { 139 | "caption": { 140 | "base_caption": "3d, blender, pixel art, realistic, blurry, lowres, error, film grain, scan artifacts, worst quality, bad quality, jpeg artifacts, very displeasing, chromatic aberration, multiple views, logo, too many watermarks, white blank page, blank page, {{{worst quality, bad quality}}}, normal quality, very displeasing, censored, displeasing, gundam, furry, simple background, logo, sign, glitch, sketch, simple coloring", 141 | "char_captions": [] 142 | }, 143 | "legacy_uc": False 144 | }, 145 | "legacy_uc": False, 146 | "seed": 9999999999, 147 | "extra_noise_seed":9999999999, 148 | "characterPrompts": [], 149 | "negative_prompt": "3d, blender, pixel art, realistic, blurry, lowres, error, film grain, scan artifacts, worst quality, bad quality, jpeg artifacts, very displeasing, chromatic aberration, multiple views, logo, too many watermarks, white blank page, blank page, {{{worst quality, bad quality}}}, normal quality, very displeasing, censored, displeasing, gundam, furry, simple background, logo, sign, glitch, sketch, simple coloring", 150 | "reference_image_multiple": [], 151 | "reference_information_extracted_multiple": [], 152 | "reference_strength_multiple": [], 153 | "deliberate_euler_ancestral_bug": False, 154 | "prefer_brownian": True 155 | } 156 | 157 | V4_5_PARAMETERS = { 158 | "params_version": 3, 159 | "width": 832, 160 | "height": 1216, 161 | "scale": 5, 162 | "sampler": "k_euler", 163 | "steps": 28, 164 | "n_samples": 1, 165 | "ucPreset": 2, 166 | "qualityToggle": True, 167 | "autoSmea": False, 168 | "dynamic_thresholding": False, 169 | "controlnet_strength": 1, 170 | "legacy": False, 171 | "add_original_image": True, 172 | "cfg_rescale": 0.4, 173 | "noise_schedule": "karras", 174 | "legacy_v3_extend": False, 175 | "skip_cfg_above_sigma": None, 176 | "use_coords": True, 177 | "legacy_uc": False, 178 | "normalize_reference_strength_multiple": True, 179 | "inpaintImg2ImgStrength": 1, 180 | "v4_prompt": 181 | { 182 | "caption": 183 | { 184 | "base_caption": "1girl, 1boy,\n\n1.1::aritst:kiira::, artist:gachigachi, 0.7::artist:tianliang duohe fangdongye::, 1.1::artist:odayaka::, year 2025, year 2024, nsfw, dynamic lighting,\n\nloli, flat chest, smooth skin, skinny,\n\nnude, spread legs, pussy, penis, vaginal sex, deep penetration, beds, bedding, rape, crying, cum inside, {{female orgasm}}, trembling girl, heart-shaped pupils, blush, twitching womb,\n\ncovering own mouth, grab legs,\n\nblack long hair,, no text, best quality, very aesthetic, absurdres", 185 | "char_captions": 186 | [] 187 | }, 188 | "use_coords": True, 189 | "use_order": True 190 | }, 191 | "v4_negative_prompt": 192 | { 193 | "caption": 194 | { 195 | "base_caption": "ribs, angry", 196 | "char_captions": 197 | [] 198 | }, 199 | "legacy_uc": False 200 | }, 201 | "characterPrompts": 202 | [], 203 | "negative_prompt": "ribs, angry" 204 | } 205 | 206 | # 모델의 정보를 저장함 207 | # key는 모델 이름으로, 드롭다운에 노출됨 208 | # value에 저장된 bool 값이 UI 노출 여부를 결정함. main_window.on_model_changed 참고. 209 | # @need_call_complete_function 값이 True인 경우, _complete_v4_parameters 함수가 호출됨. 210 | MODEL_INFO_DICT = { 211 | "NAI Diffusion Anime V3": { 212 | "model": "nai-diffusion-3", 213 | "i2i": True, 214 | "inpaint": True, 215 | "vibe": True, 216 | "sm": True, 217 | "sm_dyn": True, 218 | "variety_plus": True, 219 | "characterPrompts": False, 220 | "sampler": SAMPLER_ITEMS_V3, 221 | "inpainting_model": "nai-diffusion-3-inpainting", 222 | "default_parameters": V3_PARAMETERS, 223 | "need_call_complete_function": False 224 | }, 225 | "NAI Diffusion V4 Full": { 226 | "model": "nai-diffusion-4-full", 227 | "i2i": True, 228 | "inpaint": True, 229 | "vibe": True, 230 | "sm": False, 231 | "sm_dyn": False, 232 | "variety_plus": False, 233 | "characterPrompts": True, 234 | "sampler": SAMPLER_ITEMS_V4, 235 | "inpainting_model": "nai-diffusion-4-full-inpainting", 236 | "default_parameters": V4_PARAMETERS, 237 | "need_call_complete_function": True 238 | }, 239 | "NAI Diffusion V4 Curated": { 240 | "model": "nai-diffusion-4-curated-preview", 241 | "i2i": True, 242 | "inpaint": True, 243 | "vibe": True, 244 | "sm": False, 245 | "sm_dyn": False, 246 | "variety_plus": False, 247 | "characterPrompts": True, 248 | "sampler": SAMPLER_ITEMS_V4, 249 | "inpainting_model": "nai-diffusion-4-curated-inpainting", 250 | "default_parameters": V4_PARAMETERS, 251 | "need_call_complete_function": True 252 | }, 253 | "NAI Diffusion V4.5 Full": { 254 | "model": "nai-diffusion-4-5-full", 255 | "i2i": False, 256 | "inpaint": False, 257 | "vibe": False, 258 | "sm": False, 259 | "sm_dyn": False, 260 | "variety_plus": False, 261 | "characterPrompts": True, 262 | "sampler": SAMPLER_ITEMS_V4, 263 | "default_parameters": V4_5_PARAMETERS, 264 | "need_call_complete_function": True 265 | }, 266 | "NAI Diffusion V4.5 Curated": { 267 | "model": "nai-diffusion-4-5-curated", 268 | "i2i": False, 269 | "inpaint": False, 270 | "vibe": False, 271 | "sm": False, 272 | "sm_dyn": False, 273 | "variety_plus": False, 274 | "characterPrompts": True, 275 | "sampler": SAMPLER_ITEMS_V4, 276 | "default_parameters": V4_5_PARAMETERS, 277 | "need_call_complete_function": True 278 | } 279 | } 280 | 281 | # 모든 테이블에 필요한 옵션이 있는지 확인 282 | assert all(all(key in model_info for key in ["model", "i2i", "inpaint", "vibe", "sm", "sm_dyn", "variety_plus", "characterPrompts", "sampler", "default_parameters", "need_call_complete_function"]) for model_info in MODEL_INFO_DICT.values()), "모델에 필요한 옵션이 누락되었습니다" 283 | 284 | 285 | def argon_hash(email: str, password: str, size: int, domain: str) -> str: 286 | pre_salt = f"{password[:6]}{email}{domain}" 287 | # salt 288 | blake = blake2b(digest_size=16) 289 | blake.update(pre_salt.encode()) 290 | salt = blake.digest() 291 | raw = argon2.low_level.hash_secret_raw( 292 | password.encode(), 293 | salt, 294 | 2, 295 | int(2000000 / 1024), 296 | 1, 297 | size, 298 | argon2.low_level.Type.ID, 299 | ) 300 | hashed = urlsafe_b64encode(raw).decode() 301 | return hashed 302 | 303 | 304 | def _complete_v4_parameters(parameters): 305 | # prompt 306 | parameters["v4_prompt"]["caption"]["base_caption"] = parameters["prompt"] 307 | 308 | # prompt -> use_coords 309 | parameters["v4_prompt"]["use_coords"] = parameters["use_coords"] 310 | 311 | # negative prompt 312 | parameters["v4_negative_prompt"]["caption"]["base_caption"] = parameters["negative_prompt"] 313 | 314 | # negative prompt -> legacy_uc 315 | parameters["v4_negative_prompt"]["legacy_uc"] = parameters["legacy_uc"] 316 | 317 | # both -> char_captions 318 | for target in ["v4_prompt", "v4_negative_prompt"]: 319 | for char_dict in parameters["characterPrompts"]: 320 | new_char_dict = copy.deepcopy(DEFAULT_PARAMETER_CHARCAPTIONS) 321 | new_char_dict["char_caption"] = char_dict["prompt"] if target == "v4_prompt" else char_dict["uc"] 322 | new_char_dict["centers"][0]["x"] = char_dict["center"]["x"] 323 | new_char_dict["centers"][0]["y"] = char_dict["center"]["y"] 324 | 325 | parent = parameters[target]["caption"]["char_captions"] 326 | parent.append(new_char_dict) 327 | 328 | # 다음이 포함되어있으면 작동이 안됨. 329 | del parameters["sm"] 330 | del parameters["sm_dyn"] 331 | del parameters["variety_plus"] 332 | 333 | def get_character_prompts_from_v4_prompt(parameters): 334 | characterPrompts, use_coords = None, None 335 | 336 | if "v4_prompt" in parameters and parameters["v4_prompt"] and "v4_negative_prompt" in parameters and parameters["v4_negative_prompt"]: 337 | 338 | # use_coords 339 | if "use_coords" in parameters["v4_prompt"] and parameters["v4_prompt"]["use_coords"] is not None: 340 | use_coords = parameters["v4_prompt"]["use_coords"] 341 | 342 | # characterPrompts 만들기 343 | try: 344 | # 갯수만큼 빈껍데기 만들기 345 | characterPrompts = [copy.deepcopy(DEFAULT_PARAMETER_CHARPROMPTS) for _ in parameters["v4_prompt"]["caption"]["char_captions"]] 346 | # 내용 채우기 347 | for target in ["v4_prompt", "v4_negative_prompt"]: 348 | for i, char_dict in enumerate(parameters[target]["caption"]["char_captions"]): 349 | characterPrompts[i]["prompt" if target=="v4_prompt" else "uc"] = char_dict["char_caption"] 350 | characterPrompts[i]["center"] = { 351 | "x":char_dict["centers"][0]["x"], 352 | "y":char_dict["centers"][0]["y"] 353 | } 354 | except Exception: 355 | characterPrompts = None 356 | 357 | return characterPrompts, use_coords 358 | 359 | 360 | class NAIGenerator(): 361 | def __init__(self): 362 | self.access_token = None 363 | self.username = None 364 | self.password = None 365 | self.parameters = {} 366 | 367 | def set_param_dict(self, param_dict): 368 | allowedKeyList = TARGET_PARAMETERS.keys() 369 | for key in param_dict.keys(): 370 | assert key in allowedKeyList, "[NAIGenerator.set_param_dict] 허용되지않은 키가 있습니다. " + key 371 | 372 | self.parameters = param_dict 373 | 374 | def try_login(self, username, password): 375 | # get_access_key 376 | access_key = argon_hash(username, password, 64, 377 | "novelai_data_access_key")[:64] 378 | try: 379 | # try login 380 | response = requests.post( 381 | f"{BASE_URL_MAIN}/user/login", json={"key": access_key}) 382 | self.access_token = response.json()["accessToken"] 383 | 384 | # if success, save id/pw in 385 | self.username = username 386 | self.password = password 387 | 388 | return True 389 | except Exception as e: 390 | print(e) 391 | 392 | return False 393 | 394 | def get_anlas(self): 395 | try: 396 | response = requests.get(BASE_URL_MAIN + "/user/subscription", headers={ 397 | "Authorization": f"Bearer {self.access_token}"}) 398 | data_dict = json.loads(response.content) 399 | trainingStepsLeft = data_dict['trainingStepsLeft'] 400 | anlas = int(trainingStepsLeft['fixedTrainingStepsLeft']) + \ 401 | int(trainingStepsLeft['purchasedTrainingSteps']) 402 | 403 | return anlas 404 | except Exception as e: 405 | print(e) 406 | 407 | return None 408 | 409 | def generate_image(self): 410 | parameters = {} 411 | 412 | # model 413 | model_name = self.parameters["model"] 414 | model = MODEL_INFO_DICT[model_name]["model"] 415 | model_info = MODEL_INFO_DICT[model_name] 416 | 417 | # action 418 | action = "generate" 419 | if "image" in self.parameters and self.parameters['image']: 420 | action = "img2img" 421 | if 'mask' in self.parameters and self.parameters['mask']: 422 | action = "infill" 423 | model = model_info["inpainting_model"] 424 | 425 | # parameter 생성 426 | parameters = copy.deepcopy(model_info["default_parameters"]) 427 | parameters.update(self.parameters) 428 | 429 | # extraseed 통일 430 | parameters["extra_noise_seed"] = parameters["seed"] 431 | 432 | # smea 버그 수정 433 | if "image" in parameters or parameters['sampler'] == 'ddim_v3': 434 | parameters['sm'] = False 435 | parameters['sm_dyn'] = False 436 | 437 | # vibe 아니면 관련 제거 438 | if not model_info["vibe"]: 439 | parameters["reference_image_multiple"] = [] 440 | parameters["reference_information_extracted_multiple"]= [] 441 | parameters["reference_strength_multiple"]= [] 442 | 443 | # _complete_v4_parameters 함수가 호출됨. 444 | if model_info["need_call_complete_function"]: 445 | _complete_v4_parameters(parameters) 446 | 447 | try: 448 | response = requests.post(url=BASE_URL_IMAGE, 449 | json={ 450 | "input": parameters["prompt"], 451 | "model": model, 452 | "action": action, 453 | "parameters": parameters 454 | }, 455 | headers={ 456 | "Authorization": f"Bearer " + self.access_token} 457 | ) 458 | return response.content 459 | except Exception as e: 460 | print("[generate_image]", e) 461 | 462 | return None 463 | 464 | def check_logged_in(self): 465 | access_result = None 466 | try: 467 | access_result = requests.get(BASE_URL_MAIN + "/user/information", headers={ 468 | "Authorization": f"Bearer {self.access_token}"}, timeout=5) 469 | except Exception as e: 470 | print(e) 471 | return (access_result is not None) 472 | -------------------------------------------------------------------------------- /app/gui/layout/character_prompt_layout.py: -------------------------------------------------------------------------------- 1 | from PyQt5.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QScrollArea, QFrame, 2 | QGridLayout, QDialog, QCheckBox, QSizePolicy) 3 | from PyQt5.QtCore import Qt, pyqtSignal 4 | 5 | from core.worker.completer import CompletionTextEdit 6 | 7 | from util.ui_util import add_button 8 | 9 | from config.themes import COLOR 10 | 11 | WIDTH_CHARITEM = 400 12 | HEIGHT_CHARITEM = 270 13 | HEIGHT_CONTAINER_EXPAND = 400 14 | HEIGHT_CONTAINER_REDUCE = 80 15 | 16 | def _get_position_text(position): 17 | return f"{int(position[0] * 5)-2}, {(4 - int(position[1] * 5 )- 2)}" 18 | 19 | # stateChanged와는 다르게 값이 변경이 되든 안되든 콜백이 호출한다. 20 | class CustomCheckBox(QCheckBox): 21 | onSetCheckedCalled = pyqtSignal(bool) 22 | 23 | def setChecked(self, checked): 24 | super().setChecked(checked) 25 | self.onSetCheckedCalled.emit(checked) 26 | 27 | class PositionSelectorDialog(QDialog): 28 | """캐릭터 위치 선택 다이얼로그""" 29 | 30 | def __init__(self, parent=None): 31 | super().__init__(parent) 32 | self.setWindowTitle("캐릭터 위치 선택") 33 | self.selected_position = None 34 | self.setup_ui() 35 | 36 | def setup_ui(self): 37 | try: 38 | layout = QVBoxLayout() 39 | self.setLayout(layout) 40 | self.setStyleSheet("QVBoxLayout{background-color: "+COLOR.DARK+"};}") 41 | 42 | # 설명 라벨 43 | info_label = QLabel("원하는 위치를 클릭하세요 (5x5 그리드)") 44 | layout.addWidget(info_label) 45 | 46 | # 그리드 생성 (5x5) 47 | grid_layout = QGridLayout() 48 | self.buttons = [] 49 | 50 | for row in range(5): 51 | for col in range(5): 52 | btn = QPushButton() 53 | btn.setStyleSheet("background-color: "+ COLOR.WHITE) 54 | btn.setText(_get_position_text(((col + 0.5) / 5, (row + 0.5) / 5))) 55 | btn.setFixedSize(50, 50) 56 | # 람다 함수에 위치 정보를 명시적으로 전달 57 | btn.clicked.connect(lambda checked=False, r=row, c=col: self.on_position_selected(r, c)) 58 | grid_layout.addWidget(btn, row, col) 59 | self.buttons.append(btn) 60 | 61 | layout.addLayout(grid_layout) 62 | 63 | # 완료/취소 버튼 64 | buttons_layout = QHBoxLayout() 65 | done_button = QPushButton("완료") 66 | done_button.clicked.connect(self.accept) 67 | cancel_button = QPushButton("취소") 68 | cancel_button.clicked.connect(self.reject) 69 | 70 | buttons_layout.addStretch() 71 | buttons_layout.addWidget(done_button) 72 | buttons_layout.addWidget(cancel_button) 73 | layout.addLayout(buttons_layout) 74 | except Exception as e: 75 | print(f"위치 선택자 UI 초기화 오류: {e}") 76 | 77 | def on_position_selected(self, row, col): 78 | try: 79 | print(f"위치 선택: 행={row}, 열={col}") 80 | # 다른 버튼들은 원래 색으로 복원 81 | for btn in self.buttons: 82 | btn.setStyleSheet("background-color: "+ COLOR.WHITE) 83 | 84 | # 선택된 버튼 하이라이트 85 | index = row * 5 + col 86 | self.buttons[index].setStyleSheet(f"background-color: {COLOR.BUTTON_SELECTED};") 87 | 88 | # 위치 저장 (0~1 범위의 비율로 저장) 89 | self.selected_position = ((col + 0.5) / 5, (row + 0.5) / 5) 90 | print(f"저장된 위치: {self.selected_position}") 91 | except Exception as e: 92 | print(f"위치 선택 처리 오류: {e}") 93 | 94 | class CharacterPromptWidget(QFrame): 95 | """캐릭터 프롬프트를 입력하고 관리하는 위젯""" 96 | 97 | deleted = pyqtSignal(object) # 삭제 시그널 98 | moved = pyqtSignal(object, int) # 이동 시그널 (위젯, 방향) 99 | 100 | def __init__(self, parent=None, index=0): 101 | super().__init__(parent) 102 | self.parent = parent 103 | self.index = index 104 | self.position = None # 캐릭터 위치 (None = AI 선택) 105 | self.show_negative = False # 네거티브 프롬프트 표시 여부 106 | 107 | # 스타일 설정 108 | self.setFrameShape(QFrame.StyledPanel) 109 | self.setFrameShadow(QFrame.Raised) 110 | 111 | # 수직 레이아웃으로 설계 112 | self.setup_ui() 113 | 114 | # 위젯 크기 정책 설정 115 | self.setSizePolicy(QSizePolicy.Fixed, QSizePolicy.Expanding) 116 | self.setFixedWidth(WIDTH_CHARITEM) # 최소 너비 설정 117 | self.setFixedHeight(HEIGHT_CHARITEM) # 최소 너비 설정 118 | 119 | # 새로운 메서드 추가 120 | def update_title(self): 121 | """타이틀 업데이트""" 122 | try: 123 | header_layout = self.layout.itemAt(0).layout() 124 | if header_layout: 125 | title_label = header_layout.itemAt(0).widget() 126 | if title_label: 127 | title_label.setText(f"캐릭터 {self.index + 1}") 128 | except Exception as e: 129 | print(f"타이틀 업데이트 중 오류: {e}") 130 | 131 | def setup_ui(self): 132 | self.layout = QVBoxLayout() 133 | self.setLayout(self.layout) 134 | 135 | # 헤더 (타이틀 + 컨트롤 버튼) 136 | header_layout = QHBoxLayout() 137 | 138 | title_label = QLabel(f"캐릭터 {self.index + 1}") 139 | header_layout.addWidget(title_label) 140 | 141 | header_layout.addStretch() 142 | 143 | # 네거 버튼 144 | self.negative_btn = QPushButton("네거") 145 | self.negative_btn.setFixedWidth(50) 146 | self.negative_btn.clicked.connect(self.toggle_negative_prompt) 147 | 148 | # 캐릭터 순서 이동 버튼 149 | move_up_btn = QPushButton("◀") 150 | move_up_btn.setFixedWidth(50) 151 | move_up_btn.clicked.connect(lambda: self.moved.emit(self, -1)) 152 | 153 | move_down_btn = QPushButton("▶") 154 | move_down_btn.setFixedWidth(50) 155 | move_down_btn.clicked.connect(lambda: self.moved.emit(self, 1)) 156 | 157 | # 위치 설정 버튼 158 | self.position_btn = QPushButton("위치") 159 | self.position_btn.setFixedWidth(50) 160 | self.position_btn.setEnabled(True) 161 | self.position_btn.clicked.connect(self.show_position_dialog) 162 | 163 | # 삭제 버튼 164 | delete_btn = QPushButton("✕") 165 | delete_btn.setFixedWidth(50) 166 | delete_btn.clicked.connect(lambda: self.deleted.emit(self)) 167 | 168 | header_layout.addWidget(self.negative_btn) 169 | header_layout.addWidget(move_up_btn) 170 | header_layout.addWidget(move_down_btn) 171 | header_layout.addWidget(self.position_btn) 172 | header_layout.addWidget(delete_btn) 173 | 174 | self.layout.addLayout(header_layout) 175 | 176 | # QPlainTextEdit 대신 CompletionTextEdit 사용 177 | self.prompt_edit = CompletionTextEdit() 178 | self.prompt_edit.setPlaceholderText("프롬프트 입력...") 179 | self.prompt_edit.setMinimumHeight(80) 180 | self.layout.addWidget(self.prompt_edit) 181 | 182 | # 네거티브 프롬프트 입력에도 CompletionTextEdit 사용 183 | self.neg_prompt_edit = CompletionTextEdit() 184 | self.neg_prompt_edit.setPlaceholderText("네거티브 프롬프트 입력...") 185 | self.neg_prompt_edit.setMinimumHeight(80) 186 | self.neg_prompt_edit.setVisible(False) 187 | self.layout.addWidget(self.neg_prompt_edit) 188 | 189 | def toggle_negative_prompt(self, state): 190 | """네거티브 프롬프트 표시/숨김 전환""" 191 | promptVisible = self.prompt_edit.isVisible() 192 | npromptVisible = self.neg_prompt_edit.isVisible() 193 | self.prompt_edit.setVisible(npromptVisible) 194 | self.neg_prompt_edit.setVisible(promptVisible) 195 | 196 | nowPromptVisible = npromptVisible 197 | 198 | self.negative_btn.setStyleSheet(f"background-color: {COLOR.BUTTON if nowPromptVisible else COLOR.BUTTON_SELECTED};") 199 | 200 | def show_position_dialog(self): 201 | """캐릭터 위치 선택 다이얼로그 표시""" 202 | try: 203 | print("위치 선택 다이얼로그 열기 시도") 204 | dialog = PositionSelectorDialog(self) 205 | 206 | # 기존 위치 선택된 상태로 표시 207 | if self.position: 208 | print(f"기존 위치: {self.position}") 209 | col = int(self.position[0] * 5) 210 | row = int(self.position[1] * 5) 211 | index = row * 5 + col 212 | dialog.buttons[index].setStyleSheet(f"background-color: {COLOR.BUTTON_SELECTED};") 213 | dialog.selected_position = self.position 214 | 215 | # 다이얼로그 표시 및 결과 처리 216 | result = dialog.exec_() 217 | print(f"다이얼로그 결과: {result}, 선택된 위치: {dialog.selected_position}") 218 | 219 | if result == QDialog.Accepted and dialog.selected_position: 220 | self.position = dialog.selected_position 221 | self.position_btn.setStyleSheet(f"background-color: {COLOR.BUTTON_SELECTED};") 222 | self.position_btn.setText(_get_position_text(self.position)) 223 | print(f"위치 설정 완료: {self.position}") 224 | except Exception as e: 225 | print(f"위치 선택 다이얼로그 오류: {e}") 226 | 227 | 228 | # return { 229 | # "prompt": "girl, ", 230 | # "uc": "lowres, aliasing, ", 231 | # "center": { 232 | # "x": 0.5, 233 | # "y": 0.5 234 | # }, 235 | # } 236 | def get_data(self): 237 | """캐릭터 프롬프트 데이터 반환""" 238 | return { 239 | "prompt": self.prompt_edit.toPlainText() or "", 240 | "uc": self.neg_prompt_edit.toPlainText() or "", 241 | "center": { 242 | "x": float(self.position[0]) if self.position else 0.5, 243 | "y": float(self.position[1]) if self.position else 0.5 244 | } 245 | } 246 | 247 | def set_data(self, data): 248 | """캐릭터 프롬프트 데이터 설정""" 249 | if "prompt" in data: 250 | self.prompt_edit.setPlainText(data["prompt"]) 251 | 252 | if "uc" in data: 253 | self.neg_prompt_edit.setPlainText(data["uc"]) 254 | 255 | if "center" in data and data["center"]["x"] and data["center"]["y"]: 256 | self.position = (float(data["center"]["x"]), float(data["center"]["y"])) 257 | if self.position: 258 | self.position_btn.setStyleSheet(f"background-color: {COLOR.BUTTON_SELECTED};") 259 | self.position_btn.setText(_get_position_text(self.position)) 260 | 261 | class CharacterPromptsContainer(QWidget): 262 | """캐릭터 프롬프트 컨테이너 위젯""" 263 | 264 | def __init__(self, parent=None): 265 | super().__init__(parent) 266 | self.parent = parent 267 | self.character_widgets = [] 268 | self.not_use_ai_positions = True 269 | self.setup_ui() 270 | 271 | def setup_ui(self): 272 | self.main_layout = QVBoxLayout() 273 | self.main_layout.setContentsMargins(0,0,0,0) 274 | self.setLayout(self.main_layout) 275 | 276 | # 캐릭터 프롬프트 설명 277 | info_layout = QHBoxLayout() 278 | info_label = QLabel("캐릭터 프롬프트 (V4에만 적용)") 279 | info_label.setWordWrap(True) 280 | info_layout.addWidget(info_label) 281 | self.main_layout.addLayout(info_layout) 282 | 283 | # AI 위치 선택 여부 체크박스 284 | controls_layout = QHBoxLayout() 285 | self.ai_position_checkbox = CustomCheckBox("직접 위치 선택하기") 286 | self.ai_position_checkbox.setChecked(False) 287 | self.ai_position_checkbox.onSetCheckedCalled.connect(self.toggle_ai_positions) 288 | self.ai_position_checkbox.stateChanged.connect(self.toggle_ai_positions) 289 | controls_layout.addWidget(self.ai_position_checkbox) 290 | 291 | # 모두 삭제 버튼 292 | self.clear_button = add_button(controls_layout, "모두 삭제", self.clear_characters, 200, 200) 293 | self.clear_button.setMinimumHeight(30) 294 | 295 | # 캐릭터 추가 버튼 296 | add_button(controls_layout, "캐릭터 추가", self.add_character, 200, 200).setMinimumHeight(30) 297 | 298 | self.main_layout.addLayout(controls_layout) 299 | 300 | # 수평 스크롤 영역 301 | self.scroll_area = QScrollArea() 302 | self.scroll_area.setWidgetResizable(True) 303 | self.scroll_area.setFrameShape(QFrame.NoFrame) 304 | self.scroll_area.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff) 305 | self.scroll_area.setVerticalScrollBarPolicy(Qt.ScrollBarAlwaysOff) 306 | self.scroll_area.setStyleSheet(""" 307 | QTextEdit { 308 | background-color: """ + COLOR.DARK + """; 309 | }""") 310 | 311 | # 캐릭터 위젯들을 수평으로 배치할 컨테이너 312 | self.characters_container = QWidget() 313 | self.characters_layout = QHBoxLayout(self.characters_container) 314 | self.characters_layout.setContentsMargins(0, 0, 0, 0) 315 | self.characters_layout.setSpacing(10) # 캐릭터 위젯 사이의 간격 316 | self.characters_layout.addStretch() # 오른쪽 끝에 빈 공간 추가 317 | 318 | self.scroll_area.setWidget(self.characters_container) 319 | 320 | self.main_layout.addWidget(self.scroll_area) 321 | 322 | self.on_update_character_count() 323 | 324 | self.parent.dict_ui_settings["use_coords"] = self.ai_position_checkbox 325 | self.parent.dict_ui_settings["characterPrompts"] = self 326 | 327 | def toggle_ai_positions(self, state): 328 | """AI 위치 선택 여부 토글""" 329 | try: 330 | self.not_use_ai_positions = state == Qt.Checked 331 | 332 | print("toggle_ai_positions", self.not_use_ai_positions) 333 | 334 | # 모든 캐릭터 위젯의 위치 버튼 활성화/비활성화 335 | for widget in self.character_widgets: 336 | enabled = self.not_use_ai_positions 337 | widget.position_btn.setEnabled(enabled) 338 | 339 | # 버튼 스타일 업데이트 (시각적 피드백) 340 | if enabled: 341 | if widget.position: 342 | # 위치가 이미 설정된 경우 343 | widget.position_btn.setStyleSheet(f"background-color: {COLOR.BUTTON_SELECTED};") 344 | widget.position_btn.setText(_get_position_text(widget.position)) 345 | else: 346 | # 위치가 설정되지 않은 경우 347 | widget.position_btn.setStyleSheet(f"background-color: {COLOR.BUTTON};") 348 | widget.position_btn.setText("위치") 349 | else: 350 | # 비활성화된 경우 351 | widget.position_btn.setStyleSheet(f"background-color: {COLOR.BUTTON_DSIABLED};") 352 | widget.position_btn.setText("위치") 353 | except Exception as e: 354 | print(f"AI 위치 토글 오류: {e}") 355 | 356 | def on_update_character_count(self): 357 | """캐릭터 위젯 유무에 따라 ui를 업데이트함""" 358 | is_there_charater = len(self.character_widgets) > 0 359 | 360 | self.setFixedHeight(HEIGHT_CONTAINER_EXPAND if is_there_charater else HEIGHT_CONTAINER_REDUCE) 361 | self.scroll_area.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOn if is_there_charater else Qt.ScrollBarAlwaysOff) 362 | 363 | self.clear_button.setVisible(is_there_charater) 364 | 365 | def add_character(self): 366 | """새 캐릭터 프롬프트 추가""" 367 | if len(self.character_widgets) >= 6: 368 | from PyQt5.QtWidgets import QMessageBox 369 | QMessageBox.warning(self, "경고", "최대 6개의 캐릭터만 추가할 수 있습니다.") 370 | return 371 | 372 | widget = CharacterPromptWidget(self, len(self.character_widgets)) 373 | widget.deleted.connect(self.remove_character) 374 | widget.moved.connect(self.move_character) 375 | widget.position_btn.setEnabled(self.not_use_ai_positions) 376 | 377 | # 켜져있다면, 태그 자동 완성 적용 (로딩 전이라면 다 로딩된 다음에 따로 적용됨.) 378 | if self.parent.completiontag_list: 379 | widget.prompt_edit.start_complete_mode(self.parent.completiontag_list) 380 | widget.neg_prompt_edit.start_complete_mode(self.parent.completiontag_list) 381 | 382 | # stretch 아이템 앞에 위젯 삽입 383 | self.characters_layout.insertWidget(self.characters_layout.count() - 1, widget) 384 | self.character_widgets.append(widget) 385 | 386 | # 인덱스 업데이트 387 | self.update_indices() 388 | 389 | self.on_update_character_count() 390 | 391 | def remove_character(self, widget): 392 | """캐릭터 프롬프트 삭제""" 393 | if widget in self.character_widgets: 394 | self.characters_layout.removeWidget(widget) 395 | self.character_widgets.remove(widget) 396 | widget.deleteLater() 397 | 398 | # 인덱스 업데이트 399 | self.update_indices() 400 | 401 | self.on_update_character_count() 402 | 403 | def move_character(self, widget, direction): 404 | """캐릭터 프롬프트 순서 이동""" 405 | if widget not in self.character_widgets: 406 | return 407 | 408 | index = self.character_widgets.index(widget) 409 | new_index = index + direction 410 | 411 | if 0 <= new_index < len(self.character_widgets): 412 | # 위젯 순서 변경 413 | self.character_widgets.pop(index) 414 | self.character_widgets.insert(new_index, widget) 415 | 416 | # 레이아웃에서 제거 후 재배치 417 | for i, w in enumerate(self.character_widgets): 418 | self.characters_layout.removeWidget(w) 419 | 420 | # stretch 아이템 제거 421 | stretch_item = self.characters_layout.takeAt(self.characters_layout.count() - 1) 422 | 423 | # 위젯 추가 424 | for i, w in enumerate(self.character_widgets): 425 | self.characters_layout.addWidget(w) 426 | 427 | # stretch 아이템 다시 추가 428 | self.characters_layout.addStretch() 429 | 430 | # 인덱스 업데이트 431 | self.update_indices() 432 | 433 | def update_indices(self): 434 | """캐릭터 인덱스 업데이트""" 435 | for i, widget in enumerate(self.character_widgets): 436 | widget.index = i 437 | try: 438 | widget.update_title() 439 | except: 440 | # 예전 방식으로도 시도 441 | try: 442 | header_layout = widget.layout.itemAt(0).layout() 443 | if header_layout: 444 | title_label = header_layout.itemAt(0).widget() 445 | if title_label: 446 | title_label.setText(f"캐릭터 {i + 1}") 447 | except: 448 | pass 449 | 450 | def clear_characters(self): 451 | """모든 캐릭터 프롬프트 삭제""" 452 | for widget in self.character_widgets[:]: 453 | self.remove_character(widget) 454 | 455 | self.on_update_character_count() 456 | 457 | # 로딩 전에 생성된 오브젝트에게 list를 삽입한다. 458 | def refresh_completiontag_list(self): 459 | for char_widget in self.character_widgets: 460 | char_widget.prompt_edit.start_complete_mode(self.parent.completiontag_list) 461 | char_widget.neg_prompt_edit.start_complete_mode(self.parent.completiontag_list) 462 | 463 | # return [{ 464 | # "prompt": "girl, ", 465 | # "uc": "lowres, aliasing, ", 466 | # "center": { 467 | # "x": 0.5, 468 | # "y": 0.5 469 | # }, 470 | # }] 471 | def get_data(self): 472 | """모든 캐릭터 프롬프트 데이터 반환""" 473 | return [widget.get_data() for widget in self.character_widgets] 474 | 475 | # data = [{ 476 | # "prompt": "girl, ", 477 | # "uc": "lowres, aliasing, ", 478 | # "center": { 479 | # "x": 0.5, 480 | # "y": 0.5 481 | # }, 482 | # }] 483 | def set_data(self, data): 484 | """캐릭터 프롬프트 데이터 설정""" 485 | self.clear_characters() 486 | 487 | for char_data in data: 488 | self.add_character() 489 | self.character_widgets[-1].set_data(char_data) -------------------------------------------------------------------------------- /app/main_window.py: -------------------------------------------------------------------------------- 1 | import json 2 | import sys 3 | import os 4 | import time 5 | import random 6 | from io import BytesIO 7 | from PIL import Image 8 | from urllib import request 9 | 10 | from PyQt5.QtWidgets import QApplication, QMainWindow, QFileDialog, QMessageBox, QDialog 11 | from PyQt5.QtCore import QSettings, QPoint, QSize, QCoreApplication, QTimer 12 | 13 | from core.thread.generate_thread import GenerateThread 14 | from core.thread.token_thread import TokenValidateThread 15 | from core.thread.anlas_thread import AnlasThread 16 | from core.thread.completiontagload_thread import CompletionTagLoadThread 17 | 18 | from gui.layout.main_layout import init_main_layout 19 | from gui.widget.status_bar import StatusBar 20 | from gui.widget.menu_bar import MenuBar 21 | 22 | from gui.dialog.generate_dialog import GenerateDialog 23 | from gui.dialog.miniutil_dialog import MiniUtilDialog 24 | from gui.dialog.fileio_dialog import FileIODialog 25 | from gui.dialog.login_dialog import LoginDialog 26 | from gui.dialog.option_dialog import OptionDialog 27 | from gui.dialog.etc_dialog import show_setting_load_dialog, show_setting_save_dialog 28 | 29 | from util.common_util import strtobool, try_dumps, try_loads 30 | from util.string_util import apply_wc_and_lessthan, prettify_naidict 31 | from util.file_util import create_folder_if_not_exists 32 | from util.ui_util import set_opacity 33 | 34 | from core.worker.naiinfo_getter import get_naidict_from_file, get_naidict_from_txt, get_naidict_from_img 35 | from core.worker.nai_generator import NAIGenerator, get_character_prompts_from_v4_prompt, TARGET_PARAMETERS, SAMPLER_ITEMS_V4, SAMPLER_ITEMS_V3, MODEL_INFO_DICT, DEFAULT_MODEL_V4 36 | from core.worker.wildcard_applier import WildcardApplier 37 | from core.worker.danbooru_tagger import DanbooruTagger 38 | 39 | from config.strings import STRING 40 | from config.consts import RESOLUTION_FAMILIY_MASK, RESOLUTION_FAMILIY, TITLE_NAME, TOP_NAME, APP_NAME 41 | from config.paths import DEFAULT_PATH 42 | 43 | 44 | class NAIAutoGeneratorWindow(QMainWindow): 45 | def __init__(self, app): 46 | super().__init__() 47 | self.app = app 48 | 49 | self.init_variable() 50 | self.init_window() 51 | self.init_statusbar() 52 | self.init_menubar() 53 | self.init_content() 54 | self.check_folders() 55 | self.show() 56 | 57 | self.init_nai() 58 | self.init_wc() 59 | self.init_tagger() 60 | self.init_completion() 61 | 62 | self.load_data() 63 | 64 | def init_variable(self): 65 | self.is_expand = False 66 | self.trying_auto_login = False 67 | self.autogenerate_thread = None 68 | self.list_settings_batch_target = [] 69 | self.index_settings_batch_target = -1 70 | 71 | def init_window(self): 72 | self.setWindowTitle(TITLE_NAME) 73 | self.setAcceptDrops(True) 74 | 75 | self.settings = QSettings(TOP_NAME, APP_NAME) 76 | self.move(self.settings.value("pos", QPoint(500, 200))) 77 | self.resize(self.settings.value("size", QSize(1179, 1044))) 78 | self.settings.setValue("splitterSizes", None) 79 | font_size = self.settings.value("nag_font_size", 18) 80 | self.app.setStyleSheet("QWidget{font-size:" + str(font_size) + "px}") 81 | 82 | def init_statusbar(self): 83 | self.statusbar = StatusBar(self.statusBar()) 84 | self.statusbar.set_statusbar_text("BEFORE_LOGIN") 85 | 86 | def init_menubar(self): 87 | self.menubar = MenuBar(self) 88 | 89 | def init_content(self): 90 | init_main_layout(self) 91 | 92 | def init_nai(self): 93 | self.nai = NAIGenerator() 94 | 95 | if self.settings.value("auto_login", False): 96 | access_token = self.settings.value("access_token", "") 97 | username = self.settings.value("username", "") 98 | password = self.settings.value("password", "") 99 | if not access_token or not username or not password: 100 | return 101 | 102 | self.statusbar.set_statusbar_text("LOGGINGIN") 103 | self.nai.access_token = access_token 104 | self.nai.username = username 105 | self.nai.password = password 106 | 107 | self.trying_auto_login = True 108 | validate_thread = TokenValidateThread(self) 109 | validate_thread.validation_result.connect(self.on_login_result) 110 | validate_thread.start() 111 | 112 | def init_wc(self): 113 | self.wcapplier = WildcardApplier(self.settings.value( 114 | "path_wildcards", os.path.abspath(DEFAULT_PATH["path_wildcards"]))) 115 | 116 | def init_tagger(self): 117 | self.dtagger = DanbooruTagger(self.settings.value( 118 | "path_models", os.path.abspath(DEFAULT_PATH["path_models"]))) 119 | 120 | def init_completion(self): 121 | self.completiontag_list = None 122 | if strtobool(self.settings.value("will_complete_tag", True)): 123 | CompletionTagLoadThread(self).start() 124 | 125 | # 저장과 불러오기는 항상 str으로 한다. 126 | def save_data(self): 127 | data_dict = self.get_data_from_savetarget_ui() 128 | 129 | data_dict["seed_fix_checkbox"] = self.dict_ui_settings["seed_fix_checkbox"].isChecked() 130 | 131 | for k, v in data_dict.items(): 132 | if k == "characterPrompts": 133 | self.settings.setValue(k, try_dumps(v) or TARGET_PARAMETERS[k]) 134 | else: 135 | self.settings.setValue(k, str(v)) 136 | 137 | # 저장과 불러오기는 항상 str으로 한다. 138 | def load_data(self): 139 | data_dict = {} 140 | for key in TARGET_PARAMETERS: 141 | if key == "characterPrompts": 142 | data_dict["characterPrompts"] = self.settings.value("characterPrompts", TARGET_PARAMETERS[key]) 143 | if isinstance(data_dict["characterPrompts"], str): 144 | data_dict["characterPrompts"] = try_loads(data_dict["characterPrompts"]) or TARGET_PARAMETERS[key] 145 | else: 146 | data_dict[key] = self.settings.value(key, TARGET_PARAMETERS[key]) 147 | 148 | self.set_data(data_dict) 149 | 150 | def set_data(self, data_dict): 151 | ui_dict = self.dict_ui_settings 152 | 153 | # Combo : setCurrentText 154 | list_ui_using_setcurrenttext = ["model", "sampler"] 155 | for key in list_ui_using_setcurrenttext: 156 | if key in data_dict: 157 | ui_dict[key].setCurrentText(str(data_dict[key])) 158 | else: 159 | print("[set_data] 없는 키가 있습니다. 키 : ",key) 160 | 161 | # 혹시 몰라 model ui 세팅을 매뉴얼로 호출해줌. 162 | if "model" in data_dict: 163 | self.on_model_changed(data_dict["model"]) 164 | 165 | # EditText : setText 166 | list_ui_using_settext = ["prompt", "negative_prompt", "width", "height", 167 | "steps", "seed", "scale", "cfg_rescale"] 168 | for key in list_ui_using_settext: 169 | if key in data_dict: 170 | ui_dict[key].setText(str(data_dict[key])) 171 | else: 172 | print("[set_data] 없는 키가 있습니다. 키 : ",key) 173 | 174 | # Check : setChecked 175 | list_ui_using_setchecked = ["sm", "sm_dyn", "use_coords", "variety_plus"] 176 | for key in list_ui_using_setchecked: 177 | if key in data_dict: 178 | ui_dict[key].setChecked(strtobool(data_dict[key])) 179 | else: 180 | print("[set_data] 없는 키가 있습니다. 키 : ",key) 181 | 182 | # char : characterPrompts 얻어오기 183 | characterPrompts = None 184 | is_forced_use_coords = None 185 | if "characterPrompts" in data_dict and data_dict["characterPrompts"]: 186 | characterPrompts = data_dict["characterPrompts"] 187 | 188 | #str이면 json으로 변경해보려는 노력을 한다. (load도 이리로옴) 189 | if isinstance(characterPrompts, str): 190 | try: 191 | characterPrompts = json.loads(characterPrompts) 192 | except Exception: 193 | characterPrompts = None 194 | elif "v4_prompt" in data_dict and data_dict["v4_prompt"] and "v4_negative_prompt" in data_dict and data_dict["v4_negative_prompt"]: 195 | # 이미지에서 불러온 경우, characterPrompts, use_coords가 없고 v4_prompt만 있는 경우가 있다. 196 | 197 | # 여기서 컨버트하여 적용함. 198 | characterPrompts, use_coords = get_character_prompts_from_v4_prompt(data_dict) 199 | 200 | if characterPrompts is not None: 201 | # v4가 제대로 전환된경우 이미지에 model 지정이 없었다면 자동으로 v4 모델로 전환한다. 202 | if "model" not in data_dict: 203 | ui_dict["model"].setCurrentText(DEFAULT_MODEL_V4) 204 | self.on_model_changed(MODEL_INFO_DICT[DEFAULT_MODEL_V4]["model"]) 205 | 206 | if use_coords is not None: 207 | is_forced_use_coords = use_coords 208 | 209 | # char : 최종 적용 210 | if characterPrompts: 211 | ui_dict["characterPrompts"].set_data(characterPrompts) 212 | else: 213 | ui_dict["characterPrompts"].clear_characters() 214 | 215 | # char : use_coords를 set_data 보다 뒤에서 적용해야함 216 | if is_forced_use_coords is not None: 217 | print("is_forced_use_coords", is_forced_use_coords) 218 | ui_dict["use_coords"].setChecked(is_forced_use_coords) 219 | 220 | 221 | def check_folders(self): 222 | for key, default_path in DEFAULT_PATH.items(): 223 | path = self.settings.value(key, os.path.abspath(default_path)) 224 | create_folder_if_not_exists(path) 225 | 226 | # 저장 대상인 항목만 이곳에 담는다. 227 | def get_data_from_savetarget_ui(self): 228 | data = { 229 | "model": self.dict_ui_settings["model"].currentText(), 230 | "prompt": self.dict_ui_settings["prompt"].toPlainText(), 231 | "negative_prompt": self.dict_ui_settings["negative_prompt"].toPlainText(), 232 | "width": int(self.dict_ui_settings["width"].text()), 233 | "height": int(self.dict_ui_settings["height"].text()), 234 | "sampler": self.dict_ui_settings["sampler"].currentText(), 235 | "steps": int(self.dict_ui_settings["steps"].text()), 236 | "seed": int(self.dict_ui_settings["seed"].text()), 237 | "scale": float(self.dict_ui_settings["scale"].text()), 238 | "cfg_rescale": float(self.dict_ui_settings["cfg_rescale"].text()), 239 | "sm": self.dict_ui_settings["sm"].isChecked(), 240 | "sm_dyn": self.dict_ui_settings["sm_dyn"].isChecked(), 241 | "variety_plus": self.dict_ui_settings["variety_plus"].isChecked(), 242 | "use_coords": self.dict_ui_settings["use_coords"].isChecked(), 243 | "characterPrompts": self.dict_ui_settings["characterPrompts"].get_data() 244 | } 245 | 246 | return data 247 | 248 | # Warning! Don't interact with pyqt gui in this function 249 | # ui에서 데이터를 가져온 다음에 생성전 수정을 가한다. 250 | def _get_data_for_generate(self): 251 | data = self.get_data_from_savetarget_ui() 252 | self.save_data() 253 | 254 | # data precheck 255 | self.check_folders() 256 | data["prompt"] = apply_wc_and_lessthan(self.wcapplier, data["prompt"]) 257 | data["negative_prompt"] = apply_wc_and_lessthan(self.wcapplier, data["negative_prompt"]) 258 | 259 | if "characterPrompts" in data and data["characterPrompts"]: 260 | for char_dict in data["characterPrompts"]: 261 | if "prompt" in char_dict and char_dict["prompt"]: 262 | char_dict["prompt"] = apply_wc_and_lessthan(self.wcapplier, char_dict["prompt"]) 263 | if "uc" in char_dict and char_dict["uc"]: 264 | char_dict["uc"] = apply_wc_and_lessthan(self.wcapplier, char_dict["uc"]) 265 | 266 | print(data) 267 | 268 | # seed pick 269 | if not self.dict_ui_settings["seed_fix_checkbox"].isChecked() or data["seed"] == -1: 270 | data["seed"] = random.randint(0, 9999999999) 271 | 272 | # resolution pick 273 | if strtobool(self.checkbox_random_resolution.isChecked()): 274 | fl = self.get_now_resolution_familly_list() 275 | if fl: 276 | text = fl[random.randrange(0, len(fl))] 277 | 278 | res_text = text.split("(")[1].split(")")[0] 279 | width, height = res_text.split("x") 280 | data["width"], data["height"] = int(width), int(height) 281 | 282 | # image option check 283 | i2i_param = self.i2i_settings_group.get_nai_param() 284 | vibe_param = self.vibe_settings_group.get_nai_param() 285 | if i2i_param: 286 | data["image"] = i2i_param[0] 287 | data["strength"] = float(i2i_param[1]) 288 | data["noise"] = float(i2i_param[2]) 289 | 290 | # mask 체크 291 | # if self.i2i_settings_group.mask: 292 | # data['mask'] = convert_qimage_to_imagedata( 293 | # self.i2i_settings_group.mask) 294 | if vibe_param: 295 | image_tuple, info_tuple, strength_tuple = zip(*vibe_param) 296 | data["reference_image_multiple"]= list(image_tuple) 297 | data["reference_information_extracted_multiple"]= [float(x) for x in info_tuple] 298 | data["reference_strength_multiple"]= [float(x) for x in strength_tuple] 299 | 300 | return data 301 | 302 | def on_model_changed(self, text): 303 | if text in MODEL_INFO_DICT: 304 | model_info = MODEL_INFO_DICT[text] 305 | elif text in [info["model"] for info in MODEL_INFO_DICT.values()]: 306 | model_info = next((info for info in MODEL_INFO_DICT.values() if info["model"] == text), None) 307 | if not model_info: 308 | model_info = MODEL_INFO_DICT[DEFAULT_MODEL_V4] 309 | self.dict_ui_settings["model"].setCurrentText(DEFAULT_MODEL_V4) 310 | QMessageBox.information(self, '경고', "모델 정보를 찾을 수 없습니다. 기본 모델로 설정합니다.") 311 | 312 | # sampler 313 | sampler_ui = self.dict_ui_settings["sampler"] 314 | sampler_prevText = sampler_ui.currentText() 315 | 316 | sampler_ui.clear() 317 | targetItems = model_info["sampler"] 318 | sampler_ui.addItems(targetItems) 319 | 320 | if sampler_prevText in targetItems: 321 | sampler_ui.setCurrentText(sampler_prevText) 322 | 323 | # char off 324 | self.character_prompts_container.setVisible(model_info["characterPrompts"]) 325 | 326 | # i2i off 327 | self.i2i_settings_group.setVisible(model_info["i2i"]) 328 | 329 | # vibe off 330 | self.vibe_settings_group.setVisible(model_info["vibe"]) 331 | 332 | # settings off 333 | set_opacity(self.dict_ui_settings["sm"], model_info["sm"]) 334 | set_opacity(self.dict_ui_settings["sm_dyn"], model_info["sm_dyn"]) 335 | set_opacity(self.dict_ui_settings["variety_plus"], model_info["variety_plus"]) 336 | 337 | def _on_after_create_data_apply_gui(self): 338 | data = self.nai.parameters 339 | 340 | # resolution text 341 | fl = self.get_now_resolution_familly_list() 342 | if fl: 343 | for resol in fl: 344 | if str(data["width"]) + "x" + str(data["height"]) in resol: 345 | self.combo_resolution.setCurrentText(resol) 346 | break 347 | 348 | # seed text 349 | self.dict_ui_settings["seed"].setText(str(data["seed"])) 350 | 351 | # result text 352 | self.prompt_result.setText(prettify_naidict(data)) 353 | 354 | def on_click_generate_once(self): 355 | self.list_settings_batch_target = [] 356 | 357 | data = self._get_data_for_generate() 358 | self.nai.set_param_dict(data) 359 | self._on_after_create_data_apply_gui() 360 | 361 | generate_thread = GenerateThread(self, False) 362 | generate_thread.generate_result.connect(self._on_result_generate) 363 | generate_thread.start() 364 | 365 | self.statusbar.set_statusbar_text("GENEARTING") 366 | self.generate_buttons_layout.set_disable_button(True) 367 | self.generate_thread = generate_thread 368 | 369 | def _on_result_generate(self, error_code, result): 370 | self.generate_thread = None 371 | self.generate_buttons_layout.set_disable_button(False) 372 | self.statusbar.set_statusbar_text("IDLE") 373 | self.refresh_anlas() 374 | 375 | if error_code == 0: 376 | self.image_result.set_custom_pixmap(result) 377 | else: 378 | QMessageBox.information( 379 | self, '경고', "이미지를 생성하는데 문제가 있습니다.\n\n" + str(result)) 380 | 381 | def on_click_generate_sett(self): 382 | path_list, _ = QFileDialog().getOpenFileNames(self, 383 | caption="불러올 세팅 파일들을 선택해주세요", 384 | filter="Txt File (*.txt)") 385 | if path_list: 386 | if len(path_list) < 2: 387 | QMessageBox.information( 388 | self, '경고', "두개 이상 선택해주세요.") 389 | return 390 | 391 | for path in path_list: 392 | if not path.endswith(".txt") or not os.path.isfile(path): 393 | QMessageBox.information( 394 | self, '경고', ".txt로 된 세팅 파일만 선택해주세요.") 395 | return 396 | 397 | self.on_click_generate_auto(path_list) 398 | 399 | def proceed_settings_batch(self): 400 | self.index_settings_batch_target += 1 401 | 402 | while len(self.list_settings_batch_target) <= self.index_settings_batch_target: 403 | self.index_settings_batch_target -= len( 404 | self.list_settings_batch_target) 405 | 406 | path = self.list_settings_batch_target[self.index_settings_batch_target] 407 | is_success = self._load_settings(path) 408 | 409 | return is_success 410 | 411 | def on_click_generate_auto(self, setting_batch_target=[]): 412 | if not self.autogenerate_thread: 413 | d = GenerateDialog(self) 414 | if d.exec_() == QDialog.Accepted: 415 | self.list_settings_batch_target = setting_batch_target 416 | if setting_batch_target: 417 | self.index_settings_batch_target = -1 418 | is_success = self.proceed_settings_batch() 419 | if not is_success: 420 | QMessageBox.information( 421 | self, '경고', "세팅을 불러오는데 실패했습니다.") 422 | return 423 | 424 | agt = GenerateThread( 425 | self, True, d.count, d.delay, d.ignore_error) 426 | agt.on_data_created.connect( 427 | self._on_after_create_data_apply_gui) 428 | agt.on_error.connect(self._on_error_autogenerate) 429 | agt.on_end.connect(self._on_end_autogenerate) 430 | agt.on_statusbar_change.connect( 431 | self.statusbar.set_statusbar_text) 432 | agt.on_success.connect(self._on_success_autogenerate) 433 | agt.start() 434 | 435 | self.generate_buttons_layout.set_autogenerate_mode(True) 436 | self.autogenerate_thread = agt 437 | else: 438 | self._on_end_autogenerate() 439 | 440 | def _on_error_autogenerate(self, error_code, result): 441 | QMessageBox.information( 442 | self, '경고', "이미지를 생성하는데 문제가 있습니다.\n\n" + str(result)) 443 | self._on_end_autogenerate() 444 | 445 | def _on_end_autogenerate(self): 446 | self.autogenerate_thread.stop() 447 | self.autogenerate_thread = None 448 | self.generate_buttons_layout.set_autogenerate_mode(False) 449 | self.statusbar.set_statusbar_text("IDLE") 450 | self.refresh_anlas() 451 | 452 | def _on_success_autogenerate(self, result_str): 453 | self._on_refresh_anlas(self.nai.get_anlas() or -1) 454 | 455 | self.image_result.set_custom_pixmap(result_str) 456 | 457 | if self.list_settings_batch_target: 458 | self.proceed_settings_batch() 459 | 460 | def on_click_open_folder(self, target_pathcode): 461 | path = self.settings.value( 462 | target_pathcode, DEFAULT_PATH[target_pathcode]) 463 | path = os.path.abspath(path) 464 | create_folder_if_not_exists(path) 465 | os.startfile(path) 466 | 467 | def on_click_save_settings(self): 468 | show_setting_save_dialog(self) 469 | 470 | def on_click_load_settings(self): 471 | show_setting_load_dialog(self) 472 | 473 | def _load_settings(self, path): 474 | try: 475 | with open(path, "r", encoding="utf8") as f: 476 | json_str = f.read() 477 | json_obj = json.loads(json_str) 478 | 479 | self.set_data(json_obj) 480 | 481 | return True 482 | except Exception as e: 483 | print(e) 484 | 485 | return False 486 | 487 | def on_random_resolution_checked(self, is_checked): 488 | if is_checked == 2: 489 | fl = self.get_now_resolution_familly_list() 490 | if not fl: 491 | QMessageBox.information( 492 | self, '이미지 크기 랜덤', "랜덤이 지원되지 않는 형식입니다.\n") 493 | else: 494 | s = "" 495 | for f in fl: 496 | s += f + "\n" 497 | QMessageBox.information( 498 | self, '이미지 크기 랜덤', "다음 크기 중 하나가 랜덤으로 선택됩니다.\n\n" + s) 499 | 500 | self.settings.setValue("image_random_checkbox", is_checked == 2) 501 | 502 | def get_now_resolution_familly_list(self): 503 | family_mask = RESOLUTION_FAMILIY_MASK[self.combo_resolution.currentIndex( 504 | )] 505 | 506 | if family_mask == -1: 507 | return [] 508 | 509 | return RESOLUTION_FAMILIY[family_mask] 510 | 511 | def on_change_path(self, code, src): 512 | path = os.path.abspath(src) 513 | 514 | self.settings.setValue(code, path) 515 | 516 | create_folder_if_not_exists(path) 517 | 518 | if code == "path_wildcards": 519 | self.init_wc() 520 | elif code == "path_models": 521 | self.init_tagger() 522 | 523 | def on_click_getter(self): 524 | MiniUtilDialog(self, "getter").show() 525 | 526 | def on_click_tagger(self): 527 | MiniUtilDialog(self, "tagger").show() 528 | 529 | def on_click_expand(self): 530 | if self.is_expand: 531 | self.is_expand = False 532 | self.button_expand.setText("<<") 533 | 534 | self.settings.setValue( 535 | "splitterSizes", self.main_splitter.saveState()) 536 | self.main_splitter.setHandleWidth(0) 537 | self.main_splitter.widget(1).setMaximumSize(0, 0) 538 | else: 539 | self.is_expand = True 540 | self.button_expand.setText(">>") 541 | 542 | self.main_splitter.setHandleWidth(8) 543 | self.main_splitter.widget(1).setMaximumSize(16777215, 16777215) 544 | 545 | try: 546 | self.main_splitter.restoreState( 547 | self.settings.value("splitterSizes")) 548 | except Exception as e: 549 | self.main_splitter.setSizes([16777215, 16777215]) 550 | QTimer.singleShot(20, self.image_result.refresh_size) 551 | 552 | def install_model(self, model_name): 553 | loading_dialog = FileIODialog( 554 | "모델 다운 받는 중...\n이 작업은 오래 걸릴 수 있습니다.", lambda: str(self.dtagger.download_model(model_name))) 555 | if loading_dialog.exec_() == QDialog.Accepted: 556 | if loading_dialog.result == "True": 557 | self.option_dialog.on_model_downloaded(model_name) 558 | 559 | def get_image_info_bysrc(self, file_src): 560 | nai_dict, error_code = get_naidict_from_file(file_src) 561 | 562 | self._get_image_info_byinfo(nai_dict, error_code, file_src) 563 | 564 | def get_image_info_bytxt(self, file_src): 565 | nai_dict, error_code = get_naidict_from_txt(file_src) 566 | 567 | self._get_image_info_byinfo(nai_dict, error_code, None) 568 | 569 | def get_image_info_byimg(self, img): 570 | nai_dict, error_code = get_naidict_from_img(img) 571 | 572 | self._get_image_info_byinfo(nai_dict, error_code, img) 573 | 574 | def _get_image_info_byinfo(self, nai_dict, error_code, img_obj): 575 | if error_code == 1: 576 | QMessageBox.information(self, '경고', "EXIF가 존재하지 않는 파일입니다.") 577 | self.statusbar.set_statusbar_text("IDLE") 578 | elif error_code == 2 or error_code == 3: 579 | QMessageBox.information( 580 | self, '경고', "EXIF는 존재하나 NAI로부터 만들어진 것이 아닌 듯 합니다.") 581 | self.statusbar.set_statusbar_text("IDLE") 582 | elif error_code == 0: 583 | new_dict = { 584 | "prompt": nai_dict["prompt"], "negative_prompt": nai_dict["negative_prompt"]} 585 | new_dict.update(nai_dict["option"]) 586 | new_dict.update(nai_dict["etc"]) 587 | 588 | self.set_data(new_dict) 589 | if img_obj: 590 | self.image_result.set_custom_pixmap(img_obj) 591 | self.statusbar.set_statusbar_text("LOAD_COMPLETE") 592 | 593 | def show_login_dialog(self): 594 | self.login_dialog = LoginDialog(self) 595 | self.login_dialog.exec_() 596 | 597 | def show_option_dialog(self): 598 | self.option_dialog = OptionDialog(self) 599 | self.option_dialog.exec_() 600 | 601 | def show_about_dialog(self): 602 | QMessageBox.about(self, 'About', STRING.ABOUT) 603 | 604 | def refresh_anlas(self): 605 | anlas_thread = AnlasThread(self) 606 | anlas_thread.anlas_result.connect(self._on_refresh_anlas) 607 | anlas_thread.start() 608 | 609 | def _on_refresh_anlas(self, anlas): 610 | if anlas == -1: 611 | anlas = "?" 612 | self.label_anlas.setText("Anlas: " + str(anlas)) 613 | 614 | def on_login_result(self, error_code): 615 | if error_code == 0: 616 | self.statusbar.set_statusbar_text("LOGINED") 617 | self.label_loginstate.set_logged_in(True) 618 | self.generate_buttons_layout.set_disable_button(False) 619 | self.refresh_anlas() 620 | else: 621 | self.nai = NAIGenerator() # reset 622 | self.statusbar.set_statusbar_text("BEFORE_LOGIN") 623 | self.label_loginstate.set_logged_in(False) 624 | self.generate_buttons_layout.set_disable_button(True) 625 | self.set_auto_login(False) 626 | 627 | self.trying_auto_login = False 628 | 629 | def set_auto_login(self, is_auto_login): 630 | self.settings.setValue("auto_login", 631 | True if is_auto_login else False) 632 | self.settings.setValue("access_token", 633 | self.nai.access_token if is_auto_login else None) 634 | self.settings.setValue("username", 635 | self.nai.username if is_auto_login else None) 636 | self.settings.setValue("password", 637 | self.nai.password if is_auto_login else None) 638 | 639 | def on_logout(self): 640 | self.statusbar.set_statusbar_text("BEFORE_LOGIN") 641 | 642 | self.label_loginstate.set_logged_in(False) 643 | 644 | self.generate_buttons_layout.set_disable_button(True) 645 | 646 | self.set_auto_login(False) 647 | 648 | def dragEnterEvent(self, event): 649 | if event.mimeData().hasUrls(): 650 | event.accept() 651 | else: 652 | event.ignore() 653 | 654 | def dropEvent(self, event): 655 | files = [u for u in event.mimeData().urls()] 656 | 657 | if len(files) != 1: 658 | QMessageBox.information(self, '경고', "파일을 하나만 옮겨주세요.") 659 | return 660 | 661 | furl = files[0] 662 | if furl.isLocalFile(): 663 | fname = furl.toLocalFile() 664 | if fname.endswith(".png") or fname.endswith(".webp"): 665 | if not fname.endswith(".jpg"): 666 | self.get_image_info_bysrc(fname) 667 | return 668 | elif fname.endswith(".txt"): 669 | self.get_image_info_bytxt(fname) 670 | return 671 | 672 | QMessageBox.information( 673 | self, '경고', "세팅 불러오기는 png, webp, txt 파일만 가능합니다.") 674 | else: 675 | self.statusbar.set_statusbar_text("LOADING") 676 | try: 677 | url = furl.url() 678 | res = request.urlopen(url).read() 679 | img = Image.open(BytesIO(res)) 680 | if img: 681 | self.get_image_info_byimg(img) 682 | 683 | except Exception as e: # 바이브 이미지 684 | print(e) 685 | self.statusbar.set_statusbar_text("IDLE") 686 | QMessageBox.information(self, '경고', "이미지 파일 다운로드에 실패했습니다.") 687 | return 688 | 689 | def closeEvent(self, e): 690 | size = self.size() 691 | size.setWidth( 692 | int(size.width() / 2 if self.is_expand else size.width())) 693 | self.settings.setValue("size", size) 694 | self.settings.setValue("pos", self.pos()) 695 | self.save_data() 696 | e.accept() 697 | 698 | def quit_app(self): 699 | time.sleep(0.1) 700 | self.close() 701 | self.app.closeAllWindows() 702 | QCoreApplication.exit(0) 703 | 704 | 705 | if __name__ == '__main__': 706 | input_list = sys.argv 707 | app = QApplication(sys.argv) 708 | 709 | widget = NAIAutoGeneratorWindow(app) 710 | 711 | time.sleep(0.1) 712 | 713 | sys.exit(app.exec_()) 714 | --------------------------------------------------------------------------------