""" Главное окно Comfy-Studio: трёхпанельный интерфейс, контекстное меню, полноэкранный просмотр, интеграция с ComfyUI. """ import os import sys import logging import subprocess from pathlib import Path from PyQt6.QtWidgets import ( QMainWindow, QSplitter, QWidget, QHBoxLayout, QVBoxLayout, QStatusBar, QMessageBox, QApplication, QMenu, QInputDialog, QLineEdit, QPushButton ) from PyQt6.QtCore import Qt, QModelIndex, QMimeData, QUrl from comfy_studio.core.config import DEFAULT_WINDOW_WIDTH, DEFAULT_WINDOW_HEIGHT, APP_NAME from comfy_studio.core.i18n import tr from comfy_studio.core.settings import load_settings, save_settings from comfy_studio.db.manager import DBManager from comfy_studio.services.watcher import FolderWatcher from comfy_studio.services.comfy_api import ComfyAPIClient, ComfySendWorker from comfy_studio.ui.theme import DARK_THEME_QSS from comfy_studio.ui.widgets.left_panel import LeftPanel from comfy_studio.ui.widgets.center_grid import CenterGrid, FilePathRole, FileIdRole from comfy_studio.ui.widgets.right_panel import RightPanel from comfy_studio.ui.widgets.settings_dialog import SettingsDialog from comfy_studio.ui.widgets.viewer import FullscreenViewer from comfy_studio.ui.widgets.translatable_text_edit import set_translate_languages logger = logging.getLogger(__name__) 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"]) set_translate_languages(self.settings.get("translate_languages", ["ru", "en"])) self.viewer = None self.comfy_sender = None self.setWindowTitle(APP_NAME) self.resize(DEFAULT_WINDOW_WIDTH, DEFAULT_WINDOW_HEIGHT) self.setStyleSheet(DARK_THEME_QSS) self.init_ui() self.bind_events() self.left_panel.set_tracked_folders(self.settings["tracked_paths"]) self.left_panel.update_tags(self.db_manager.get_all_tags()) 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(tr("settings_btn")) 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([240, 760, 380]) main_layout.addWidget(self.splitter) self.status_bar = QStatusBar() self.setStatusBar(self.status_bar) self.status_bar.showMessage(tr("app_ready")) 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.left_panel.search_changed.connect(self._on_search) self.left_panel.filters_changed.connect(self._reload_current) # Центральный грид 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.fullscreen_requested.connect(self._open_fullscreen) 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.right_panel.favorite_toggled.connect(self._on_favorite_toggled) self.right_panel.tag_added.connect(self._on_tag_added) self.right_panel.tag_removed.connect(self._on_tag_removed) self.right_panel.json_saved.connect(self._on_json_saved) self.settings_btn.clicked.connect(self._open_settings) # Реактивное обновление от watcher self.watcher.signals.file_added.connect(self._on_file_changed_externally) self.watcher.signals.file_removed.connect(self._on_file_changed_externally) # ──────────────────────────────── # Загрузка изображений # ──────────────────────────────── def load_folder_images(self, folder_path: str, search_query: str = ""): self.current_root_path = folder_path self.center_grid.path_label.setText(f"{tr('folder_label')}: {folder_path}") filters = self.left_panel.get_filters() files = self.db_manager.get_files_in_folder( folder_path, search_query, favorites_only=filters["favorites_only"], tag=filters["tag"], ) self.center_grid.model.set_files(files) self.status_bar.showMessage(f"{tr('files_shown')}: {len(files)}") def _reload_current(self): if self.current_root_path: self.load_folder_images(self.current_root_path, self.left_panel.search_input.text()) def _on_subfolder_selected(self, index: QModelIndex): folder_path = self.left_panel.folder_model.filePath(index) self.load_folder_images(folder_path) def _on_search(self, query: str): if self.current_root_path: self.load_folder_images(self.current_root_path, query) def _on_image_selected(self): indexes = self.center_grid.list_view.selectedIndexes() if not indexes: self.right_panel.clear_fields() return file_id = self.center_grid.model.data(indexes[0], FileIdRole) if file_id: details = self.db_manager.get_file_details(file_id) self.right_panel.display_metadata(details) def _on_file_changed_externally(self, filepath: str): if not self.current_root_path: return file_dir = str(Path(filepath).parent.resolve().as_posix()) current_dir = str(Path(self.current_root_path).resolve().as_posix()) if file_dir == current_dir: self._reload_current() # ──────────────────────────────── # Полноэкранный просмотр # ──────────────────────────────── def _open_fullscreen(self, row: int): paths = [f["filepath"] for f in self.center_grid.model.files] if not paths: return self.viewer = FullscreenViewer(paths, row) self.viewer.open_fullscreen() # ──────────────────────────────── # Контекстное меню и CRUD # ──────────────────────────────── def _on_context_menu(self, position): indexes = self.center_grid.list_view.selectedIndexes() if not indexes: return menu = QMenu(self) fullscreen_action = menu.addAction(tr("open_fullscreen")) fullscreen_action.setEnabled(len(indexes) == 1) favorite_action = menu.addAction(tr("toggle_favorite")) show_action = menu.addAction(tr("show_on_disk")) show_action.setEnabled(len(indexes) == 1) menu.addSeparator() rename_action = menu.addAction(tr("rename")) rename_action.setEnabled(len(indexes) == 1) delete_action = menu.addAction(tr("delete_selected")) action = menu.exec(self.center_grid.list_view.mapToGlobal(position)) if action == fullscreen_action: self._open_fullscreen(indexes[0].row()) elif action == favorite_action: for idx in indexes: file_id = self.center_grid.model.data(idx, FileIdRole) details = self.db_manager.get_file_details(file_id) if details: self.db_manager.set_favorite(file_id, not details["favorite"]) self._reload_current() self._on_image_selected() elif action == show_action: filepath = self.center_grid.model.data(indexes[0], 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, tr("rename_title"), f"{tr('rename_prompt')} {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): """Открывает файловый менеджер ОС с выделенным файлом.""" norm_path = os.path.normpath(filepath) if not os.path.exists(norm_path): return try: if sys.platform == "win32": subprocess.run(f'explorer /select,"{norm_path}"', shell=True) elif sys.platform == "darwin": subprocess.run(["open", "-R", norm_path]) else: subprocess.run(["xdg-open", os.path.dirname(norm_path)]) except Exception as e: logger.error(f"Не удалось открыть файловый менеджер: {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) try: os.rename(str(old_path), str(new_filepath)) self.status_bar.showMessage(f"{tr('renamed_ok')}: {new_filepath.name}") self._reload_current() except Exception as e: logger.error(f"Ошибка переименования {old_path.name}: {e}") QMessageBox.critical(self, tr("error"), f"{tr('rename_failed')}: {e}") def _on_copy_file(self, filepath: str): if os.path.exists(filepath): mime_data = QMimeData() mime_data.setUrls([QUrl.fromLocalFile(filepath)]) QApplication.clipboard().setMimeData(mime_data) self.status_bar.showMessage(tr("copied_to_clipboard")) def _on_delete_files(self, filepaths: list): if not filepaths: return confirm = QMessageBox.question( self, tr("delete_confirm_title"), tr("delete_confirm_text", n=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 except Exception as e: logger.error(f"Ошибка удаления файла {path}: {e}") self.status_bar.showMessage(f"{tr('deleted_from_disk')}: {deleted_count}") self._reload_current() # ──────────────────────────────── # Метаданные / избранное / теги # ──────────────────────────────── def _on_save_metadata(self, file_id: int, payload: dict): success = self.db_manager.update_file_details( file_id, payload["positive_prompt"], payload["negative_prompt"], payload["rating"] ) if success: self.status_bar.showMessage(tr("params_updated")) details = self.db_manager.get_file_details(file_id) self.right_panel.display_metadata(details) self._reload_current() else: QMessageBox.critical(self, tr("error"), tr("db_save_error")) def _on_json_saved(self, file_id: int, field: str, json_text: str): """Сохранение отредактированного workflow/prompt JSON в БД.""" success = self.db_manager.update_json_field(file_id, field, json_text or None) if success: self.status_bar.showMessage(tr("json_saved")) details = self.db_manager.get_file_details(file_id) self.right_panel.display_metadata(details) self._reload_current() else: QMessageBox.critical(self, tr("error"), tr("db_save_error")) def _on_favorite_toggled(self, file_id: int, favorite: bool): self.db_manager.set_favorite(file_id, favorite) self._reload_current() def _on_tag_added(self, file_id: int, tag: str): self.db_manager.add_tag_to_file(file_id, tag) details = self.db_manager.get_file_details(file_id) self.right_panel.display_metadata(details) self.left_panel.update_tags(self.db_manager.get_all_tags()) def _on_tag_removed(self, file_id: int, tag: str): self.db_manager.remove_tag_from_file(file_id, tag) details = self.db_manager.get_file_details(file_id) self.right_panel.display_metadata(details) self.left_panel.update_tags(self.db_manager.get_all_tags()) self._reload_current() # ──────────────────────────────── # Настройки # ──────────────────────────────── def _open_settings(self): dialog = SettingsDialog(self.settings, self) if dialog.exec() == SettingsDialog.DialogCode.Accepted: save_settings(self.settings) self.api_client = ComfyAPIClient(self.settings["comfyui_url"]) set_translate_languages(self.settings.get("translate_languages", ["ru", "en"])) self.left_panel.set_tracked_folders(self.settings["tracked_paths"]) self.status_bar.showMessage(tr("watcher_restarting")) 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(tr("settings_applied")) # ──────────────────────────────── # ComfyUI # ──────────────────────────────── def _on_send_to_comfy(self, prompt_json: str): self.status_bar.showMessage(tr("sending_prompt")) 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(tr("prompt_queued")) QMessageBox.information(self, tr("success"), message) else: self.status_bar.showMessage(tr("prompt_send_failed")) QMessageBox.critical(self, tr("network_error_title"), message)