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 from modules.prompt_window import PromptWindow 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.prompt_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("Введите описание изображения...") prompt_layout = QVBoxLayout() prompt_header = QHBoxLayout() prompt_header.addWidget(prompt_label) # Добавляем кнопку промптов рядом с полем промпта self.prompt_button = QPushButton("Готовые промпты") self.prompt_button.clicked.connect(self.toggle_prompt_window) prompt_header.addWidget(self.prompt_button) prompt_header.addStretch() prompt_layout.addLayout(prompt_header) prompt_layout.addWidget(self.prompt_input) settings_layout.addLayout(prompt_layout) # Негативный промпт 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 toggle_prompt_window(self): """Открыть/закрыть окно с промптами""" if self.prompt_window is None: self.prompt_window = PromptWindow(parent=self) self.prompt_window.prompt_selected.connect(self.add_prompt) self.prompt_window.show() else: self.prompt_window.close() def add_prompt(self, prompt): """Добавить промпт в поле ввода""" current_text = self.prompt_input.toPlainText() if current_text: current_text += ", " self.prompt_input.setText(current_text + prompt) 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())