updates UI

tour-comfy/database.py

@@ -112,6 +112,23 @@ def migrate_schema():
                 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",
@@ -142,6 +159,10 @@ def migrate_schema():
         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()
@@ -214,6 +235,14 @@ def initialize_tags_in_db():
                 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:
@@ -1006,3 +1035,134 @@ def list_db_prompts(prompt_type: str = None, limit: int = 100):
     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}")