150ef6dab0
• Implemented global offline mode for all HuggingFace-dependent modules to eliminate unauthenticated request warnings and prevent external connection attempts during startup and runtime. Changes • Global Offline Configuration: Added HF_HUB_OFFLINE=1 and TRANSFORMERS_OFFLINE=1 to the environment variables at the top of edit_api.py, embeddings.py, watcher.py, and pose_llm/pose_llm_api.py. • Explicit Local Files Only: Updated all huggingface_hub.hf_hub_download calls in edit_api.py to include local_files_only=True, ensuring they never attempt to reach the Hub if models are already cached. • LLM Service Optimization: Updated AutoTokenizer and AutoModelForCausalLM calls in pose_llm/pose_llm_api.py to use local_files_only=True. • Consistency: Ensured that shared modules like embeddings.py which uses open_clip are also locked to offline mode to prevent background downloads.
542 lines
22 KiB
Python
542 lines
22 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:
|
|
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",
|
|
]:
|
|
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}")
|
|
|
|
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")
|
|
|
|
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):
|
|
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)
|
|
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)
|
|
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);
|
|
""", (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))
|
|
conn.commit()
|
|
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()
|
|
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()
|
|
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()
|
|
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()
|
|
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
|
|
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
|
|
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):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("DELETE FROM person WHERE filename = %s", (filename,))
|
|
conn.commit()
|
|
finally:
|
|
cur.close()
|
|
_put_db_connection(conn)
|
|
|
|
def delete_group(group_id):
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
try:
|
|
cur.execute("DELETE FROM person WHERE group_id = %s", (group_id,))
|
|
conn.commit()
|
|
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)
|
|
|
|
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()
|
|
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)
|