1195 lines
44 KiB
Python
1195 lines
44 KiB
Python
import psycopg2
|
|
from psycopg2 import pool as _pgpool
|
|
import threading
|
|
import json
|
|
import time
|
|
import random
|
|
|
|
DB_CONFIG = {
|
|
"host": "192.168.1.160",
|
|
"port": 5433,
|
|
"dbname": "dv",
|
|
"user": "dev",
|
|
"password": "dev"
|
|
}
|
|
|
|
# A pooled connection is reused across requests instead of paying a fresh TCP +
|
|
# auth round-trip (~40 ms) on every single query. Under load (a generation
|
|
# running, a burst of reorders) the old open-per-call design starved the web
|
|
# threadpool and occasionally tripped Postgres' connection ceiling → 500s.
|
|
# The pool bounds connections and keeps them warm. If the pool can't be
|
|
# created or is momentarily exhausted, get_db_connection() falls back to a
|
|
# direct connect so callers never block or fail.
|
|
# Cover the web server's sync threadpool (uvicorn/anyio default = 40 workers)
|
|
# with headroom for background workers, while staying well under Postgres'
|
|
# max_connections (100).
|
|
_POOL_MIN = 2
|
|
_POOL_MAX = 48
|
|
_pool = None
|
|
_pool_lock = threading.Lock()
|
|
|
|
|
|
def _get_pool():
|
|
global _pool
|
|
if _pool is not None:
|
|
return _pool
|
|
with _pool_lock:
|
|
if _pool is None:
|
|
try:
|
|
_pool = _pgpool.ThreadedConnectionPool(
|
|
_POOL_MIN, _POOL_MAX, **DB_CONFIG)
|
|
except Exception as e:
|
|
print(f"[db] pool init failed, using direct connections: {e}")
|
|
_pool = False # sentinel: don't retry on every call
|
|
return _pool
|
|
|
|
|
|
def get_db_connection():
|
|
"""Return a live DB connection, preferring the pool.
|
|
|
|
Always pair with _put_db_connection() (the existing finally: conn.close()
|
|
callsites are rewritten to call it) so pooled connections are returned
|
|
rather than dropped.
|
|
"""
|
|
p = _get_pool()
|
|
if p:
|
|
try:
|
|
conn = p.getconn()
|
|
# Guard against a stale/dead pooled connection.
|
|
if getattr(conn, "closed", 0):
|
|
try:
|
|
p.putconn(conn, close=True)
|
|
except Exception:
|
|
pass
|
|
else:
|
|
return conn
|
|
except _pgpool.PoolError:
|
|
# Pool momentarily exhausted — expected under burst; fall back to a
|
|
# direct connection silently rather than blocking or failing.
|
|
pass
|
|
except Exception as e:
|
|
print(f"[db] pool getconn failed, direct connect: {e}")
|
|
return psycopg2.connect(**DB_CONFIG)
|
|
|
|
|
|
def _put_db_connection(conn):
|
|
"""Return a connection to the pool (rolling back any open txn) or close it.
|
|
|
|
Safe for both pooled and direct/fallback connections: putconn raises for a
|
|
connection the pool doesn't own, in which case we just close it.
|
|
"""
|
|
if conn is None:
|
|
return
|
|
p = _pool
|
|
try:
|
|
if p:
|
|
try:
|
|
conn.rollback() # clear any aborted/idle-in-txn state before reuse
|
|
except Exception:
|
|
pass
|
|
p.putconn(conn)
|
|
return
|
|
except Exception:
|
|
pass
|
|
try:
|
|
conn.close()
|
|
except Exception:
|
|
pass
|
|
|
|
def migrate_schema():
|
|
"""Add new columns to person table if they don't exist. Safe to call repeatedly."""
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
# Create prompt table first
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS prompt (
|
|
id SERIAL PRIMARY KEY,
|
|
type TEXT NOT NULL,
|
|
prompt_text TEXT NOT NULL,
|
|
metadata JSONB DEFAULT '{}'::jsonb,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
CONSTRAINT unique_prompt_type_text UNIQUE (type, prompt_text)
|
|
)
|
|
""")
|
|
# Create turntable table
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS turntable (
|
|
group_id TEXT PRIMARY KEY,
|
|
preferred_filename TEXT,
|
|
source_image TEXT,
|
|
seed INTEGER,
|
|
n_views INTEGER,
|
|
steps INTEGER,
|
|
video_path TEXT,
|
|
completed BOOLEAN DEFAULT FALSE,
|
|
started_at DOUBLE PRECISION,
|
|
completed_at DOUBLE PRECISION,
|
|
angles TEXT, -- JSON list of angles
|
|
views TEXT -- JSON representation of views dict
|
|
)
|
|
""")
|
|
for sql in [
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS prompt TEXT",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS pose TEXT",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS sort_order INTEGER",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS group_name TEXT",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS hidden BOOLEAN DEFAULT FALSE",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS has_background BOOLEAN DEFAULT TRUE",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS source_refs TEXT",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS has_clothing BOOLEAN DEFAULT NULL",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS content_type TEXT DEFAULT 'image'",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS faceswap_source_video TEXT",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS archived BOOLEAN DEFAULT FALSE",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS face_embedding vector(512)",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS is_source BOOLEAN DEFAULT FALSE",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS pose_description TEXT",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS pose_skeleton TEXT",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS people_count INTEGER DEFAULT NULL",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS anatomical_completeness BOOLEAN DEFAULT NULL",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS facial_direction TEXT DEFAULT NULL",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS objects JSONB DEFAULT NULL",
|
|
"ALTER TABLE person ADD COLUMN IF NOT EXISTS description TEXT",
|
|
]:
|
|
cur.execute(sql)
|
|
conn.commit()
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
try:
|
|
initialize_tags_in_db()
|
|
except Exception as e:
|
|
print(f"[db] initialize_tags_in_db error: {e}")
|
|
try:
|
|
migrate_turntable_states_from_disk_to_db()
|
|
except Exception as e:
|
|
print(f"[db] migrate_turntable_states_from_disk_to_db error: {e}")
|
|
|
|
def initialize_tags_in_db():
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("""
|
|
SELECT filename, archived, hidden, is_source, sort_order, has_background, tags
|
|
FROM person
|
|
""")
|
|
rows = cur.fetchall()
|
|
for filename, archived, hidden, is_source, sort_order, has_background, tags_val in rows:
|
|
tags_list = []
|
|
if tags_val:
|
|
if isinstance(tags_val, str):
|
|
try:
|
|
tags_list = json.loads(tags_val)
|
|
except Exception:
|
|
tags_list = []
|
|
elif isinstance(tags_val, list):
|
|
tags_list = tags_val
|
|
|
|
is_archived = bool(archived) if archived is not None else False
|
|
is_hidden = bool(hidden) if hidden is not None else False
|
|
is_src = bool(is_source) if is_source is not None else False
|
|
has_bg = bool(has_background) if has_background is not None else True
|
|
|
|
def add_tag_if_missing(tag):
|
|
if tag not in tags_list:
|
|
tags_list.append(tag)
|
|
|
|
if is_archived:
|
|
add_tag_if_missing("ARCHIVED")
|
|
if "VISIBLE" in tags_list:
|
|
tags_list.remove("VISIBLE")
|
|
else:
|
|
if not is_hidden:
|
|
add_tag_if_missing("VISIBLE")
|
|
if "ARCHIVED" in tags_list:
|
|
tags_list.remove("ARCHIVED")
|
|
|
|
if is_hidden:
|
|
add_tag_if_missing("HIDDEN")
|
|
if "VISIBLE" in tags_list:
|
|
tags_list.remove("VISIBLE")
|
|
else:
|
|
if "HIDDEN" in tags_list:
|
|
tags_list.remove("HIDDEN")
|
|
|
|
if is_src:
|
|
add_tag_if_missing("SOURCE")
|
|
else:
|
|
if "SOURCE" in tags_list:
|
|
tags_list.remove("SOURCE")
|
|
|
|
if sort_order == 0:
|
|
add_tag_if_missing("GROUP_ANCHOR")
|
|
if "GROUP_MEMBER" in tags_list:
|
|
tags_list.remove("GROUP_MEMBER")
|
|
else:
|
|
add_tag_if_missing("GROUP_MEMBER")
|
|
if "GROUP_ANCHOR" in tags_list:
|
|
tags_list.remove("GROUP_ANCHOR")
|
|
|
|
if not has_bg:
|
|
add_tag_if_missing("BACKGROUND_REMOVED")
|
|
if "BACKGROUND" in tags_list:
|
|
tags_list.remove("BACKGROUND")
|
|
else:
|
|
add_tag_if_missing("BACKGROUND")
|
|
if "BACKGROUND_REMOVED" in tags_list:
|
|
tags_list.remove("BACKGROUND_REMOVED")
|
|
|
|
# ORBIT tag logic: only true orbits are tagged with ORBIT
|
|
is_orb = filename.startswith("_turntable/")
|
|
if is_orb:
|
|
add_tag_if_missing("ORBIT")
|
|
else:
|
|
if "ORBIT" in tags_list:
|
|
tags_list.remove("ORBIT")
|
|
|
|
cur.execute("UPDATE person SET tags = %s WHERE filename = %s", (json.dumps(tags_list), filename))
|
|
conn.commit()
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def upsert_person(filename, filepath=None, name=None, group_id=None, tags=None,
|
|
embedding=None, clip_description=None, prompt=None, pose=None,
|
|
sort_order=None, group_name=None, hidden=None,
|
|
has_background=None, source_refs=None, has_clothing=None,
|
|
content_type=None, faceswap_source_video=None, archived=None,
|
|
face_embedding=None, is_source=None,
|
|
pose_description=None, pose_skeleton=None,
|
|
people_count=None, anatomical_completeness=None, facial_direction=None, objects=None,
|
|
description=None):
|
|
if tags is not None:
|
|
if isinstance(tags, str):
|
|
try:
|
|
tags_list = json.loads(tags)
|
|
except Exception:
|
|
tags_list = []
|
|
elif isinstance(tags, list):
|
|
tags_list = tags
|
|
else:
|
|
tags_list = []
|
|
|
|
# Enforce that ORBIT is only on _turntable/ files
|
|
if filename.startswith("_turntable/"):
|
|
if "ORBIT" not in tags_list:
|
|
tags_list.append("ORBIT")
|
|
else:
|
|
tags_list = [t for t in tags_list if t != "ORBIT"]
|
|
tags = tags_list
|
|
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
face_embedding_str = ("[" + ",".join(map(str, face_embedding)) + "]") if face_embedding is not None else None
|
|
try:
|
|
cur.execute("""
|
|
INSERT INTO person (filename, filepath, name, group_id, tags, embedding,
|
|
clip_description, prompt, pose, sort_order, group_name, hidden,
|
|
has_background, source_refs, has_clothing,
|
|
content_type, faceswap_source_video, archived, face_embedding, is_source,
|
|
pose_description, pose_skeleton, people_count, anatomical_completeness, facial_direction, objects,
|
|
description)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT (filename) DO UPDATE
|
|
SET filepath = COALESCE(EXCLUDED.filepath, person.filepath),
|
|
name = COALESCE(EXCLUDED.name, person.name),
|
|
group_id = COALESCE(EXCLUDED.group_id, person.group_id),
|
|
tags = COALESCE(EXCLUDED.tags, person.tags),
|
|
embedding = COALESCE(EXCLUDED.embedding, person.embedding),
|
|
clip_description = COALESCE(EXCLUDED.clip_description, person.clip_description),
|
|
prompt = COALESCE(EXCLUDED.prompt, person.prompt),
|
|
pose = COALESCE(EXCLUDED.pose, person.pose),
|
|
sort_order = COALESCE(EXCLUDED.sort_order, person.sort_order),
|
|
group_name = COALESCE(EXCLUDED.group_name, person.group_name),
|
|
hidden = COALESCE(EXCLUDED.hidden, person.hidden),
|
|
has_background = COALESCE(EXCLUDED.has_background, person.has_background),
|
|
source_refs = COALESCE(EXCLUDED.source_refs, person.source_refs),
|
|
has_clothing = COALESCE(EXCLUDED.has_clothing, person.has_clothing),
|
|
content_type = COALESCE(EXCLUDED.content_type, person.content_type),
|
|
faceswap_source_video = COALESCE(EXCLUDED.faceswap_source_video, person.faceswap_source_video),
|
|
archived = COALESCE(EXCLUDED.archived, person.archived),
|
|
face_embedding = COALESCE(EXCLUDED.face_embedding, person.face_embedding),
|
|
is_source = COALESCE(EXCLUDED.is_source, person.is_source),
|
|
pose_description = COALESCE(EXCLUDED.pose_description, person.pose_description),
|
|
pose_skeleton = COALESCE(EXCLUDED.pose_skeleton, person.pose_skeleton),
|
|
people_count = COALESCE(EXCLUDED.people_count, person.people_count),
|
|
anatomical_completeness = COALESCE(EXCLUDED.anatomical_completeness, person.anatomical_completeness),
|
|
facial_direction = COALESCE(EXCLUDED.facial_direction, person.facial_direction),
|
|
objects = COALESCE(EXCLUDED.objects, person.objects),
|
|
description = COALESCE(EXCLUDED.description, person.description);
|
|
""", (filename, filepath, name, group_id,
|
|
json.dumps(tags) if tags else None,
|
|
embedding, clip_description, prompt, pose, sort_order, group_name, hidden,
|
|
has_background, source_refs, has_clothing,
|
|
content_type, faceswap_source_video, archived, face_embedding_str, is_source,
|
|
pose_description, pose_skeleton, people_count, anatomical_completeness, facial_direction,
|
|
json.dumps(objects) if objects else None,
|
|
description))
|
|
conn.commit()
|
|
# Sync after commit
|
|
try:
|
|
sync_by_filename_async(filename)
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def set_archived(filename, archived: bool):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("UPDATE person SET archived = %s WHERE filename = %s", (archived, filename))
|
|
conn.commit()
|
|
try:
|
|
sync_by_filename_async(filename)
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def set_filenames_archived(filenames, archived: bool):
|
|
if not filenames:
|
|
return []
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("UPDATE person SET archived = %s WHERE filename = ANY(%s) RETURNING filename", (archived, list(filenames)))
|
|
rows = cur.fetchall()
|
|
updated = [r[0] for r in rows]
|
|
conn.commit()
|
|
for fn in updated:
|
|
try:
|
|
sync_by_filename_async(fn)
|
|
except Exception:
|
|
pass
|
|
return updated
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def set_hidden(filename, hidden: bool):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("UPDATE person SET hidden = %s WHERE filename = %s", (hidden, filename))
|
|
conn.commit()
|
|
try:
|
|
sync_by_filename_async(filename)
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def set_person_tags(filename, tags):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("UPDATE person SET tags = %s WHERE filename = %s", (json.dumps(tags) if tags is not None else None, filename))
|
|
conn.commit()
|
|
try:
|
|
sync_by_filename_async(filename)
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def get_person(filename):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("""
|
|
SELECT name, group_id, tags, embedding, clip_description, filepath,
|
|
prompt, pose, sort_order, group_name, hidden, has_background, source_refs,
|
|
has_clothing, is_source, pose_description, pose_skeleton,
|
|
people_count, anatomical_completeness, facial_direction, objects, description
|
|
FROM person WHERE filename = %s
|
|
""", (filename,))
|
|
return cur.fetchone()
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def list_persons(include_archived=False):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
where = "" if include_archived else "WHERE archived IS NOT TRUE"
|
|
cur.execute(f"""
|
|
SELECT filename, name, group_id, clip_description,
|
|
prompt, pose, sort_order, group_name, hidden, has_background, source_refs,
|
|
has_clothing, content_type, faceswap_source_video, archived, is_source, tags,
|
|
pose_description, pose_skeleton, people_count, anatomical_completeness, facial_direction, objects, description,
|
|
filepath
|
|
FROM person
|
|
{where}
|
|
""")
|
|
return cur.fetchall()
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def search_similar(embedding, limit=10):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
embedding_str = "[" + ",".join(map(str, embedding)) + "]"
|
|
cur.execute("""
|
|
SELECT filename, name, group_id, clip_description, embedding <=> %s AS distance
|
|
FROM person
|
|
WHERE embedding IS NOT NULL
|
|
ORDER BY distance ASC
|
|
LIMIT %s;
|
|
""", (embedding_str, limit))
|
|
return cur.fetchall()
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def delete_person(filename):
|
|
person = get_person(filename)
|
|
gid = person[1] if person else None
|
|
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("DELETE FROM person WHERE filename = %s", (filename,))
|
|
conn.commit()
|
|
# Delete JSON file on disk
|
|
try:
|
|
out_dir = _get_output_dir()
|
|
base, _ = os.path.splitext(filename)
|
|
json_path = os.path.join(out_dir, base + ".json")
|
|
if os.path.exists(json_path):
|
|
os.remove(json_path)
|
|
except Exception:
|
|
pass
|
|
if gid:
|
|
try:
|
|
sync_group_to_json_async(gid)
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def delete_group(group_id):
|
|
files = get_group_files(group_id)
|
|
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("DELETE FROM person WHERE group_id = %s", (group_id,))
|
|
conn.commit()
|
|
# Delete JSON files on disk
|
|
out_dir = _get_output_dir()
|
|
for f in files:
|
|
try:
|
|
base, _ = os.path.splitext(f[0])
|
|
json_path = os.path.join(out_dir, base + ".json")
|
|
if os.path.exists(json_path):
|
|
os.remove(json_path)
|
|
except Exception:
|
|
pass
|
|
try:
|
|
group_json_path = os.path.join(out_dir, "_data", f"group_{group_id}.json")
|
|
if os.path.exists(group_json_path):
|
|
os.remove(group_json_path)
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def get_group_files(group_id):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("SELECT filename, filepath FROM person WHERE group_id = %s", (group_id,))
|
|
return cur.fetchall()
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def set_group_order(group_id, ordered_filenames):
|
|
"""Assign sort_order 0,1,2,... to filenames in the given order, resiliently."""
|
|
max_retries = 5
|
|
for attempt in range(max_retries):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
# Deterministic lock ordering to prevent deadlocks:
|
|
# Fetch and lock all unique filenames alphabetically
|
|
unique_filenames = sorted(list(set(ordered_filenames)))
|
|
if unique_filenames:
|
|
cur.execute(
|
|
"SELECT filename FROM person WHERE filename = ANY(%s) ORDER BY filename FOR UPDATE",
|
|
(unique_filenames,)
|
|
)
|
|
for idx, fname in enumerate(ordered_filenames):
|
|
cur.execute(
|
|
"UPDATE person SET sort_order = %s WHERE filename = %s",
|
|
(idx, fname)
|
|
)
|
|
conn.commit()
|
|
break
|
|
except (psycopg2.OperationalError, psycopg2.extensions.TransactionRollbackError) as e:
|
|
pgcode = getattr(e, 'pgcode', None)
|
|
if pgcode in ('40P01', '40001') or "deadlock" in str(e).lower():
|
|
try:
|
|
conn.rollback()
|
|
except Exception:
|
|
pass
|
|
if attempt == max_retries - 1:
|
|
raise
|
|
sleep_time = (0.1 * (2 ** attempt)) + random.uniform(0.01, 0.1)
|
|
time.sleep(sleep_time)
|
|
else:
|
|
raise
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
try:
|
|
sync_group_to_json_async(group_id)
|
|
for fname in ordered_filenames:
|
|
sync_person_to_json_async(fname)
|
|
except Exception:
|
|
pass
|
|
|
|
def get_group_order(group_id):
|
|
"""Return [(filename, sort_order), ...] sorted by sort_order NULLS LAST."""
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("""
|
|
SELECT filename, sort_order
|
|
FROM person
|
|
WHERE group_id = %s
|
|
ORDER BY sort_order NULLS LAST, filename
|
|
""", (group_id,))
|
|
return cur.fetchall()
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def set_group_name(group_id, name):
|
|
"""Set group_name for every file in the group."""
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("UPDATE person SET group_name = %s WHERE group_id = %s", (name, group_id))
|
|
conn.commit()
|
|
try:
|
|
sync_group_to_json_async(group_id)
|
|
files = get_group_files(group_id)
|
|
for f in files:
|
|
sync_person_to_json_async(f[0])
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def get_next_sort_order(group_id):
|
|
"""Return max(sort_order)+1 for the group, or 1 if no sorted members exist."""
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute(
|
|
"SELECT COALESCE(MAX(sort_order), 0) + 1 FROM person WHERE group_id = %s",
|
|
(group_id,)
|
|
)
|
|
return cur.fetchone()[0]
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
|
|
def get_all_group_names():
|
|
"""Return {group_id: group_name} for groups that have a name set."""
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("""
|
|
SELECT DISTINCT ON (group_id) group_id, group_name
|
|
FROM person
|
|
WHERE group_id IS NOT NULL AND group_name IS NOT NULL
|
|
""")
|
|
return {row[0]: row[1] for row in cur.fetchall()}
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
|
|
def get_face_embedding(filename):
|
|
"""Return the face_embedding as a list of floats for a filename, or None."""
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("SELECT face_embedding FROM person WHERE filename = %s", (filename,))
|
|
row = cur.fetchone()
|
|
if row and row[0] is not None:
|
|
val = row[0]
|
|
# psycopg2 without a pgvector adapter returns vectors as plain strings "[f,f,...]"
|
|
if isinstance(val, str):
|
|
return [float(x) for x in val.strip("[]").split(",")]
|
|
return list(val)
|
|
return None
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
|
|
def search_similar_face(embedding, limit=12, exclude_group_id=None):
|
|
"""Cosine search on face_embedding (stored only for *_face.png rows).
|
|
|
|
Returns [(filename, group_id, distance), ...] sorted ascending by distance.
|
|
Rows belonging to exclude_group_id are skipped so a group doesn't match itself.
|
|
"""
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
embedding_str = "[" + ",".join(map(str, embedding)) + "]"
|
|
if exclude_group_id:
|
|
cur.execute("""
|
|
SELECT filename, group_id, face_embedding <=> %s AS distance
|
|
FROM person
|
|
WHERE face_embedding IS NOT NULL
|
|
AND (group_id IS NULL OR group_id != %s)
|
|
ORDER BY distance ASC
|
|
LIMIT %s
|
|
""", (embedding_str, exclude_group_id, limit))
|
|
else:
|
|
cur.execute("""
|
|
SELECT filename, group_id, face_embedding <=> %s AS distance
|
|
FROM person
|
|
WHERE face_embedding IS NOT NULL
|
|
ORDER BY distance ASC
|
|
LIMIT %s
|
|
""", (embedding_str, limit))
|
|
return cur.fetchall()
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
|
|
def invalidate_all_metadata():
|
|
"""Reset metadata columns to NULL so they can be re-analyzed by the backfill worker."""
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("""
|
|
UPDATE person
|
|
SET people_count = NULL,
|
|
anatomical_completeness = NULL,
|
|
facial_direction = NULL,
|
|
objects = NULL,
|
|
pose_description = NULL,
|
|
pose_skeleton = NULL
|
|
WHERE NOT (filename LIKE '_turntable/%')
|
|
""")
|
|
conn.commit()
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
import os
|
|
import math
|
|
|
|
def _get_output_dir():
|
|
# Load config.json to find output_dir
|
|
config_paths = ["config.json", "tour-comfy/config.json", "../config.json"]
|
|
for p in config_paths:
|
|
if os.path.exists(p):
|
|
try:
|
|
with open(p, "r") as f:
|
|
cfg = json.load(f)
|
|
if "output_dir" in cfg:
|
|
return cfg["output_dir"]
|
|
except Exception:
|
|
pass
|
|
return "../output" # Default fallback
|
|
|
|
def sync_person_to_json(filename):
|
|
"""Write/sync a detailed .json file with the DB fields of the specified filename."""
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("""
|
|
SELECT filename, filepath, name, group_id, tags, embedding,
|
|
clip_description, prompt, pose, sort_order, group_name, hidden,
|
|
has_background, source_refs, has_clothing, content_type,
|
|
faceswap_source_video, archived, face_embedding, is_source,
|
|
pose_description, pose_skeleton, people_count, anatomical_completeness,
|
|
facial_direction, objects, description
|
|
FROM person
|
|
WHERE filename = %s
|
|
""", (filename,))
|
|
row = cur.fetchone()
|
|
if not row:
|
|
return
|
|
|
|
# Parse tags
|
|
tags_val = row[4]
|
|
tags_list = []
|
|
if tags_val:
|
|
if isinstance(tags_val, str):
|
|
try: tags_list = json.loads(tags_val)
|
|
except Exception: tags_list = []
|
|
elif isinstance(tags_val, list):
|
|
tags_list = tags_val
|
|
|
|
# Parse objects
|
|
obj_val = row[25]
|
|
obj_list = []
|
|
if obj_val:
|
|
if isinstance(obj_val, str):
|
|
try: obj_list = json.loads(obj_val)
|
|
except Exception: obj_list = []
|
|
elif isinstance(obj_val, list):
|
|
obj_list = obj_val
|
|
|
|
# Parse source_refs
|
|
ref_val = row[13]
|
|
ref_list = []
|
|
if ref_val:
|
|
if isinstance(ref_val, str):
|
|
try: ref_list = json.loads(ref_val)
|
|
except Exception: ref_list = []
|
|
elif isinstance(ref_val, list):
|
|
ref_list = ref_val
|
|
|
|
# Represent embeddings as lists of floats (for json compatibility)
|
|
embedding_list = None
|
|
if row[5] is not None:
|
|
if isinstance(row[5], str):
|
|
embedding_list = [float(x) for x in row[5].strip("[]").split(",") if x.strip()]
|
|
else:
|
|
embedding_list = list(row[5])
|
|
|
|
face_embedding_list = None
|
|
if row[18] is not None:
|
|
if isinstance(row[18], str):
|
|
face_embedding_list = [float(x) for x in row[18].strip("[]").split(",") if x.strip()]
|
|
else:
|
|
face_embedding_list = list(row[18])
|
|
|
|
# Parse pose_skeleton
|
|
pose_skel = None
|
|
if row[21]:
|
|
try:
|
|
pose_skel = json.loads(row[21])
|
|
except Exception:
|
|
pose_skel = row[21]
|
|
|
|
data = {
|
|
"filename": row[0],
|
|
"filepath": row[1],
|
|
"name": row[2],
|
|
"group_id": row[3],
|
|
"tags": tags_list,
|
|
"embedding": embedding_list,
|
|
"clip_description": row[6],
|
|
"prompt": row[7],
|
|
"pose": row[8],
|
|
"sort_order": row[9],
|
|
"group_name": row[10],
|
|
"hidden": bool(row[11]) if row[11] is not None else False,
|
|
"has_background": bool(row[12]) if row[12] is not None else True,
|
|
"source_refs": ref_list,
|
|
"has_clothing": row[14],
|
|
"content_type": row[15] or "image",
|
|
"faceswap_source_video": row[16],
|
|
"archived": bool(row[17]) if row[17] is not None else False,
|
|
"face_embedding": face_embedding_list,
|
|
"is_source": bool(row[19]) if row[19] is not None else False,
|
|
"pose_description": row[20],
|
|
"pose_skeleton": pose_skel,
|
|
"people_count": row[22],
|
|
"anatomical_completeness": row[23],
|
|
"facial_direction": row[24],
|
|
"objects": obj_list,
|
|
"description": row[26]
|
|
}
|
|
|
|
# Determine json path
|
|
out_dir = _get_output_dir()
|
|
if row[1] and os.path.isabs(row[1]):
|
|
# Use same directory as absolute filepath
|
|
base, _ = os.path.splitext(row[1])
|
|
json_path = base + ".json"
|
|
else:
|
|
# Fallback to output_dir
|
|
base, _ = os.path.splitext(row[0])
|
|
json_path = os.path.join(out_dir, base + ".json")
|
|
|
|
os.makedirs(os.path.dirname(json_path), exist_ok=True)
|
|
tmp_json_path = json_path + ".tmp"
|
|
with open(tmp_json_path, "w") as f:
|
|
json.dump(data, f, indent=2)
|
|
os.replace(tmp_json_path, json_path)
|
|
|
|
except Exception as e:
|
|
print(f"[db] Error syncing person {filename} to JSON: {e}")
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def sync_group_to_json(group_id):
|
|
"""Sync all members of a group, calculate group rating and similarity variation,
|
|
and save to _data/group_<group_id>.json.
|
|
"""
|
|
if not group_id:
|
|
return
|
|
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
# Fetch all members
|
|
cur.execute("""
|
|
SELECT filename, filepath, name, tags, embedding,
|
|
clip_description, prompt, pose, sort_order, group_name, hidden,
|
|
has_background, source_refs, content_type, faceswap_source_video,
|
|
archived, face_embedding, is_source, pose_description, pose_skeleton
|
|
FROM person
|
|
WHERE group_id = %s
|
|
ORDER BY sort_order ASC, filename ASC
|
|
""", (group_id,))
|
|
rows = cur.fetchall()
|
|
|
|
members = []
|
|
likes = 0
|
|
dislikes = 0
|
|
clip_embs = []
|
|
face_embs = []
|
|
|
|
for r in rows:
|
|
tags_val = r[3]
|
|
tags_list = []
|
|
if tags_val:
|
|
if isinstance(tags_val, str):
|
|
try: tags_list = json.loads(tags_val)
|
|
except Exception: tags_list = []
|
|
elif isinstance(tags_val, list):
|
|
tags_list = tags_val
|
|
|
|
if "LIKE" in tags_list:
|
|
likes += 1
|
|
elif "DISLIKE" in tags_list:
|
|
dislikes += 1
|
|
|
|
# Parse embeddings for similarity
|
|
if r[4] is not None:
|
|
if isinstance(r[4], str):
|
|
clip_list = [float(x) for x in r[4].strip("[]").split(",") if x.strip()]
|
|
else:
|
|
clip_list = list(r[4])
|
|
if clip_list:
|
|
clip_embs.append(clip_list)
|
|
|
|
if r[16] is not None:
|
|
if isinstance(r[16], str):
|
|
face_list = [float(x) for x in r[16].strip("[]").split(",") if x.strip()]
|
|
else:
|
|
face_list = list(r[16])
|
|
if face_list:
|
|
face_embs.append(face_list)
|
|
|
|
members.append({
|
|
"filename": r[0],
|
|
"filepath": r[1],
|
|
"name": r[2],
|
|
"tags": tags_list,
|
|
"clip_description": r[5],
|
|
"prompt": r[6],
|
|
"pose": r[7],
|
|
"sort_order": r[8],
|
|
"hidden": bool(r[10]),
|
|
"archived": bool(r[15]),
|
|
"is_source": bool(r[17])
|
|
})
|
|
|
|
def cosine_similarity(v1, v2):
|
|
if not v1 or not v2: return 0.0
|
|
dot = sum(a * b for a, b in zip(v1, v2))
|
|
norm1 = math.sqrt(sum(a * a for a in v1))
|
|
norm2 = math.sqrt(sum(b * b for b in v2))
|
|
if norm1 == 0 or norm2 == 0: return 0.0
|
|
return dot / (norm1 * norm2)
|
|
|
|
def calc_stats(embs):
|
|
if len(embs) < 2:
|
|
return {"average": 1.0, "min": 1.0, "max": 1.0, "variation": 0.0, "pairs_count": 0}
|
|
similarities = []
|
|
for i in range(len(embs)):
|
|
for j in range(i + 1, len(embs)):
|
|
similarities.append(cosine_similarity(embs[i], embs[j]))
|
|
if not similarities:
|
|
return {"average": 1.0, "min": 1.0, "max": 1.0, "variation": 0.0, "pairs_count": 0}
|
|
avg_sim = sum(similarities) / len(similarities)
|
|
min_sim = min(similarities)
|
|
max_sim = max(similarities)
|
|
variance = sum((s - avg_sim) ** 2 for s in similarities) / len(similarities)
|
|
return {
|
|
"average": float(avg_sim),
|
|
"min": float(min_sim),
|
|
"max": float(max_sim),
|
|
"variation": float(math.sqrt(variance)),
|
|
"pairs_count": len(similarities)
|
|
}
|
|
|
|
group_name = rows[0][9] if rows else group_id
|
|
|
|
group_data = {
|
|
"group_id": group_id,
|
|
"group_name": group_name,
|
|
"rating": {
|
|
"likes": likes,
|
|
"dislikes": dislikes,
|
|
"score": likes - dislikes,
|
|
"normalized_score": (likes - dislikes) / (likes + dislikes + 2.0) if (likes + dislikes) > 0 else 0.0,
|
|
"ratio": likes / (likes + dislikes) if (likes + dislikes) > 0 else 0.0
|
|
},
|
|
"clip_similarity_stats": calc_stats(clip_embs),
|
|
"face_similarity_stats": calc_stats(face_embs),
|
|
"members": members
|
|
}
|
|
|
|
out_dir = _get_output_dir()
|
|
data_dir = os.path.join(out_dir, "_data")
|
|
os.makedirs(data_dir, exist_ok=True)
|
|
|
|
json_path = os.path.join(data_dir, f"group_{group_id}.json")
|
|
tmp_json_path = json_path + ".tmp"
|
|
with open(tmp_json_path, "w") as f:
|
|
json.dump(group_data, f, indent=2)
|
|
os.replace(tmp_json_path, json_path)
|
|
|
|
except Exception as e:
|
|
print(f"[db] Error syncing group {group_id} to JSON: {e}")
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def sync_by_filename(filename):
|
|
sync_person_to_json(filename)
|
|
person = get_person(filename)
|
|
if person and person[1]:
|
|
sync_group_to_json(person[1])
|
|
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
_sync_executor = ThreadPoolExecutor(max_workers=2, thread_name_prefix="db_sync_worker")
|
|
|
|
def sync_by_filename_async(filename):
|
|
_sync_executor.submit(sync_by_filename, filename)
|
|
|
|
def sync_group_to_json_async(group_id):
|
|
_sync_executor.submit(sync_group_to_json, group_id)
|
|
|
|
def sync_person_to_json_async(filename):
|
|
_sync_executor.submit(sync_person_to_json, filename)
|
|
|
|
|
|
def save_db_prompt(prompt_type: str, prompt_text: str, metadata: dict = None):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("""
|
|
INSERT INTO prompt (type, prompt_text, metadata, created_at)
|
|
VALUES (%s, %s, %s, NOW())
|
|
ON CONFLICT (type, prompt_text) DO UPDATE
|
|
SET metadata = EXCLUDED.metadata,
|
|
created_at = NOW()
|
|
""", (prompt_type, prompt_text, json.dumps(metadata) if metadata is not None else None))
|
|
conn.commit()
|
|
except Exception as e:
|
|
print(f"[db] error saving prompt {prompt_type}: {e}")
|
|
try:
|
|
conn.rollback()
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def list_db_prompts(prompt_type: str = None, limit: int = 100):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
if prompt_type:
|
|
cur.execute("""
|
|
SELECT id, type, prompt_text, metadata, created_at
|
|
FROM prompt
|
|
WHERE type = %s
|
|
ORDER BY created_at DESC
|
|
LIMIT %s
|
|
""", (prompt_type, limit))
|
|
else:
|
|
cur.execute("""
|
|
SELECT id, type, prompt_text, metadata, created_at
|
|
FROM prompt
|
|
ORDER BY created_at DESC
|
|
LIMIT %s
|
|
""", (limit,))
|
|
|
|
rows = cur.fetchall()
|
|
result = []
|
|
for r in rows:
|
|
meta = r[3]
|
|
if isinstance(meta, str):
|
|
try:
|
|
meta = json.loads(meta)
|
|
except Exception:
|
|
meta = {}
|
|
elif meta is None:
|
|
meta = {}
|
|
result.append({
|
|
"id": r[0],
|
|
"type": r[1],
|
|
"prompt_text": r[2],
|
|
"metadata": meta,
|
|
"created_at": r[4].isoformat() if r[4] else None
|
|
})
|
|
return result
|
|
except Exception as e:
|
|
print(f"[db] error listing prompts: {e}")
|
|
return []
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
|
|
def save_turntable(group_id, state):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("""
|
|
INSERT INTO turntable (group_id, preferred_filename, source_image, seed, n_views, steps, video_path, completed, started_at, completed_at, angles, views)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
ON CONFLICT (group_id) DO UPDATE SET
|
|
preferred_filename = EXCLUDED.preferred_filename,
|
|
source_image = EXCLUDED.source_image,
|
|
seed = EXCLUDED.seed,
|
|
n_views = EXCLUDED.n_views,
|
|
steps = EXCLUDED.steps,
|
|
video_path = EXCLUDED.video_path,
|
|
completed = EXCLUDED.completed,
|
|
started_at = EXCLUDED.started_at,
|
|
completed_at = EXCLUDED.completed_at,
|
|
angles = EXCLUDED.angles,
|
|
views = EXCLUDED.views
|
|
""", (
|
|
group_id,
|
|
state.get("preferred_filename"),
|
|
state.get("source_image"),
|
|
state.get("seed"),
|
|
state.get("n_views"),
|
|
state.get("steps"),
|
|
state.get("video_path"),
|
|
state.get("completed", False),
|
|
state.get("started_at"),
|
|
state.get("completed_at"),
|
|
json.dumps(state.get("angles", [])),
|
|
json.dumps(state.get("views", {}))
|
|
))
|
|
conn.commit()
|
|
except Exception as e:
|
|
print(f"[db] error saving turntable for {group_id}: {e}")
|
|
try:
|
|
conn.rollback()
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
|
|
def load_turntable(group_id):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("""
|
|
SELECT group_id, preferred_filename, source_image, seed, n_views, steps, video_path, completed, started_at, completed_at, angles, views
|
|
FROM turntable
|
|
WHERE group_id = %s
|
|
""", (group_id,))
|
|
row = cur.fetchone()
|
|
if not row:
|
|
return None
|
|
|
|
angles_val = []
|
|
if row[10]:
|
|
try:
|
|
angles_val = json.loads(row[10])
|
|
except Exception:
|
|
angles_val = []
|
|
|
|
views_val = {}
|
|
if row[11]:
|
|
try:
|
|
views_val = json.loads(row[11])
|
|
except Exception:
|
|
views_val = {}
|
|
|
|
return {
|
|
"group_id": row[0],
|
|
"preferred_filename": row[1],
|
|
"source_image": row[2],
|
|
"seed": row[3],
|
|
"n_views": row[4],
|
|
"steps": row[5],
|
|
"video_path": row[6],
|
|
"completed": bool(row[7]) if row[7] is not None else False,
|
|
"started_at": row[8],
|
|
"completed_at": row[9],
|
|
"angles": angles_val,
|
|
"views": views_val
|
|
}
|
|
except Exception as e:
|
|
print(f"[db] error loading turntable for {group_id}: {e}")
|
|
return None
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
|
|
def delete_turntable(group_id):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("DELETE FROM turntable WHERE group_id = %s", (group_id,))
|
|
conn.commit()
|
|
except Exception as e:
|
|
print(f"[db] error deleting turntable for {group_id}: {e}")
|
|
try:
|
|
conn.rollback()
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
|
|
def migrate_turntable_states_from_disk_to_db():
|
|
try:
|
|
out_dir = _get_output_dir()
|
|
td = os.path.join(out_dir, "_turntable")
|
|
if not os.path.isdir(td):
|
|
return
|
|
gids = [d for d in os.listdir(td) if os.path.isfile(os.path.join(td, d, "state.json"))]
|
|
for gid in gids:
|
|
p = os.path.join(td, gid, "state.json")
|
|
try:
|
|
with open(p, "r") as f:
|
|
state = json.load(f)
|
|
if state:
|
|
save_turntable(gid, state)
|
|
except Exception as e:
|
|
print(f"[db] Error migrating turntable state {gid} from disk: {e}")
|
|
except Exception as e:
|
|
print(f"[db] Error in migrate_turntable_states_from_disk_to_db: {e}")
|