updates UI
This commit is contained in:
@@ -101,6 +101,17 @@ def migrate_schema():
|
||||
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)
|
||||
)
|
||||
""")
|
||||
for sql in [
|
||||
"ALTER TABLE person ADD COLUMN IF NOT EXISTS prompt TEXT",
|
||||
"ALTER TABLE person ADD COLUMN IF NOT EXISTS pose TEXT",
|
||||
@@ -262,6 +273,11 @@ def upsert_person(filename, filepath=None, name=None, group_id=None, tags=None,
|
||||
pose_description, pose_skeleton, people_count, anatomical_completeness, facial_direction,
|
||||
json.dumps(objects) if objects else None))
|
||||
conn.commit()
|
||||
# Sync after commit
|
||||
try:
|
||||
sync_by_filename_async(filename)
|
||||
except Exception:
|
||||
pass
|
||||
finally:
|
||||
cur.close()
|
||||
_put_db_connection(conn)
|
||||
@@ -272,6 +288,10 @@ def set_archived(filename, archived: bool):
|
||||
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)
|
||||
@@ -286,6 +306,11 @@ def set_filenames_archived(filenames, archived: bool):
|
||||
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()
|
||||
@@ -297,6 +322,10 @@ def set_hidden(filename, hidden: bool):
|
||||
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)
|
||||
@@ -307,6 +336,10 @@ def set_person_tags(filename, tags):
|
||||
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)
|
||||
@@ -363,21 +396,56 @@ def search_similar(embedding, limit=10):
|
||||
_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)
|
||||
@@ -430,6 +498,13 @@ def set_group_order(group_id, ordered_filenames):
|
||||
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."""
|
||||
@@ -454,6 +529,13 @@ def set_group_name(group_id, name):
|
||||
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)
|
||||
@@ -560,3 +642,367 @@ def invalidate_all_metadata():
|
||||
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
|
||||
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
|
||||
}
|
||||
|
||||
# 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)
|
||||
|
||||
Reference in New Issue
Block a user