152 lines
4.5 KiB
Python
152 lines
4.5 KiB
Python
"""
|
|
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
|