reorder

tour-comfy/database.py

@@ -119,6 +119,87 @@ def migrate_schema():
     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,
@@ -201,6 +282,16 @@ def set_hidden(filename, hidden: bool):
         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()
@@ -224,7 +315,7 @@ def list_persons(include_archived=False):
         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
+                   has_clothing, content_type, faceswap_source_video, archived, is_source, tags
             FROM person
             {where}
         """)