Initial release: Comfy-Studio v1.0.0

Десктопная галерея и менеджер изображений ComfyUI (PyQt6 + SQLite): - мультипапочный мониторинг (watchdog), грид миниатюр, drag-and-drop - теги, избранное, рейтинг, поиск по промтам - парсер метаданных ComfyUI (PNG/WebP/JPEG), редактор workflow/prompt JSON - перевод выделенного текста через контекстное меню (настраиваемые языки) - отправка промта в очередь ComfyUI с внедрением правок и нового seed - полноэкранный просмотр, локализация ru/en, тёмная тема Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-07-02 14:09:57 +08:00
commit 34a40ae759
29 changed files with 3109 additions and 0 deletions
@@ -0,0 +1,389 @@
+"""
+Менеджер базы данных SQLite: папки, файлы, метаданные, теги, избранное.
+"""
+
+import sqlite3
+import logging
+from pathlib import Path
+from typing import Optional, List, Dict, Any
+from contextlib import contextmanager
+
+from comfy_studio.db.schema import PRAGMA_FOREIGN_KEYS, CREATE_TABLES_SQL, CREATE_INDEXES_SQL
+
+logger = logging.getLogger(__name__)
+
+
+def _normalize(path: str) -> str:
+    return str(Path(path).resolve().as_posix())
+
+
+class DBManager:
+    def __init__(self, db_path: str = "comfystudio.db"):
+        self.db_path = str(db_path)
+        self._initialize_db()
+
+    @contextmanager
+    def _get_connection(self):
+        conn = None
+        try:
+            conn = sqlite3.connect(self.db_path)
+            conn.execute(PRAGMA_FOREIGN_KEYS)
+            conn.row_factory = sqlite3.Row
+            yield conn
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка соединения с БД {self.db_path}: {e}")
+            if conn:
+                conn.rollback()
+            raise
+        finally:
+            if conn:
+                conn.close()
+
+    def _initialize_db(self) -> bool:
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.executescript(CREATE_TABLES_SQL)
+                cursor.executescript(CREATE_INDEXES_SQL)
+                conn.commit()
+                return True
+        except sqlite3.Error as e:
+            logger.critical(f"Критическая ошибка инициализации БД: {e}")
+            return False
+
+    # ────────────────────────────────
+    # Папки
+    # ────────────────────────────────
+
+    def add_folder(self, folder_path: str, parent_id: Optional[int] = None) -> Optional[int]:
+        if not folder_path:
+            return None
+        normalized_path = _normalize(folder_path)
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    "INSERT INTO folders (path, parent_id) VALUES (?, ?) "
+                    "ON CONFLICT(path) DO UPDATE SET parent_id=excluded.parent_id RETURNING id",
+                    (normalized_path, parent_id)
+                )
+                result = cursor.fetchone()
+                if not result:
+                    cursor.execute("SELECT id FROM folders WHERE path = ?", (normalized_path,))
+                    result = cursor.fetchone()
+                conn.commit()
+                return result[0] if result else None
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка при добавлении папки {normalized_path}: {e}")
+            return None
+
+    def remove_folder_by_path(self, folder_path: str) -> bool:
+        normalized_path = _normalize(folder_path)
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute("DELETE FROM folders WHERE path = ?", (normalized_path,))
+                conn.commit()
+                return cursor.rowcount > 0
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка при удалении папки {normalized_path}: {e}")
+            return False
+
+    # ────────────────────────────────
+    # Файлы
+    # ────────────────────────────────
+
+    def register_file(self, folder_id: int, filename: str, filepath: str,
+                      size: int, mtime: float) -> Optional[int]:
+        if not filename or not filepath:
+            return None
+        normalized_filepath = _normalize(filepath)
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    """
+                    INSERT INTO files (folder_id, filename, filepath, size, mtime)
+                    VALUES (?, ?, ?, ?, ?)
+                    ON CONFLICT(filepath) DO UPDATE SET
+                        size = excluded.size,
+                        mtime = excluded.mtime,
+                        folder_id = excluded.folder_id
+                    RETURNING id
+                    """,
+                    (folder_id, filename, normalized_filepath, size, mtime)
+                )
+                result = cursor.fetchone()
+                conn.commit()
+                return result[0] if result else None
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка регистрации файла {normalized_filepath}: {e}")
+            return None
+
+    def remove_file_by_path(self, filepath: str) -> bool:
+        normalized_filepath = _normalize(filepath)
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute("DELETE FROM files WHERE filepath = ?", (normalized_filepath,))
+                conn.commit()
+                return cursor.rowcount > 0
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка удаления файла {normalized_filepath}: {e}")
+            return False
+
+    # ────────────────────────────────
+    # Метаданные
+    # ────────────────────────────────
+
+    def save_metadata(self, file_id: int, meta_payload: Dict[str, Any]) -> bool:
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    """
+                    INSERT INTO metadata (
+                        file_id, prompt_json, workflow_json, positive_prompt,
+                        negative_prompt, seed, model_name, sampler, steps, cfg
+                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+                    ON CONFLICT(file_id) DO UPDATE SET
+                        prompt_json = excluded.prompt_json,
+                        workflow_json = excluded.workflow_json,
+                        positive_prompt = excluded.positive_prompt,
+                        negative_prompt = excluded.negative_prompt,
+                        seed = excluded.seed,
+                        model_name = excluded.model_name,
+                        sampler = excluded.sampler,
+                        steps = excluded.steps,
+                        cfg = excluded.cfg
+                    """,
+                    (
+                        file_id,
+                        meta_payload.get("prompt_json"),
+                        meta_payload.get("workflow_json"),
+                        meta_payload.get("positive_prompt"),
+                        meta_payload.get("negative_prompt"),
+                        meta_payload.get("seed"),
+                        meta_payload.get("model_name"),
+                        meta_payload.get("sampler"),
+                        meta_payload.get("steps"),
+                        meta_payload.get("cfg")
+                    )
+                )
+                conn.commit()
+                return True
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка сохранения метаданных файла ID {file_id}: {e}")
+            return False
+
+    # ────────────────────────────────
+    # Выборки
+    # ────────────────────────────────
+
+    def get_files_in_folder(self, folder_path: str, search_query: str = "",
+                            favorites_only: bool = False,
+                            tag: Optional[str] = None) -> List[Dict[str, Any]]:
+        """Файлы папки с фильтрами: поиск, избранное, тег."""
+        normalized_path = _normalize(folder_path)
+        sql = """
+            SELECT f.id, f.filename, f.filepath, f.rating, f.favorite,
+                   CASE WHEN m.workflow_json IS NOT NULL AND m.workflow_json != ''
+                        THEN 1 ELSE 0 END as has_workflow
+            FROM files f
+            JOIN folders fo ON f.folder_id = fo.id
+            LEFT JOIN metadata m ON f.id = m.file_id
+            WHERE fo.path = ?
+        """
+        params: List[Any] = [normalized_path]
+
+        if search_query:
+            like = f"%{search_query}%"
+            sql += " AND (f.filename LIKE ? OR m.positive_prompt LIKE ? OR m.negative_prompt LIKE ?)"
+            params.extend([like, like, like])
+
+        if favorites_only:
+            sql += " AND f.favorite = 1"
+
+        if tag:
+            sql += """ AND f.id IN (
+                SELECT ft.file_id FROM file_tags ft
+                JOIN tags t ON ft.tag_id = t.id WHERE t.name = ?
+            )"""
+            params.append(tag)
+
+        sql += " ORDER BY f.filename ASC"
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute(sql, params)
+                return [dict(row) for row in cursor.fetchall()]
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка получения файлов в папке {normalized_path}: {e}")
+            return []
+
+    def get_file_details(self, file_id: int) -> Optional[Dict[str, Any]]:
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    """
+                    SELECT f.id, f.filename, f.filepath, f.rating, f.favorite,
+                           m.prompt_json, m.workflow_json, m.positive_prompt, m.negative_prompt,
+                           m.seed, m.model_name, m.sampler, m.steps, m.cfg
+                    FROM files f
+                    LEFT JOIN metadata m ON f.id = m.file_id
+                    WHERE f.id = ?
+                    """,
+                    (file_id,)
+                )
+                row = cursor.fetchone()
+                if not row:
+                    return None
+                details = dict(row)
+                details["tags"] = self.get_file_tags(file_id)
+                return details
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка получения деталей файла ID {file_id}: {e}")
+            return None
+
+    def update_file_details(self, file_id: int, positive: str, negative: str, rating: int) -> bool:
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute("UPDATE files SET rating = ? WHERE id = ?", (rating, file_id))
+                cursor.execute(
+                    "UPDATE metadata SET positive_prompt = ?, negative_prompt = ? WHERE file_id = ?",
+                    (positive, negative, file_id)
+                )
+                conn.commit()
+                return True
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка сохранения правок для файла ID {file_id}: {e}")
+            return False
+
+    def update_json_field(self, file_id: int, field: str, json_text: Optional[str]) -> bool:
+        """Обновляет workflow_json или prompt_json файла.
+        При изменении prompt_json перечитывает производные параметры генерации."""
+        if field not in ("workflow_json", "prompt_json"):
+            logger.error(f"Недопустимое JSON-поле: {field}")
+            return False
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                # Гарантируем существование строки metadata
+                cursor.execute(
+                    "INSERT OR IGNORE INTO metadata (file_id) VALUES (?)", (file_id,))
+                cursor.execute(
+                    f"UPDATE metadata SET {field} = ? WHERE file_id = ?",
+                    (json_text or None, file_id)
+                )
+                if field == "prompt_json":
+                    # Пересчитываем параметры из нового prompt-графа
+                    from comfy_studio.parsers.metadata_parser import MetadataParser
+                    params = MetadataParser.parse_comfy_parameters(json_text)
+                    cursor.execute(
+                        """
+                        UPDATE metadata SET
+                            positive_prompt = ?, negative_prompt = ?, seed = ?,
+                            model_name = ?, sampler = ?, steps = ?, cfg = ?
+                        WHERE file_id = ?
+                        """,
+                        (
+                            params["positive_prompt"], params["negative_prompt"],
+                            params["seed"], params["model_name"], params["sampler"],
+                            params["steps"], params["cfg"], file_id
+                        )
+                    )
+                conn.commit()
+                return True
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка сохранения {field} для файла ID {file_id}: {e}")
+            return False
+
+    # ────────────────────────────────
+    # Избранное
+    # ────────────────────────────────
+
+    def set_favorite(self, file_id: int, favorite: bool) -> bool:
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute("UPDATE files SET favorite = ? WHERE id = ?",
+                               (1 if favorite else 0, file_id))
+                conn.commit()
+                return cursor.rowcount > 0
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка изменения избранного для ID {file_id}: {e}")
+            return False
+
+    # ────────────────────────────────
+    # Теги
+    # ────────────────────────────────
+
+    def add_tag_to_file(self, file_id: int, tag_name: str) -> bool:
+        tag_name = tag_name.strip().lower()
+        if not tag_name:
+            return False
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute("INSERT OR IGNORE INTO tags (name) VALUES (?)", (tag_name,))
+                cursor.execute("SELECT id FROM tags WHERE name = ?", (tag_name,))
+                tag_id = cursor.fetchone()[0]
+                cursor.execute(
+                    "INSERT OR IGNORE INTO file_tags (file_id, tag_id) VALUES (?, ?)",
+                    (file_id, tag_id)
+                )
+                conn.commit()
+                return True
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка добавления тега '{tag_name}' к файлу ID {file_id}: {e}")
+            return False
+
+    def remove_tag_from_file(self, file_id: int, tag_name: str) -> bool:
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    """
+                    DELETE FROM file_tags WHERE file_id = ?
+                      AND tag_id = (SELECT id FROM tags WHERE name = ?)
+                    """,
+                    (file_id, tag_name.strip().lower())
+                )
+                # Подчищаем теги-сироты
+                cursor.execute(
+                    "DELETE FROM tags WHERE id NOT IN (SELECT DISTINCT tag_id FROM file_tags)"
+                )
+                conn.commit()
+                return True
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка удаления тега '{tag_name}' у файла ID {file_id}: {e}")
+            return False
+
+    def get_file_tags(self, file_id: int) -> List[str]:
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute(
+                    """
+                    SELECT t.name FROM tags t
+                    JOIN file_tags ft ON t.id = ft.tag_id
+                    WHERE ft.file_id = ? ORDER BY t.name
+                    """,
+                    (file_id,)
+                )
+                return [row[0] for row in cursor.fetchall()]
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка получения тегов файла ID {file_id}: {e}")
+            return []
+
+    def get_all_tags(self) -> List[str]:
+        try:
+            with self._get_connection() as conn:
+                cursor = conn.cursor()
+                cursor.execute("SELECT name FROM tags ORDER BY name")
+                return [row[0] for row in cursor.fetchall()]
+        except sqlite3.Error as e:
+            logger.error(f"Ошибка получения списка тегов: {e}")
+            return []
@@ -0,0 +1,60 @@
+"""SQL-схема базы данных Comfy-Studio."""
+
+PRAGMA_FOREIGN_KEYS = "PRAGMA foreign_keys = ON;"
+
+CREATE_TABLES_SQL = """
+CREATE TABLE IF NOT EXISTS folders (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    path TEXT UNIQUE NOT NULL,
+    parent_id INTEGER,
+    last_scanned TIMESTAMP,
+    FOREIGN KEY (parent_id) REFERENCES folders(id) ON DELETE CASCADE
+);
+
+CREATE TABLE IF NOT EXISTS files (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    folder_id INTEGER NOT NULL,
+    filename TEXT NOT NULL,
+    filepath TEXT UNIQUE NOT NULL,
+    size INTEGER NOT NULL,
+    mtime REAL NOT NULL,
+    rating INTEGER DEFAULT 0 CHECK (rating >= 0 AND rating <= 5),
+    favorite INTEGER DEFAULT 0 CHECK (favorite IN (0, 1)),
+    FOREIGN KEY (folder_id) REFERENCES folders(id) ON DELETE CASCADE
+);
+
+CREATE TABLE IF NOT EXISTS metadata (
+    file_id INTEGER PRIMARY KEY,
+    prompt_json TEXT,
+    workflow_json TEXT,
+    positive_prompt TEXT,
+    negative_prompt TEXT,
+    seed INTEGER,
+    model_name TEXT,
+    sampler TEXT,
+    steps INTEGER,
+    cfg REAL,
+    FOREIGN KEY (file_id) REFERENCES files(id) ON DELETE CASCADE
+);
+
+CREATE TABLE IF NOT EXISTS tags (
+    id INTEGER PRIMARY KEY AUTOINCREMENT,
+    name TEXT UNIQUE NOT NULL
+);
+
+CREATE TABLE IF NOT EXISTS file_tags (
+    file_id INTEGER NOT NULL,
+    tag_id INTEGER NOT NULL,
+    PRIMARY KEY (file_id, tag_id),
+    FOREIGN KEY (file_id) REFERENCES files(id) ON DELETE CASCADE,
+    FOREIGN KEY (tag_id) REFERENCES tags(id) ON DELETE CASCADE
+);
+"""
+
+CREATE_INDEXES_SQL = """
+CREATE INDEX IF NOT EXISTS idx_folders_path ON folders(path);
+CREATE INDEX IF NOT EXISTS idx_files_filepath ON files(filepath);
+CREATE INDEX IF NOT EXISTS idx_files_folder_id ON files(folder_id);
+CREATE INDEX IF NOT EXISTS idx_files_favorite ON files(favorite);
+CREATE INDEX IF NOT EXISTS idx_tags_name ON tags(name);
+"""