diff --git a/main.py b/main.py new file mode 100644 index 0000000..e47fd40 --- /dev/null +++ b/main.py @@ -0,0 +1,435 @@ +import sys +import json +import base64 +import requests +from time import sleep +from PyQt6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, + QHBoxLayout, QLabel, QLineEdit, QPushButton, + QComboBox, QSpinBox, QDoubleSpinBox, QProgressBar, + QScrollArea, QTextEdit, QFileDialog) +from PyQt6.QtCore import Qt, QThread, pyqtSignal, QTimer +from PyQt6.QtGui import QPixmap, QImage +from PIL import Image +from io import BytesIO + +# Импортируем новые модули +from modules.console_window import ConsoleWindow +from modules.model_manager import ModelManager +from modules.sampler_manager import SamplerManager + +class GenerationThread(QThread): + progress_updated = pyqtSignal(float, str) + generation_complete = pyqtSignal(list) + error_occurred = pyqtSignal(str) + + def __init__(self, api_url, params): + super().__init__() + self.api_url = api_url + self.params = params + self.is_generating = False + + def run(self): + try: + self.is_generating = True + self.progress_updated.emit(0, "Начало генерации...") + + # Отправляем запрос на генерацию + response = requests.post(f"{self.api_url}/sdapi/v1/txt2img", json=self.params) + + if response.status_code != 200: + self.error_occurred.emit(f"Ошибка API: {response.status_code}") + return + + # Получаем результат + result = response.json() + images = result.get('images', []) + if not images: + self.error_occurred.emit("Не удалось получить изображения") + return + + self.generation_complete.emit(images) + + except Exception as e: + self.error_occurred.emit(f"Ошибка: {str(e)}") + finally: + self.is_generating = False + + def stop(self): + self.is_generating = False + +class MainWindow(QMainWindow): + def __init__(self): + super().__init__() + self.api_url = "http://localhost:7860" # URL по умолчанию + self.console_window = None + self.progress_timer = None + self.init_ui() + + def init_ui(self): + self.setWindowTitle("Stable Diffusion UI") + self.setStyleSheet(""" + QMainWindow { + background-color: #2b2b2b; + } + QWidget { + color: #ffffff; + font-size: 14px; + } + QLineEdit, QTextEdit, QComboBox, QSpinBox, QDoubleSpinBox { + background-color: #3b3b3b; + border: 1px solid #555555; + border-radius: 4px; + padding: 5px; + color: #ffffff; + min-width: 100px; + } + QSpinBox::up-button, QDoubleSpinBox::up-button { + background-color: #4a4a4a; + border: none; + border-left: 1px solid #555555; + border-bottom: 1px solid #555555; + border-top-right-radius: 4px; + width: 20px; + height: 12px; + } + QSpinBox::down-button, QDoubleSpinBox::down-button { + background-color: #4a4a4a; + border: none; + border-left: 1px solid #555555; + border-bottom-right-radius: 4px; + width: 20px; + height: 12px; + } + QSpinBox::up-button:hover, QDoubleSpinBox::up-button:hover, + QSpinBox::down-button:hover, QDoubleSpinBox::down-button:hover { + background-color: #5a5a5a; + } + QSpinBox::up-button:pressed, QDoubleSpinBox::up-button:pressed, + QSpinBox::down-button:pressed, QDoubleSpinBox::down-button:pressed { + background-color: #666666; + } + QSpinBox::up-arrow, QDoubleSpinBox::up-arrow { + image: url(data:image/svg+xml;base64,PHN2ZyB3aWR0aD0iMTYiIGhlaWdodD0iMTYiIHZpZXdCb3g9IjAgMCAxNiAxNiIgZmlsbD0ibm9uZSIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIj48cGF0aCBkPSJNOCA0TDQgOEgxMkw4IDRaIiBmaWxsPSIjZmZmZmZmIi8+PC9zdmc+); + width: 12px; + height: 6px; + } + QSpinBox::down-arrow, QDoubleSpinBox::down-arrow { + image: url(data:image/svg+xml;base64,PHN2ZyB3aWR0aD0iMTYiIGhlaWdodD0iMTYiIHZpZXdCb3g9IjAgMCAxNiAxNiIgZmlsbD0ibm9uZSIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIj48cGF0aCBkPSJNOCAxMkwxMiA4SDRMOCAxMloiIGZpbGw9IiNmZmZmZmYiLz48L3N2Zz4=); + width: 12px; + height: 6px; + } + QPushButton { + background-color: #0d6efd; + border: none; + border-radius: 4px; + padding: 8px 16px; + color: white; + font-weight: bold; + } + QPushButton:hover { + background-color: #0b5ed7; + } + QPushButton:pressed { + background-color: #0a58ca; + } + QProgressBar { + border: 1px solid #555555; + border-radius: 4px; + text-align: center; + } + QProgressBar::chunk { + background-color: #0d6efd; + } + """) + + # Создаем центральный виджет + central_widget = QWidget() + self.setCentralWidget(central_widget) + main_layout = QHBoxLayout(central_widget) + + # Левая панель с настройками + settings_widget = QWidget() + settings_layout = QVBoxLayout(settings_widget) + settings_layout.setSpacing(10) + + # URL API + api_layout = QHBoxLayout() + api_label = QLabel("API URL:") + self.api_url_input = QLineEdit(self.api_url) + self.api_url_input.textChanged.connect(self.update_api_url) + api_layout.addWidget(api_label) + api_layout.addWidget(self.api_url_input) + settings_layout.addLayout(api_layout) + + # Промпт + prompt_label = QLabel("Промпт:") + self.prompt_input = QTextEdit() + self.prompt_input.setPlaceholderText("Введите описание изображения...") + settings_layout.addWidget(prompt_label) + settings_layout.addWidget(self.prompt_input) + + # Негативный промпт + neg_prompt_label = QLabel("Негативный промпт:") + self.neg_prompt_input = QTextEdit() + self.neg_prompt_input.setPlaceholderText("Введите нежелательные элементы...") + settings_layout.addWidget(neg_prompt_label) + settings_layout.addWidget(self.neg_prompt_input) + + # Добавляем менеджер моделей + self.model_manager = ModelManager(self.api_url) + self.model_manager.model_changed.connect(self.on_model_changed) + settings_layout.addWidget(self.model_manager) + + # Добавляем менеджер сэмплеров + self.sampler_manager = SamplerManager(self.api_url) + settings_layout.addWidget(self.sampler_manager) + + # Параметры генерации + params_layout = QVBoxLayout() + + # Размеры + size_layout = QHBoxLayout() + width_label = QLabel("Ширина:") + self.width_input = QSpinBox() + self.width_input.setRange(64, 2048) + self.width_input.setValue(512) + height_label = QLabel("Высота:") + self.height_input = QSpinBox() + self.height_input.setRange(64, 2048) + self.height_input.setValue(512) + size_layout.addWidget(width_label) + size_layout.addWidget(self.width_input) + size_layout.addWidget(height_label) + size_layout.addWidget(self.height_input) + params_layout.addLayout(size_layout) + + # Сид + seed_layout = QHBoxLayout() + seed_label = QLabel("Сид:") + self.seed_input = QSpinBox() + self.seed_input.setRange(-1, 999999999) + self.seed_input.setValue(-1) + seed_layout.addWidget(seed_label) + seed_layout.addWidget(self.seed_input) + params_layout.addLayout(seed_layout) + + # CFG Scale + cfg_layout = QHBoxLayout() + cfg_label = QLabel("CFG Scale:") + self.cfg_input = QDoubleSpinBox() + self.cfg_input.setRange(1.0, 30.0) + self.cfg_input.setValue(7.0) + cfg_layout.addWidget(cfg_label) + cfg_layout.addWidget(self.cfg_input) + params_layout.addLayout(cfg_layout) + + # Шаги + steps_layout = QHBoxLayout() + steps_label = QLabel("Шаги:") + self.steps_input = QSpinBox() + self.steps_input.setRange(1, 150) + self.steps_input.setValue(20) + steps_layout.addWidget(steps_label) + steps_layout.addWidget(self.steps_input) + params_layout.addLayout(steps_layout) + + settings_layout.addLayout(params_layout) + + # Добавляем кнопку консоли + self.console_button = QPushButton("Открыть консоль") + self.console_button.clicked.connect(self.toggle_console) + settings_layout.addWidget(self.console_button) + + # Кнопка генерации + self.generate_button = QPushButton("Сгенерировать") + self.generate_button.clicked.connect(self.start_generation) + settings_layout.addWidget(self.generate_button) + + # Прогресс бар + self.progress_bar = QProgressBar() + settings_layout.addWidget(self.progress_bar) + + # Добавляем левую панель в основной layout + main_layout.addWidget(settings_widget, stretch=1) + + # Правая панель для результатов + results_widget = QWidget() + results_layout = QVBoxLayout(results_widget) + + # Область для отображения изображений + scroll_area = QScrollArea() + scroll_area.setWidgetResizable(True) + scroll_area.setMinimumWidth(512) + self.results_content = QWidget() + self.results_content_layout = QVBoxLayout(self.results_content) + scroll_area.setWidget(self.results_content) + results_layout.addWidget(scroll_area) + + # Добавляем правую панель в основной layout + main_layout.addWidget(results_widget, stretch=2) + + # Устанавливаем размер окна + self.setMinimumSize(1200, 800) + + def toggle_console(self): + """Открыть/закрыть окно консоли""" + if self.console_window is None: + self.console_window = ConsoleWindow(parent=self) + self.console_window.console_closed.connect(self.on_console_closed) + self.console_window.show() + else: + self.console_window.close() + + def on_console_closed(self): + """Обработчик закрытия окна консоли""" + self.console_window = None + + def on_model_changed(self, model_name): + """Обработчик смены модели""" + print(f"Модель изменена на: {model_name}") + + def log_to_console(self, message): + """Логирование в консоль""" + if self.console_window: + self.console_window.log(message) + print(message) + + def start_generation(self): + self.generate_button.setEnabled(False) + self.progress_bar.setValue(0) + + params = { + "prompt": self.prompt_input.toPlainText(), + "negative_prompt": self.neg_prompt_input.toPlainText(), + "width": self.width_input.value(), + "height": self.height_input.value(), + "seed": self.seed_input.value(), + "cfg_scale": self.cfg_input.value(), + "steps": self.steps_input.value(), + "sampler_name": self.sampler_manager.get_current_sampler(), + "batch_size": 1 + } + + self.log_to_console(f"Начало генерации с параметрами: {json.dumps(params, indent=2, ensure_ascii=False)}") + + # Запускаем таймер для отслеживания прогресса + self.progress_timer = QTimer() + self.progress_timer.timeout.connect(self.check_progress) + self.progress_timer.start(100) # Проверяем каждые 100 мс + + self.generation_thread = GenerationThread(self.api_url_input.text(), params) + self.generation_thread.generation_complete.connect(self.handle_generation_complete) + self.generation_thread.error_occurred.connect(self.handle_error) + self.generation_thread.start() + + def check_progress(self): + """Проверка прогресса генерации""" + try: + response = requests.get(f"{self.api_url_input.text()}/sdapi/v1/progress") + if response.status_code == 200: + data = response.json() + progress = data.get('progress', 0) + current_step = data.get('current_step', 0) + total_steps = data.get('total_steps', 1) + + percentage = int(progress * 100) + self.progress_bar.setValue(percentage) + self.log_to_console(f"Прогресс: {percentage}% - Шаг {current_step}/{total_steps}") + + # Если генерация завершена, останавливаем таймер + if progress >= 1.0: + self.progress_timer.stop() + self.progress_timer = None + except Exception as e: + print(f"Ошибка при получении прогресса: {str(e)}") + + def handle_generation_complete(self, images): + """Обработка завершения генерации""" + # Останавливаем таймер, если он ещё активен + if self.progress_timer: + self.progress_timer.stop() + self.progress_timer = None + + self.generate_button.setEnabled(True) + self.progress_bar.setValue(100) + self.log_to_console("Генерация завершена успешно") + + # Очищаем предыдущие результаты + for i in reversed(range(self.results_content_layout.count())): + self.results_content_layout.itemAt(i).widget().setParent(None) + + for img_data in images: + try: + # Декодируем base64 в байты + img_bytes = base64.b64decode(img_data) + self.log_to_console(f"Размер полученных данных: {len(img_bytes)} байт") + + # Создаем объект изображения + img = Image.open(BytesIO(img_bytes)) + self.log_to_console(f"Размер изображения: {img.size}, формат: {img.format}") + + # Конвертируем в RGB если нужно + if img.mode != 'RGB': + img = img.convert('RGB') + + # Сохраняем во временный буфер как PNG + buffer = BytesIO() + img.save(buffer, format='PNG') + buffer.seek(0) + + # Создаем QPixmap + pixmap = QPixmap() + if not pixmap.loadFromData(buffer.getvalue()): + raise Exception("Не удалось создать QPixmap из данных") + + # Проверяем размер QPixmap + if pixmap.isNull(): + raise Exception("Создан пустой QPixmap") + + self.log_to_console(f"Размер QPixmap: {pixmap.width()}x{pixmap.height()}") + + # Создаем и добавляем QLabel с изображением + img_label = QLabel() + img_label.setPixmap(pixmap) + img_label.setAlignment(Qt.AlignmentFlag.AlignCenter) + + # Создаем кнопку сохранения + save_button = QPushButton("Сохранить") + save_button.clicked.connect(lambda checked, img=img: self.save_image(img)) + + # Создаем контейнер для изображения и кнопки + container = QWidget() + container_layout = QVBoxLayout(container) + container_layout.addWidget(img_label) + container_layout.addWidget(save_button) + + self.results_content_layout.addWidget(container) + self.log_to_console("Изображение успешно добавлено в интерфейс") + except Exception as e: + error_msg = f"Ошибка отображения изображения: {str(e)}" + self.log_to_console(error_msg) + error_label = QLabel(error_msg) + self.results_content_layout.addWidget(error_label) + + def handle_error(self, error_message): + """Обработка ошибок""" + self.generate_button.setEnabled(True) + self.log_to_console(f"Ошибка: {error_message}") + + def save_image(self, img): + file_path, _ = QFileDialog.getSaveFileName(self, "Сохранить изображение", "", "PNG Files (*.png);;JPEG Files (*.jpg)") + if file_path: + img.save(file_path) + + def update_api_url(self): + """Обновление URL API во всех компонентах""" + new_url = self.api_url_input.text() + self.model_manager.update_api_url(new_url) + self.sampler_manager.update_api_url(new_url) + self.log_to_console(f"API URL обновлен: {new_url}") + +if __name__ == '__main__': + app = QApplication(sys.argv) + window = MainWindow() + window.show() + sys.exit(app.exec()) \ No newline at end of file