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_.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}")