530 lines
19 KiB
Python
530 lines
19 KiB
Python
"""Database helper functions for SWIPER.
|
||
All DB reads / writes are centralised here so the rest of the
|
||
codebase never needs to know SQL details.
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import os
|
||
import sqlite3
|
||
import json
|
||
import time
|
||
from typing import List, Dict, Any
|
||
|
||
# Optional progress bar
|
||
try:
|
||
from tqdm import tqdm # type: ignore
|
||
except ImportError: # pragma: no cover
|
||
tqdm = None
|
||
|
||
from PIL import Image
|
||
|
||
from config import DB_PATH, IMAGE_DIRS, find_image_file, NSFW_KEYWORDS
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Core helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def _get_conn() -> sqlite3.Connection:
|
||
"""Return a new connection with row access by column name."""
|
||
conn = sqlite3.connect(DB_PATH)
|
||
conn.row_factory = sqlite3.Row
|
||
return conn
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Schema setup & sync helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def init_db() -> None:
|
||
"""Create missing tables and perform schema migrations if necessary."""
|
||
conn = _get_conn()
|
||
cur = conn.cursor()
|
||
|
||
# Ensure image_selections table exists
|
||
cur.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS image_selections (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
image_path TEXT NOT NULL UNIQUE,
|
||
action TEXT NOT NULL,
|
||
timestamp INTEGER NOT NULL
|
||
)
|
||
"""
|
||
)
|
||
|
||
# Ensure image_metadata table exists
|
||
cur.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS image_metadata (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
path TEXT NOT NULL UNIQUE,
|
||
resolution_x INTEGER NOT NULL,
|
||
resolution_y INTEGER NOT NULL,
|
||
name TEXT NOT NULL,
|
||
orientation TEXT NOT NULL,
|
||
creation_date INTEGER NOT NULL,
|
||
prompt_data TEXT,
|
||
actioned TEXT DEFAULT NULL,
|
||
nsfw INTEGER NOT NULL DEFAULT 0
|
||
)
|
||
"""
|
||
)
|
||
|
||
# Ensure prompt_details table exists and new columns
|
||
cur.execute(
|
||
"""
|
||
CREATE TABLE IF NOT EXISTS prompt_details (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
image_path TEXT NOT NULL UNIQUE,
|
||
model_name TEXT,
|
||
positive_prompt TEXT,
|
||
negative_prompt TEXT,
|
||
sampler TEXT,
|
||
steps INTEGER,
|
||
cfg_scale REAL,
|
||
seed INTEGER,
|
||
clip_skip INTEGER,
|
||
loras TEXT,
|
||
textual_inversions TEXT,
|
||
other_parameters TEXT,
|
||
FOREIGN KEY (image_path) REFERENCES image_metadata (path)
|
||
)
|
||
"""
|
||
)
|
||
|
||
# ------------------------------------------------------------------
|
||
# Ensure newer optional columns exist / are migrated for prompt_details
|
||
# ------------------------------------------------------------------
|
||
cur.execute("PRAGMA table_info(prompt_details)")
|
||
pd_columns = {row['name'] for row in cur.fetchall()}
|
||
if 'loras' not in pd_columns:
|
||
cur.execute('ALTER TABLE prompt_details ADD COLUMN loras TEXT')
|
||
if 'textual_inversions' not in pd_columns:
|
||
cur.execute('ALTER TABLE prompt_details ADD COLUMN textual_inversions TEXT')
|
||
|
||
# Check and migrate the 'actioned' column
|
||
cur.execute("PRAGMA table_info(image_metadata)")
|
||
columns = {row['name']: row['type'] for row in cur.fetchall()}
|
||
|
||
# ------------------------------------------------------------------
|
||
# Ensure newer optional columns exist / are migrated
|
||
# ------------------------------------------------------------------
|
||
if 'actioned' not in columns:
|
||
# Add the column if it doesn't exist
|
||
cur.execute('ALTER TABLE image_metadata ADD COLUMN actioned TEXT DEFAULT NULL')
|
||
elif columns.get('actioned') != 'TEXT':
|
||
# Migrate the column if it has the wrong type (e.g., INTEGER)
|
||
print("Migrating 'actioned' column to TEXT type...")
|
||
cur.execute('ALTER TABLE image_metadata RENAME TO image_metadata_old')
|
||
cur.execute(
|
||
"""
|
||
CREATE TABLE image_metadata (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
path TEXT NOT NULL UNIQUE,
|
||
resolution_x INTEGER NOT NULL,
|
||
resolution_y INTEGER NOT NULL,
|
||
name TEXT NOT NULL,
|
||
orientation TEXT NOT NULL,
|
||
creation_date INTEGER NOT NULL,
|
||
prompt_data TEXT,
|
||
actioned TEXT DEFAULT NULL,
|
||
nsfw INTEGER DEFAULT 0
|
||
)
|
||
"""
|
||
)
|
||
# Copy data, omitting the old 'actioned' column
|
||
cur.execute(
|
||
"""
|
||
INSERT INTO image_metadata (id, path, resolution_x, resolution_y, name, orientation, creation_date, prompt_data, nsfw)
|
||
SELECT id, path, resolution_x, resolution_y, name, orientation, creation_date, prompt_data, 0
|
||
FROM image_metadata_old
|
||
"""
|
||
)
|
||
cur.execute('DROP TABLE image_metadata_old')
|
||
print("Migration complete.")
|
||
|
||
# Ensure nsfw column exists for older installations
|
||
if 'nsfw' not in columns:
|
||
cur.execute('ALTER TABLE image_metadata ADD COLUMN nsfw INTEGER DEFAULT 0')
|
||
# Migrate the column if it has the wrong type (e.g., INTEGER)
|
||
print("Migrating 'actioned' column to TEXT type...")
|
||
cur.execute('ALTER TABLE image_metadata RENAME TO image_metadata_old')
|
||
cur.execute(
|
||
"""
|
||
CREATE TABLE image_metadata (
|
||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||
path TEXT NOT NULL UNIQUE,
|
||
resolution_x INTEGER NOT NULL,
|
||
resolution_y INTEGER NOT NULL,
|
||
name TEXT NOT NULL,
|
||
orientation TEXT NOT NULL,
|
||
creation_date INTEGER NOT NULL,
|
||
prompt_data TEXT,
|
||
actioned TEXT DEFAULT NULL,
|
||
nsfw INTEGER DEFAULT 0
|
||
)
|
||
"""
|
||
)
|
||
# Copy data, omitting the old 'actioned' column
|
||
cur.execute(
|
||
"""
|
||
INSERT INTO image_metadata (id, path, resolution_x, resolution_y, name, orientation, creation_date, prompt_data)
|
||
SELECT id, path, resolution_x, resolution_y, name, orientation, creation_date, prompt_data
|
||
FROM image_metadata_old
|
||
"""
|
||
)
|
||
cur.execute('DROP TABLE image_metadata_old')
|
||
print("Migration complete.")
|
||
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
|
||
def parse_and_store_prompt_details(image_path: str, prompt_data: str) -> None:
|
||
"""Parse prompt data and store it in the prompt_details table."""
|
||
if not prompt_data:
|
||
return
|
||
|
||
conn = _get_conn()
|
||
cur = conn.cursor()
|
||
|
||
# Simple parsing logic, can be expanded
|
||
details = {
|
||
"image_path": image_path,
|
||
"positive_prompt": "",
|
||
"negative_prompt": "",
|
||
"model_name": None,
|
||
"sampler": None,
|
||
"steps": None,
|
||
"cfg_scale": None,
|
||
"seed": None,
|
||
"clip_skip": None,
|
||
"loras": None,
|
||
"textual_inversions": None,
|
||
"other_parameters": "{}"
|
||
}
|
||
|
||
try:
|
||
import re
|
||
text = prompt_data
|
||
|
||
# Positive & negative prompt
|
||
neg_match = re.search(r"Negative prompt:\s*(.*?)\s*(Steps:|Sampler:|Schedule type:|CFG scale:|Seed:|Size:|Model hash:|Model:|Lora hashes:|TI:|Version:|$)", text, re.S)
|
||
if neg_match:
|
||
details["negative_prompt"] = neg_match.group(1).strip()
|
||
details["positive_prompt"] = text[:neg_match.start()].strip().rstrip(',')
|
||
else:
|
||
details["positive_prompt"] = text.split("Steps:")[0].strip()
|
||
|
||
# Key–value param pairs (e.g. "Steps: 20," )
|
||
for key, val in re.findall(r"(\w[\w ]*?):\s*([^,\n]+)", text):
|
||
val = val.strip().rstrip(',')
|
||
if key == "Model":
|
||
details["model_name"] = val
|
||
elif key == "Sampler":
|
||
details["sampler"] = val
|
||
elif key == "Steps":
|
||
details["steps"] = int(val or 0)
|
||
elif key == "CFG scale":
|
||
details["cfg_scale"] = float(val or 0)
|
||
elif key == "Seed":
|
||
details["seed"] = int(val or 0)
|
||
elif key == "CLIP skip":
|
||
details["clip_skip"] = int(val or 0)
|
||
|
||
# Store other params as JSON
|
||
std_keys = {"Model", "Sampler", "Steps", "CFG scale", "Seed", "CLIP skip"}
|
||
other_params = {k: v.strip().rstrip(',') for k, v in re.findall(r"(\w[\w ]*?):\s*([^,\n]+)", text) if k not in std_keys}
|
||
details["other_parameters"] = json.dumps(other_params)
|
||
|
||
# Extract Loras and Textual Inversions (TIs)
|
||
lora_match = re.search(r"Lora hashes:\s*\"?([^\"]+)\"?", text, re.I)
|
||
if lora_match:
|
||
details["loras"] = lora_match.group(1).strip()
|
||
|
||
ti_match = re.search(r"TI:\s*\"?([^\"]+)\"?", text, re.I)
|
||
if ti_match:
|
||
details["textual_inversions"] = ti_match.group(1).strip()
|
||
|
||
# Fallback: look for <lora:...> tags if no explicit Lora hashes line
|
||
if not details.get("loras"):
|
||
tag_matches = re.findall(r"<lora:([^>]+)>", text, re.I)
|
||
if tag_matches:
|
||
details["loras"] = ", ".join(tag_matches)
|
||
|
||
# Extract TI keywords (<token> patterns) if still empty
|
||
if not details.get("textual_inversions"):
|
||
ti_tokens = re.findall(r"<([\w-]{3,32})>", details["positive_prompt"]) # crude heuristic
|
||
if ti_tokens:
|
||
details["textual_inversions"] = ", ".join(sorted(set(ti_tokens)))
|
||
|
||
# Final clean-up
|
||
if details.get("textual_inversions"):
|
||
details["textual_inversions"] = details["textual_inversions"].strip().strip('"').rstrip(',').strip()
|
||
|
||
except Exception as e:
|
||
print(f"Error parsing prompt for {image_path}: {e}")
|
||
# still insert with what we have
|
||
|
||
cur.execute(
|
||
"""
|
||
INSERT OR REPLACE INTO prompt_details (image_path, model_name, positive_prompt, negative_prompt, sampler, steps, cfg_scale, seed, clip_skip, loras, textual_inversions, other_parameters)
|
||
VALUES (:image_path, :model_name, :positive_prompt, :negative_prompt, :sampler, :steps, :cfg_scale, :seed, :clip_skip, :loras, :textual_inversions, :other_parameters)
|
||
""",
|
||
details
|
||
)
|
||
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
|
||
def sync_image_database() -> None:
|
||
"""Scan the image folder and ensure every image is present in image_metadata."""
|
||
|
||
print("Syncing image database…", flush=True)
|
||
conn = _get_conn()
|
||
cur = conn.cursor()
|
||
|
||
# Already-known images
|
||
cur.execute("SELECT path FROM image_metadata")
|
||
known = {row[0] for row in cur.fetchall()}
|
||
|
||
# Images on disk (expects <resolution>/<filename>) across all configured dirs
|
||
disk_images: set[str] = set()
|
||
for base in IMAGE_DIRS:
|
||
for res in [d for d in os.listdir(base) if os.path.isdir(os.path.join(base, d))]:
|
||
res_dir = os.path.join(base, res)
|
||
for file in os.listdir(res_dir):
|
||
if file.lower().endswith((".png", ".jpg", ".jpeg")):
|
||
disk_images.add(f"{res}/{file}")
|
||
|
||
new_images = disk_images - known
|
||
if not new_images:
|
||
print("Database already up to date.")
|
||
conn.close()
|
||
return
|
||
|
||
total_new_images = len(new_images)
|
||
|
||
# Choose iterator with progress bar if available
|
||
image_iter = (
|
||
tqdm(sorted(new_images), desc="Syncing images") if tqdm else sorted(new_images)
|
||
)
|
||
|
||
processed_count = 0
|
||
rows = []
|
||
for rel_path in image_iter:
|
||
res, filename = rel_path.split("/", 1)
|
||
abs_path = find_image_file(rel_path)
|
||
if not abs_path:
|
||
processed_count += 1
|
||
if not tqdm and (processed_count % 100 == 0 or processed_count == total_new_images):
|
||
percentage = (processed_count / total_new_images) * 100
|
||
print(
|
||
f"Processed {processed_count} / {total_new_images} images ({percentage:.2f}%)",
|
||
flush=True,
|
||
)
|
||
continue
|
||
try:
|
||
with Image.open(abs_path) as img:
|
||
w, h = img.size
|
||
orient = (
|
||
"square" if w == h else "landscape" if w > h else "portrait"
|
||
)
|
||
prompt = None
|
||
if img.format == "PNG":
|
||
prompt = img.info.get("parameters") or img.info.get("Parameters")
|
||
ts = int(os.path.getmtime(abs_path))
|
||
# Detect NSFW based on prompt keywords
|
||
nsfw_flag = 0
|
||
if prompt:
|
||
lower_prompt = prompt.lower()
|
||
nsfw_flag = 1 if any(k.lower() in lower_prompt for k in NSFW_KEYWORDS) else 0
|
||
rows.append((rel_path, w, h, filename, orient, ts, prompt, nsfw_flag))
|
||
except Exception as exc:
|
||
print(f"Failed reading {abs_path}: {exc}")
|
||
|
||
if rows:
|
||
cur.executemany(
|
||
"""
|
||
INSERT INTO image_metadata (path, resolution_x, resolution_y, name, orientation, creation_date, prompt_data, nsfw)
|
||
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
|
||
""",
|
||
rows,
|
||
)
|
||
conn.commit()
|
||
print(f"Inserted {len(rows)} new images.")
|
||
|
||
# Now, parse and store prompt details for the new images
|
||
for row in rows:
|
||
# row structure: (path, w, h, name, orientation, ts, prompt, nsfw_flag)
|
||
rel_path = row[0]
|
||
prompt = row[6]
|
||
if prompt:
|
||
parse_and_store_prompt_details(rel_path, prompt)
|
||
|
||
conn.close()
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Selection helpers
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def add_selection(image_path: str, action: str) -> None:
|
||
"""Add or update a selection and the metadata actioned status."""
|
||
conn = _get_conn()
|
||
cur = conn.cursor()
|
||
timestamp = int(time.time())
|
||
|
||
# Upsert the selection
|
||
cur.execute(
|
||
"""
|
||
INSERT INTO image_selections (image_path, action, timestamp)
|
||
VALUES (?, ?, ?)
|
||
ON CONFLICT(image_path) DO UPDATE SET
|
||
action = excluded.action,
|
||
timestamp = excluded.timestamp
|
||
""",
|
||
(image_path, action, timestamp),
|
||
)
|
||
|
||
# Update the metadata table with the action name
|
||
cur.execute(
|
||
"UPDATE image_metadata SET actioned = ? WHERE path = ?", (action, image_path)
|
||
)
|
||
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
|
||
def get_selections() -> List[Dict[str, Any]]:
|
||
"""Return selection list with metadata and parsed prompt details."""
|
||
conn = _get_conn()
|
||
cur = conn.cursor()
|
||
cur.execute(
|
||
"""
|
||
SELECT
|
||
sel.id, sel.image_path, sel.action, sel.timestamp,
|
||
meta.resolution_x, meta.resolution_y, meta.orientation,
|
||
meta.nsfw,
|
||
meta.creation_date, meta.prompt_data, meta.name,
|
||
pd.model_name, pd.positive_prompt, pd.negative_prompt,
|
||
pd.sampler, pd.steps, pd.cfg_scale, pd.seed, pd.clip_skip,
|
||
pd.loras, pd.textual_inversions,
|
||
pd.other_parameters
|
||
FROM image_selections sel
|
||
LEFT JOIN image_metadata meta ON sel.image_path = meta.path
|
||
LEFT JOIN prompt_details pd ON sel.image_path = pd.image_path
|
||
ORDER BY sel.timestamp DESC
|
||
"""
|
||
)
|
||
rows = cur.fetchall()
|
||
conn.close()
|
||
|
||
results: List[Dict[str, Any]] = []
|
||
for row in rows:
|
||
item: Dict[str, Any] = {k: row[k] for k in row.keys()}
|
||
|
||
# Parse other_parameters if it exists and is a string
|
||
other_params_str = item.get("other_parameters")
|
||
if isinstance(other_params_str, str):
|
||
try:
|
||
item["other_parameters"] = json.loads(other_params_str)
|
||
except json.JSONDecodeError:
|
||
item["other_parameters"] = {}
|
||
|
||
# Derive resolution & orientation if missing (back-compat)
|
||
if not item.get("resolution"):
|
||
try:
|
||
path_part = item["image_path"].lstrip("/images/")
|
||
item["resolution"] = path_part.split("/")[0]
|
||
except Exception:
|
||
item["resolution"] = "unknown"
|
||
if not item.get("orientation"):
|
||
try:
|
||
abs_path = find_image_file(item["image_path"].lstrip("/images/"))
|
||
if abs_path:
|
||
with Image.open(abs_path) as img:
|
||
w, h = img.size
|
||
item["orientation"] = (
|
||
"square" if w == h else "landscape" if w > h else "portrait"
|
||
)
|
||
except Exception:
|
||
item["orientation"] = "unknown"
|
||
results.append(item)
|
||
return results
|
||
|
||
|
||
def update_selection(selection_id: int, action: str) -> bool:
|
||
"""Update an existing selection and the corresponding metadata."""
|
||
conn = _get_conn()
|
||
cur = conn.cursor()
|
||
|
||
# First, get the image_path for the given selection_id
|
||
cur.execute("SELECT image_path FROM image_selections WHERE id = ?", (selection_id,))
|
||
row = cur.fetchone()
|
||
if not row:
|
||
conn.close()
|
||
return False
|
||
|
||
image_path = row['image_path']
|
||
|
||
# Update the action in the image_selections table
|
||
cur.execute(
|
||
"UPDATE image_selections SET action = ?, timestamp = ? WHERE id = ?",
|
||
(action, int(time.time()), selection_id),
|
||
)
|
||
changed = cur.rowcount > 0
|
||
|
||
if changed:
|
||
# Also update the actioned column in the image_metadata table
|
||
cur.execute(
|
||
"UPDATE image_metadata SET actioned = ? WHERE path = ?", (action, image_path)
|
||
)
|
||
|
||
conn.commit()
|
||
conn.close()
|
||
return changed
|
||
|
||
|
||
def delete_selection(selection_id: int) -> None:
|
||
"""Delete a selection and reset the metadata actioned link."""
|
||
conn = _get_conn()
|
||
cur = conn.cursor()
|
||
|
||
# Find the image path before deleting
|
||
cur.execute("SELECT image_path FROM image_selections WHERE id = ?", (selection_id,))
|
||
row = cur.fetchone()
|
||
if not row:
|
||
conn.close()
|
||
return # Or raise error
|
||
|
||
image_path = row['image_path']
|
||
|
||
# Delete the selection
|
||
cur.execute("DELETE FROM image_selections WHERE id = ?", (selection_id,))
|
||
|
||
# Update the metadata table
|
||
cur.execute("UPDATE image_metadata SET actioned = NULL WHERE path = ?", (image_path,))
|
||
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
|
||
def reset_database() -> None:
|
||
"""Clear all selections and reset all actioned links in metadata."""
|
||
conn = _get_conn()
|
||
cur = conn.cursor()
|
||
cur.execute("DELETE FROM image_selections")
|
||
cur.execute("UPDATE image_metadata SET actioned = NULL")
|
||
conn.commit()
|
||
conn.close()
|
||
|
||
|
||
def reset_database() -> None:
|
||
"""Clear all selections and reset all actioned links in metadata."""
|
||
conn = _get_conn()
|
||
cur = conn.cursor()
|
||
cur.execute("DELETE FROM image_selections")
|
||
cur.execute("UPDATE image_metadata SET actioned = NULL")
|
||
conn.commit()
|
||
conn.close()
|