Files
art/modules/model_manager.py

70 lines
2.6 KiB
Python

from PyQt6.QtWidgets import (QWidget, QVBoxLayout, QHBoxLayout,
QComboBox, QPushButton, QLabel)
from PyQt6.QtCore import pyqtSignal
import requests
class ModelManager(QWidget):
model_changed = pyqtSignal(str)
def __init__(self, api_url, parent=None):
super().__init__(parent)
self.api_url = api_url
self.setup_ui()
self.load_models()
def setup_ui(self):
layout = QHBoxLayout(self)
layout.setContentsMargins(0, 0, 0, 0)
# Метка
label = QLabel("Модель:")
layout.addWidget(label)
# Выпадающий список моделей
self.model_combo = QComboBox()
self.model_combo.currentTextChanged.connect(self.on_model_changed)
layout.addWidget(self.model_combo)
# Кнопка обновления
refresh_button = QPushButton("🔄")
refresh_button.setToolTip("Обновить список моделей")
refresh_button.clicked.connect(self.load_models)
layout.addWidget(refresh_button)
layout.addStretch()
def load_models(self):
"""Загрузить список доступных моделей"""
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()
self.model_combo.clear()
for model in models:
self.model_combo.addItem(model.get('title', ''))
# Восстанавливаем выбранную модель, если она всё ещё доступна
index = self.model_combo.findText(current_model)
if index >= 0:
self.model_combo.setCurrentIndex(index)
except Exception as e:
print(f"Ошибка при загрузке моделей: {str(e)}")
def on_model_changed(self, model_name):
"""Обработка смены модели"""
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)
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()