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>
This commit is contained in:
dinlo
@@ -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 []
|
||||
Reference in New Issue
Block a user