art/main.py

import sys
import json
import base64
import requests
import os
from time import sleep
from PyQt6.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                            QHBoxLayout, QLabel, QLineEdit, QPushButton, 
                            QComboBox, QSpinBox, QDoubleSpinBox, QProgressBar,
                            QScrollArea, QTextEdit, QFileDialog, QMenuBar, QMenu)
from PyQt6.QtCore import Qt, QThread, pyqtSignal, QTimer
from PyQt6.QtGui import QPixmap, QImage, QAction, QFontDatabase, QFont
from PIL import Image
from io import BytesIO
from pathlib import Path
from dotenv import load_dotenv

# Загружаем переменные окружения
load_dotenv()

# Импортируем модули
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
from modules.size_preset_manager import SizePresetManager
from modules.settings_menu import SettingsMenu
from modules.styles.base_window import BaseWindow
from modules.styles.navigation import NavigationPanel
from modules.styles.page import PageContainer
from modules.pages.home_page import HomePage
from modules.pages.models_page import ModelsPage
from modules.pages.settings_page import SettingsPage
from scripts.download_resources import download_font

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(BaseWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Stable Diffusion UI")
        self.setMinimumSize(1200, 800)
        
        # Получаем URL API из переменных окружения
        self.api_url = os.getenv('SD_API_URL', 'http://localhost:3000')
        
        # Инициализируем менеджеры
        self.model_manager = ModelManager(api_url=self.api_url)
        self.sampler_manager = SamplerManager(api_url=self.api_url)
        self.size_preset_manager = SizePresetManager()
        
        # Инициализируем окна
        self.prompt_window = PromptWindow()
        self.negative_prompt_window = NegativePromptWindow()
        self.console_window = ConsoleWindow()
        
        self.setup_ui()

    def setup_ui(self):
        # Загружаем шрифт
        font_path = Path("resources/fonts/InterVariable.ttf")
        if not font_path.exists():
            download_font()
        font_id = QFontDatabase.addApplicationFont(str(font_path))
        if font_id != -1:
            font_family = QFontDatabase.applicationFontFamilies(font_id)[0]
            app_font = QFont(font_family)
            QApplication.setFont(app_font)
        
        # Создаем центральный виджет
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout(central_widget)
        main_layout.setContentsMargins(0, 0, 0, 0)
        main_layout.setSpacing(0)
        
        # Создаем верхнюю панель
        top_panel = QHBoxLayout()
        
        # Кнопки для окон промптов
        prompt_btn = QPushButton("Библиотека промптов")
        prompt_btn.clicked.connect(self.show_prompt_window)
        top_panel.addWidget(prompt_btn)
        
        negative_prompt_btn = QPushButton("Негативные промпты")
        negative_prompt_btn.clicked.connect(self.show_negative_prompt_window)
        top_panel.addWidget(negative_prompt_btn)
        
        # Пресеты разрешения
        resolution_combo = QComboBox()
        resolution_combo.addItems(self.size_preset_manager.get_preset_names())
        resolution_combo.currentTextChanged.connect(self.on_resolution_changed)
        top_panel.addWidget(resolution_combo)
        
        # Кнопка консоли
        console_btn = QPushButton("Консоль")
        console_btn.clicked.connect(self.show_console)
        top_panel.addWidget(console_btn)
        
        top_panel.addStretch()
        main_layout.addLayout(top_panel)
        
        # Создаем горизонтальный layout для навигации и контента
        content_layout = QHBoxLayout()
        
        # Добавляем навигационную панель
        self.nav_panel = NavigationPanel()
        content_layout.addWidget(self.nav_panel)
        
        # Создаем контейнер для страниц
        self.page_container = PageContainer()
        content_layout.addWidget(self.page_container)
        
        main_layout.addLayout(content_layout)
        
        # Добавляем страницы
        self.home_page = HomePage(
            model_manager=self.model_manager,
            sampler_manager=self.sampler_manager,
            size_preset_manager=self.size_preset_manager
        )
        self.models_page = ModelsPage(self.model_manager)
        self.settings_page = SettingsPage()
        
        self.page_container.add_page("home", self.home_page)
        self.page_container.add_page("models", self.models_page)
        self.page_container.add_page("settings", self.settings_page)
        
        # Подключаем сигнал смены страницы
        self.nav_panel.page_changed.connect(self.switch_page)
        
        # Показываем домашнюю страницу по умолчанию
        self.page_container.show_page("home")
        
        # Загружаем данные
        self.load_data()
    
    def load_data(self):
        """Загрузка данных с сервера"""
        try:
            # Загружаем список моделей
            self.model_manager.load_models()
            
            # Загружаем список сэмплеров
            self.sampler_manager.load_samplers()
            
            # Обновляем интерфейс
            self.home_page.update_ui()
            self.models_page.update_ui()
            
        except Exception as e:
            self.console_window.log(f"Ошибка загрузки данных: {str(e)}", "ERROR")
    
    def show_prompt_window(self):
        """Показать окно промптов"""
        self.prompt_window.show()
    
    def show_negative_prompt_window(self):
        """Показать окно негативных промптов"""
        self.negative_prompt_window.show()
    
    def show_console(self):
        """Показать консоль"""
        self.console_window.show()
    
    def on_resolution_changed(self, preset_name: str):
        """Обработка изменения пресета разрешения"""
        preset = self.size_preset_manager.get_preset(preset_name)
        if preset:
            self.home_page.set_resolution(preset.width, preset.height)
    
    def switch_page(self, page_name):
        self.page_container.show_page(page_name)

def main():
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec())

if __name__ == "__main__":
    main()