dinlo
354 lines
17 KiB
Python
354 lines
17 KiB
Python
# ui/main_window.py
|
||
import os
|
||
import json
|
||
import logging
|
||
import subprocess
|
||
from pathlib import Path
|
||
import requests
|
||
from PyQt6.QtWidgets import (
|
||
QMainWindow, QSplitter, QWidget, QHBoxLayout, QVBoxLayout,
|
||
QStatusBar, QMessageBox, QApplication, QMenu, QInputDialog, QLineEdit, QPushButton
|
||
)
|
||
from PyQt6.QtCore import Qt, QModelIndex, QMimeData, QUrl, QThread, pyqtSignal
|
||
|
||
from ui.styles import DARK_THEME_QSS
|
||
from ui.widgets.left_panel import LeftPanel
|
||
from ui.widgets.center_grid import CenterGrid
|
||
from ui.widgets.right_panel import RightPanel
|
||
from ui.widgets.settings_dialog import SettingsDialog
|
||
|
||
from database.db_manager import DBManager
|
||
from watcher.fs_watcher import FolderWatcher
|
||
from utils.api_client import ComfyAPIClient
|
||
from utils.settings import load_settings, save_settings
|
||
|
||
FilePathRole = Qt.ItemDataRole.UserRole + 2
|
||
|
||
logger = logging.getLogger("ComfyGallery.MainWindow")
|
||
|
||
|
||
class ComfySendWorker(QThread):
|
||
finished = pyqtSignal(bool, str) # (success, message)
|
||
|
||
def __init__(self, api_client: ComfyAPIClient, prompt_json_str: str):
|
||
super().__init__()
|
||
self.api_client = api_client
|
||
self.prompt_json_str = prompt_json_str
|
||
|
||
def run(self):
|
||
logger.info("Запущена фоновая задача отправки промта в ComfyUI.")
|
||
try:
|
||
prompt_data = json.loads(self.prompt_json_str)
|
||
payload = {"prompt": prompt_data}
|
||
url = f"{self.api_client.base_url}/prompt"
|
||
|
||
logger.debug(f"Отправка POST запроса на {url}")
|
||
response = requests.post(url, json=payload, timeout=5)
|
||
|
||
if response.status_code == 200:
|
||
res_data = response.json()
|
||
prompt_id = res_data.get("prompt_id", "Неизвестно")
|
||
logger.info(f"Задача успешно принята сервером ComfyUI. ID задачи: {prompt_id}")
|
||
self.finished.emit(
|
||
True,
|
||
f"Промт успешно добавлен в очередь генерации ComfyUI!\n\nID задачи: {prompt_id}"
|
||
)
|
||
else:
|
||
logger.error(f"Сервер ComfyUI отклонил запрос. Код: {response.status_code}, Ответ: {response.text}")
|
||
self.finished.emit(
|
||
False,
|
||
f"Сервер ComfyUI вернул ошибку {response.status_code}:\n{response.text}"
|
||
)
|
||
except json.JSONDecodeError as je:
|
||
logger.error(f"Не удалось распарсить prompt_json метаданных: {je}")
|
||
self.finished.emit(False, f"Ошибка метаданных изображения (неверный формат JSON):\n{je}")
|
||
except requests.exceptions.RequestException as re:
|
||
logger.error(f"Сетевое исключение при отправке на {self.api_client.base_url}: {re}")
|
||
self.finished.emit(
|
||
False,
|
||
f"Не удалось подключиться к ComfyUI по адресу {self.api_client.base_url}.\n"
|
||
f"Убедитесь, что сервер ComfyUI запущен на этом порту.\n\nПодробности:\n{re}"
|
||
)
|
||
except Exception as e:
|
||
logger.error(f"Непредвиденная ошибка в сетевом воркере: {e}")
|
||
self.finished.emit(False, f"Произошла непредвиденная ошибка при отправке: {e}")
|
||
|
||
|
||
class MainWindow(QMainWindow):
|
||
def __init__(self, db_manager: DBManager, watcher: FolderWatcher):
|
||
super().__init__()
|
||
self.db_manager = db_manager
|
||
self.watcher = watcher
|
||
|
||
self.settings = load_settings()
|
||
self.current_root_path = None
|
||
self.api_client = ComfyAPIClient(self.settings["comfyui_url"])
|
||
|
||
self.setWindowTitle("ComfyGallery")
|
||
self.resize(1300, 850)
|
||
self.setStyleSheet(DARK_THEME_QSS)
|
||
|
||
self.init_ui()
|
||
self.bind_events()
|
||
|
||
self.left_panel.set_tracked_folders(self.settings["tracked_paths"])
|
||
logger.info("Главное окно успешно инициализировано.")
|
||
|
||
def init_ui(self):
|
||
central_widget = QWidget()
|
||
self.setCentralWidget(central_widget)
|
||
|
||
main_layout = QVBoxLayout(central_widget)
|
||
main_layout.setContentsMargins(8, 8, 8, 8)
|
||
|
||
top_bar = QHBoxLayout()
|
||
self.settings_btn = QPushButton("Настройки подключения и папок")
|
||
top_bar.addWidget(self.settings_btn)
|
||
top_bar.addStretch()
|
||
main_layout.addLayout(top_bar)
|
||
|
||
self.splitter = QSplitter(Qt.Orientation.Horizontal)
|
||
self.left_panel = LeftPanel()
|
||
self.center_grid = CenterGrid()
|
||
self.right_panel = RightPanel()
|
||
|
||
self.splitter.addWidget(self.left_panel)
|
||
self.splitter.addWidget(self.center_grid)
|
||
self.splitter.addWidget(self.right_panel)
|
||
|
||
self.splitter.setSizes([230, 720, 350])
|
||
main_layout.addWidget(self.splitter)
|
||
|
||
self.status_bar = QStatusBar()
|
||
self.setStatusBar(self.status_bar)
|
||
self.status_bar.showMessage("Готов к работе")
|
||
|
||
def bind_events(self):
|
||
self.left_panel.tree_view.clicked.connect(self._on_subfolder_selected)
|
||
self.left_panel.root_folder_changed.connect(self.load_folder_images)
|
||
self.center_grid.list_view.selectionModel().selectionChanged.connect(self._on_image_selected)
|
||
|
||
self.center_grid.list_view.copy_pressed.connect(self._on_copy_file)
|
||
self.center_grid.list_view.delete_pressed.connect(self._on_delete_files)
|
||
|
||
# Контекстное меню
|
||
self.center_grid.list_view.customContextMenuRequested.connect(self._on_context_menu)
|
||
|
||
self.right_panel.save_clicked.connect(self._on_save_metadata)
|
||
self.right_panel.send_to_comfy_clicked.connect(self._on_send_to_comfy)
|
||
self.left_panel.search_changed.connect(self._on_search)
|
||
|
||
self.settings_btn.clicked.connect(self._open_settings)
|
||
|
||
# Автоматическое реактивное обновление
|
||
self.watcher.signals.file_added.connect(self._on_file_added_externally)
|
||
self.watcher.signals.file_removed.connect(self._on_file_removed_externally)
|
||
|
||
def load_folder_images(self, folder_path: str, search_query: str = ""):
|
||
logger.info(f"Загрузка изображений из папки: {folder_path} (фильтр: '{search_query}')")
|
||
self.current_root_path = folder_path
|
||
self.center_grid.path_label.setText(f"Папка: {folder_path}")
|
||
|
||
files = self.db_manager.get_files_in_folder(folder_path, search_query)
|
||
self.center_grid.model.set_files(files)
|
||
self.status_bar.showMessage(f"Файлов отображено: {len(files)}")
|
||
|
||
def _on_subfolder_selected(self, index: QModelIndex):
|
||
folder_path = self.left_panel.folder_model.filePath(index)
|
||
self.load_folder_images(folder_path)
|
||
|
||
def _on_image_selected(self):
|
||
indexes = self.center_grid.list_view.selectedIndexes()
|
||
if not indexes:
|
||
self.right_panel.clear_fields()
|
||
return
|
||
|
||
idx = indexes[0]
|
||
file_id = self.center_grid.model.data(idx, Qt.ItemDataRole.UserRole + 1)
|
||
if file_id:
|
||
logger.debug(f"Выделен файл с ID: {file_id}")
|
||
details = self.db_manager.get_file_details(file_id)
|
||
self.right_panel.display_metadata(details)
|
||
|
||
def _on_search(self, query: str):
|
||
if self.current_root_path:
|
||
self.load_folder_images(self.current_root_path, query)
|
||
|
||
# --- РЕАКТИВНЫЕ СОБЫТИЯ WATCHDOG ---
|
||
|
||
def _on_file_added_externally(self, filepath: str):
|
||
logger.debug(f"Событие Watcher: обнаружен новый файл {filepath}")
|
||
file_dir = str(Path(filepath).parent.resolve().as_posix())
|
||
current_dir = str(Path(self.current_root_path).resolve().as_posix()) if self.current_root_path else ""
|
||
|
||
if file_dir == current_dir:
|
||
self.load_folder_images(self.current_root_path)
|
||
|
||
def _on_file_removed_externally(self, filepath: str):
|
||
logger.debug(f"Событие Watcher: файл удален {filepath}")
|
||
file_dir = str(Path(filepath).parent.resolve().as_posix())
|
||
current_dir = str(Path(self.current_root_path).resolve().as_posix()) if self.current_root_path else ""
|
||
|
||
if file_dir == current_dir:
|
||
self.load_folder_images(self.current_root_path)
|
||
|
||
# --- КОНТЕКСТНОЕ МЕНЮ И CRUD ---
|
||
|
||
def _on_context_menu(self, position):
|
||
indexes = self.center_grid.list_view.selectedIndexes()
|
||
if not indexes: return
|
||
|
||
menu = QMenu(self)
|
||
|
||
# Новый пункт контекстного меню
|
||
show_action = menu.addAction("Показать на диске")
|
||
show_action.setEnabled(len(indexes) == 1)
|
||
|
||
menu.addSeparator()
|
||
|
||
rename_action = menu.addAction("Переименовать")
|
||
rename_action.setEnabled(len(indexes) == 1)
|
||
|
||
delete_action = menu.addAction("Удалить выбранные")
|
||
|
||
action = menu.exec(self.center_grid.list_view.mapToGlobal(position))
|
||
|
||
if action == show_action:
|
||
idx = indexes[0]
|
||
filepath = self.center_grid.model.data(idx, FilePathRole)
|
||
self._show_in_explorer(filepath)
|
||
|
||
elif action == rename_action:
|
||
idx = indexes[0]
|
||
filepath = self.center_grid.model.data(idx, FilePathRole)
|
||
filename = self.center_grid.model.data(idx, Qt.ItemDataRole.DisplayRole)
|
||
|
||
new_name, ok = QInputDialog.getText(
|
||
self, "Переименование", f"Новое имя файла для {filename}:",
|
||
QLineEdit.EchoMode.Normal, filename
|
||
)
|
||
if ok and new_name.strip() and new_name != filename:
|
||
self._rename_file(filepath, new_name.strip())
|
||
|
||
elif action == delete_action:
|
||
paths = [self.center_grid.model.data(idx, FilePathRole) for idx in indexes]
|
||
self._on_delete_files(paths)
|
||
|
||
def _show_in_explorer(self, filepath: str):
|
||
""" Открывает Проводник Windows и подсвечивает (выделяет) данный файл. """
|
||
norm_path = os.path.normpath(filepath)
|
||
if not os.path.exists(norm_path):
|
||
logger.warning(f"Попытка показать несуществующий файл: {norm_path}")
|
||
return
|
||
|
||
logger.info(f"Открытие проводника для файла: {norm_path}")
|
||
try:
|
||
# Команда explorer /select открывает папку и фокусом выделяет файл
|
||
subprocess.run(f'explorer /select,"{norm_path}"', shell=True)
|
||
except Exception as e:
|
||
logger.error(f"Не удалось открыть Проводник Windows: {e}")
|
||
|
||
def _rename_file(self, old_filepath: str, new_filename: str):
|
||
old_path = Path(old_filepath)
|
||
new_filepath = old_path.parent / new_filename
|
||
|
||
if not new_filepath.suffix:
|
||
new_filepath = new_filepath.with_suffix(old_path.suffix)
|
||
|
||
logger.info(f"Переименование файла: {old_path.name} -> {new_filepath.name}")
|
||
try:
|
||
os.rename(str(old_path), str(new_filepath))
|
||
self.status_bar.showMessage(f"Файл успешно переименован: {new_filepath.name}")
|
||
if self.current_root_path:
|
||
self.load_folder_images(self.current_root_path)
|
||
except Exception as e:
|
||
logger.error(f"Ошибка переименования файла {old_path.name}: {e}")
|
||
QMessageBox.critical(self, "Ошибка", f"Не удалось переименовать файл: {e}")
|
||
|
||
def _on_copy_file(self, filepath: str):
|
||
logger.info(f"Копирование файла в буфер обмена: {filepath}")
|
||
if os.path.exists(filepath):
|
||
mime_data = QMimeData()
|
||
mime_data.setUrls([QUrl.fromLocalFile(filepath)])
|
||
QApplication.clipboard().setMimeData(mime_data)
|
||
self.status_bar.showMessage("Изображение скопировано в буфер обмена.")
|
||
|
||
def _on_delete_files(self, filepaths: list):
|
||
if not filepaths: return
|
||
|
||
logger.info(f"Запрос удаления группы файлов: {len(filepaths)} шт.")
|
||
confirm = QMessageBox.question(
|
||
self, "Удаление", f"Вы действительно хотите безвозвратно удалить {len(filepaths)} файлов?",
|
||
QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No
|
||
)
|
||
if confirm == QMessageBox.StandardButton.Yes:
|
||
deleted_count = 0
|
||
for path in filepaths:
|
||
try:
|
||
if os.path.exists(path):
|
||
os.remove(path)
|
||
deleted_count += 1
|
||
logger.debug(f"Файл успешно удален физически: {path}")
|
||
except Exception as e:
|
||
logger.error(f"Ошибка физического удаления файла {path}: {e}")
|
||
self.status_bar.showMessage(f"Удалено файлов с диска: {deleted_count}")
|
||
if self.current_root_path:
|
||
self.load_folder_images(self.current_root_path)
|
||
|
||
# --- НАСТРОЙКИ ---
|
||
|
||
def _open_settings(self):
|
||
logger.info("Открытие диалогового окна настроек.")
|
||
dialog = SettingsDialog(self.settings, self)
|
||
if dialog.exec() == SettingsDialog.DialogCode.Accepted:
|
||
save_settings(self.settings)
|
||
logger.info("Пользователь сохранил новые настройки.")
|
||
|
||
self.api_client = ComfyAPIClient(self.settings["comfyui_url"])
|
||
self.left_panel.set_tracked_folders(self.settings["tracked_paths"])
|
||
|
||
self.status_bar.showMessage("Перезапуск службы отслеживания...")
|
||
self.watcher.stop_monitoring()
|
||
self.watcher.start_monitoring(self.settings["tracked_paths"])
|
||
|
||
if self.settings["tracked_paths"]:
|
||
first_path = self.settings["tracked_paths"][0]
|
||
self.left_panel.set_root_path(first_path)
|
||
self.load_folder_images(first_path)
|
||
|
||
self.status_bar.showMessage("Настройки успешно применены.")
|
||
|
||
# --- ВЗАИМОДЕЙСТВИЕ С COMFYUI ---
|
||
|
||
def _on_save_metadata(self, file_id: int, payload: dict):
|
||
logger.info(f"Сохранение измененных пользователем метаданных для ID {file_id}")
|
||
success = self.db_manager.update_file_details(
|
||
file_id, payload["positive_prompt"], payload["negative_prompt"], payload["rating"]
|
||
)
|
||
if success:
|
||
self.status_bar.showMessage("Параметры успешно обновлены.")
|
||
details = self.db_manager.get_file_details(file_id)
|
||
self.right_panel.display_metadata(details)
|
||
if self.current_root_path:
|
||
self.load_folder_images(self.current_root_path)
|
||
else:
|
||
logger.error(f"Не удалось обновить метаданные в БД для ID {file_id}")
|
||
QMessageBox.critical(self, "Ошибка", "Не удалось сохранить изменения в БД.")
|
||
|
||
def _on_send_to_comfy(self, prompt_json: str):
|
||
self.status_bar.showMessage("Отправка промта в ComfyUI...")
|
||
self.right_panel.send_btn.setEnabled(False)
|
||
|
||
# Создаем и запускаем фоновый поток воркера
|
||
self.comfy_sender = ComfySendWorker(self.api_client, prompt_json)
|
||
self.comfy_sender.finished.connect(self._on_send_finished)
|
||
self.comfy_sender.start()
|
||
|
||
def _on_send_finished(self, success: bool, message: str):
|
||
self.right_panel.send_btn.setEnabled(True)
|
||
if success:
|
||
self.status_bar.showMessage("Промт успешно добавлен в очередь ComfyUI!")
|
||
QMessageBox.information(self, "Успешно отправлено", message)
|
||
else:
|
||
self.status_bar.showMessage("Не удалось отправить промт.")
|
||
QMessageBox.critical(self, "Ошибка сети / API", message) |