Воркед

modules/model_manager.py

@@ -2,6 +2,17 @@ from PyQt6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout,
                             QComboBox, QPushButton, QLabel)
 from PyQt6.QtCore import pyqtSignal
 import requests
+from typing import List, Dict
+import os
+from dotenv import load_dotenv
+
+# Загружаем переменные окружения
+load_dotenv()
+
+class Model:
+    def __init__(self, name: str, title: str = None):
+        self.name = name
+        self.title = title or name
 
 class ModelManager(QWidget):
     model_changed = pyqtSignal(str)
@@ -9,6 +20,8 @@ class ModelManager(QWidget):
     def __init__(self, api_url, parent=None):
         super().__init__(parent)
         self.api_url = api_url
+        self.models = []
+        self.current_model = None
         self.setup_ui()
         self.load_models()
 
@@ -97,33 +110,86 @@ class ModelManager(QWidget):
         try:
             response = requests.get(f"{self.api_url}/sdapi/v1/sd-models")
             if response.status_code == 200:
-                models = response.json()
-                current_model = self.model_combo.currentText()
+                models_data = response.json()
+                self.models = [
+                    Model(model.get('title', ''), model.get('model_name', ''))
+                    for model in models_data
+                ]
                 
+                # Получаем текущую модель
+                options_response = requests.get(f"{self.api_url}/sdapi/v1/options")
+                if options_response.status_code == 200:
+                    current_model_name = options_response.json().get('sd_model_checkpoint')
+                    self.current_model = next(
+                        (model for model in self.models if model.name == current_model_name),
+                        self.models[0] if self.models else None
+                    )
+                
+                # Обновляем UI
                 self.model_combo.clear()
-                for model in models:
-                    self.model_combo.addItem(model.get('title', ''))
+                for model in self.models:
+                    self.model_combo.addItem(model.title)
+                
+                # Устанавливаем текущую модель
+                if self.current_model:
+                    index = self.model_combo.findText(self.current_model.title)
+                    if index >= 0:
+                        self.model_combo.setCurrentIndex(index)
                 
-                # Восстанавливаем выбранную модель, если она всё ещё доступна
-                index = self.model_combo.findText(current_model)
-                if index >= 0:
-                    self.model_combo.setCurrentIndex(index)
         except Exception as e:
             print(f"Ошибка при загрузке моделей: {str(e)}")
+            self.models = []
+            self.current_model = None
 
-    def on_model_changed(self, model_name):
+    def on_model_changed(self, model_title):
         """Обработка смены модели"""
+        model = next(
+            (m for m in self.models if m.title == model_title),
+            None
+        )
+        if model:
+            try:
+                response = requests.post(
+                    f"{self.api_url}/sdapi/v1/options",
+                    json={"sd_model_checkpoint": model.name}
+                )
+                if response.status_code == 200:
+                    self.current_model = model
+                    self.model_changed.emit(model.name)
+            except Exception as e:
+                print(f"Ошибка при смене модели: {str(e)}")
+
+    def update_api_url(self, new_url):
+        """Обновить URL API"""
+        self.api_url = new_url
+        self.load_models()
+
+    def get_models(self):
+        """Возвращает список доступных моделей"""
+        return self.models
+
+    def get_model_names(self):
+        """Возвращает список имен моделей"""
+        return [model.name for model in self.models]
+
+    def get_current_model(self):
+        """Возвращает текущую модель"""
+        return self.current_model
+
+    def switch_model(self, model_name: str) -> bool:
+        """Переключение на другую модель"""
         try:
             response = requests.post(
                 f"{self.api_url}/sdapi/v1/options",
                 json={"sd_model_checkpoint": model_name}
             )
             if response.status_code == 200:
-                self.model_changed.emit(model_name)
+                self.current_model = next(
+                    (model for model in self.models if model.name == model_name),
+                    None
+                )
+                return True
+            return False
         except Exception as e:
-            print(f"Ошибка при смене модели: {str(e)}")
-
-    def update_api_url(self, new_url):
-        """Обновить URL API"""
-        self.api_url = new_url
-        self.load_models() 
+            print(f"Ошибка переключения модели: {str(e)}")
+            return False