reorder

tour-comfy/turntable_cache.py

@@ -0,0 +1,151 @@
+"""
+turntable_cache.py — persistent state for Qwen turntable generation.
+
+State stored as JSON: {output_dir}/_turntable/{group_id}/state.json
+Views stored alongside: {output_dir}/_turntable/{group_id}/views/view_NNN_DDDdeg.png
+Final video: {output_dir}/_turntable/{group_id}/turntable.mp4
+
+One state file per group tracks completed angles so background generation can
+resume after restart without re-rendering anything that's already on disk.
+"""
+
+import os
+import json
+import time
+from typing import Optional
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+
+
+def cache_dir(output_dir: str, group_id: str) -> str:
+    return os.path.join(output_dir, "_turntable", str(group_id))
+
+
+def state_path(output_dir: str, group_id: str) -> str:
+    return os.path.join(cache_dir(output_dir, group_id), "state.json")
+
+
+def load_state(output_dir: str, group_id: str) -> Optional[dict]:
+    p = state_path(output_dir, group_id)
+    if not os.path.exists(p):
+        return None
+    try:
+        with open(p) as f:
+            return json.load(f)
+    except Exception:
+        return None
+
+
+def save_state(output_dir: str, group_id: str, state: dict):
+    os.makedirs(cache_dir(output_dir, group_id), exist_ok=True)
+    p = state_path(output_dir, group_id)
+    tmp = p + ".tmp"
+    with open(tmp, "w") as f:
+        json.dump(state, f, indent=2)
+    os.replace(tmp, p)  # atomic
+
+
+def init_state(
+    output_dir: str,
+    group_id: str,
+    source_image: str,
+    preferred_filename: str,
+    n_views: int = 24,
+    seed: int = 42,
+    steps: int = 8,
+) -> dict:
+    """Create a fresh state dict and save it. Wipes any existing partial state."""
+    import sys
+    if _HERE not in sys.path:
+        sys.path.insert(0, _HERE)
+    from orbit_qwen import _angles_for
+
+    angles = _angles_for("turntable", n_views, 180.0)
+    state = {
+        "group_id": str(group_id),
+        "preferred_filename": preferred_filename,
+        "source_image": source_image,
+        "seed": seed,
+        "n_views": n_views,
+        "steps": steps,
+        "angles": angles,
+        "views": {},        # deg_key (str) -> abs path
+        "video_path": None,
+        "completed": False,
+        "started_at": time.time(),
+        "completed_at": None,
+    }
+    save_state(output_dir, group_id, state)
+    return state
+
+
+def deg_key(deg: float) -> str:
+    return f"{deg:.1f}"
+
+
+def mark_view_done(output_dir: str, group_id: str, state: dict, deg: float, path: str):
+    state["views"][deg_key(deg)] = path
+    save_state(output_dir, group_id, state)
+
+
+def mark_completed(output_dir: str, group_id: str, state: dict, video_path: str):
+    state["completed"] = True
+    state["video_path"] = video_path
+    state["completed_at"] = time.time()
+    save_state(output_dir, group_id, state)
+
+
+def next_missing_angle(state: dict) -> Optional[float]:
+    """Return first angle not yet in state['views'], or None if all done."""
+    done = state.get("views", {})
+    for deg in state.get("angles", []):
+        if deg_key(deg) not in done:
+            return deg
+    return None
+
+
+def list_cached_group_ids(output_dir: str) -> list:
+    td = os.path.join(output_dir, "_turntable")
+    if not os.path.isdir(td):
+        return []
+    return [
+        d for d in os.listdir(td)
+        if os.path.isfile(os.path.join(td, d, "state.json"))
+    ]
+
+
+def get_status_summary(output_dir: str) -> dict:
+    """Return {group_id: status_dict} for all groups that have a state file."""
+    result = {}
+    for gid in list_cached_group_ids(output_dir):
+        st = load_state(output_dir, gid)
+        if st:
+            result[gid] = {
+                "completed": st.get("completed", False),
+                "n_done": len(st.get("views", {})),
+                "n_total": st.get("n_views", 24),
+                "video_path": st.get("video_path"),
+                "preferred_filename": st.get("preferred_filename"),
+                "started_at": st.get("started_at"),
+                "completed_at": st.get("completed_at"),
+            }
+    return result
+
+
+def delete_state(output_dir: str, group_id: str):
+    """Wipe all cached views, state, and video for this group."""
+    import shutil
+    d = cache_dir(output_dir, group_id)
+    if os.path.isdir(d):
+        shutil.rmtree(d)
+
+
+def get_group_video(output_dir: str, group_id: str) -> Optional[str]:
+    """Return the video path if the turntable is complete and the file exists."""
+    st = load_state(output_dir, group_id)
+    if not st or not st.get("completed"):
+        return None
+    vp = st.get("video_path")
+    if vp and os.path.exists(vp):
+        return vp
+    return None