510 lines
21 KiB
Python
510 lines
21 KiB
Python
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
|
|
from modules.negative_prompt_window import NegativePromptWindow
|
|
|
|
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.negative_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_layout = QVBoxLayout()
|
|
prompt_header = QHBoxLayout()
|
|
prompt_label = QLabel("Промпт:")
|
|
self.prompt_input = QTextEdit()
|
|
self.prompt_input.setMinimumHeight(100)
|
|
prompt_buttons = QHBoxLayout()
|
|
self.prompt_button = QPushButton("Готовые промпты")
|
|
self.prompt_button.clicked.connect(self.toggle_prompt_window)
|
|
self.negative_prompt_button = QPushButton("Негативные промпты")
|
|
self.negative_prompt_button.clicked.connect(self.toggle_negative_prompt_window)
|
|
prompt_buttons.addWidget(self.prompt_button)
|
|
prompt_buttons.addWidget(self.negative_prompt_button)
|
|
prompt_header.addWidget(prompt_label)
|
|
prompt_header.addLayout(prompt_buttons)
|
|
prompt_layout.addLayout(prompt_header)
|
|
prompt_layout.addWidget(self.prompt_input)
|
|
settings_layout.addLayout(prompt_layout)
|
|
|
|
# Негативный промпт
|
|
negative_prompt_layout = QVBoxLayout()
|
|
negative_prompt_label = QLabel("Негативный промпт:")
|
|
self.negative_prompt_input = QTextEdit()
|
|
self.negative_prompt_input.setMinimumHeight(100)
|
|
negative_prompt_layout.addWidget(negative_prompt_label)
|
|
negative_prompt_layout.addWidget(self.negative_prompt_input)
|
|
settings_layout.addLayout(negative_prompt_layout)
|
|
|
|
# Добавляем менеджер моделей
|
|
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 toggle_negative_prompt_window(self):
|
|
"""Открыть/закрыть окно с негативными промптами"""
|
|
if self.negative_prompt_window is None:
|
|
self.negative_prompt_window = NegativePromptWindow(parent=self)
|
|
self.negative_prompt_window.prompt_selected.connect(self.add_negative_prompt)
|
|
self.negative_prompt_window.show()
|
|
else:
|
|
self.negative_prompt_window.close()
|
|
|
|
def add_prompt(self, text):
|
|
"""Добавить промпт в поле ввода"""
|
|
current_text = self.prompt_input.toPlainText()
|
|
if current_text:
|
|
current_text += ", "
|
|
self.prompt_input.setText(current_text + text)
|
|
|
|
def add_negative_prompt(self, text):
|
|
"""Добавить негативный промпт в поле ввода"""
|
|
current_text = self.negative_prompt_input.toPlainText()
|
|
if current_text:
|
|
current_text += ", "
|
|
self.negative_prompt_input.setText(current_text + text)
|
|
|
|
def start_generation(self):
|
|
self.generate_button.setEnabled(False)
|
|
self.progress_bar.setValue(0)
|
|
|
|
params = {
|
|
"prompt": self.prompt_input.toPlainText(),
|
|
"negative_prompt": self.negative_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))
|
|
|
|
# Создаем кнопку копирования в буфер обмена
|
|
copy_button = QPushButton("Копировать")
|
|
copy_button.clicked.connect(lambda checked, img=img: self.copy_to_clipboard(img))
|
|
|
|
# Создаем контейнер для изображения и кнопок
|
|
container = QWidget()
|
|
container_layout = QVBoxLayout(container)
|
|
container_layout.addWidget(img_label)
|
|
|
|
# Создаем горизонтальный layout для кнопок
|
|
buttons_layout = QHBoxLayout()
|
|
buttons_layout.addWidget(save_button)
|
|
buttons_layout.addWidget(copy_button)
|
|
container_layout.addLayout(buttons_layout)
|
|
|
|
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 copy_to_clipboard(self, img):
|
|
"""Копирование изображения в буфер обмена"""
|
|
# Конвертируем изображение в QPixmap
|
|
buffer = BytesIO()
|
|
img.save(buffer, format='PNG')
|
|
buffer.seek(0)
|
|
pixmap = QPixmap()
|
|
pixmap.loadFromData(buffer.getvalue())
|
|
|
|
# Копируем в буфер обмена
|
|
clipboard = QApplication.clipboard()
|
|
clipboard.setPixmap(pixmap)
|
|
|
|
# Показываем уведомление
|
|
self.log_to_console("Изображение скопировано в буфер обмена")
|
|
|
|
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()) |