"""Local persistence for generation runs under /data. Each run stores the uploaded input image, one output video per provider, and its own run.json metadata file inside its run directory (rather than an entry appended to one shared index file), so concurrent runs never contend for the same file on disk. On the actual Hugging Face Space, /data is the mounted persistent storage directory. Locally (outside the Space container) /data usually doesn't exist or isn't writable, so ensure_data_dirs() falls back to a ./data directory next to this file for development. """ from __future__ import annotations import json import logging import os import shutil import time import uuid from dataclasses import asdict, dataclass, field logger = logging.getLogger(__name__) DATA_DIR = "/data" _LOCAL_FALLBACK_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "data") RUNS_DIR = os.path.join(DATA_DIR, "runs") @dataclass class RunResult: provider_name: str output_path: str | None status: str # "ok" | "error" error: str | None duration_seconds: float params_used: dict @dataclass class Run: run_id: str timestamp: float input_path: str prompt: str | None = None results: list[RunResult] = field(default_factory=list) def ensure_data_dirs() -> None: """Creates the runs directory, falling back to a local ./data dir if /data isn't writable (i.e. not running inside the Space container).""" global DATA_DIR, RUNS_DIR try: os.makedirs(RUNS_DIR, exist_ok=True) except OSError: logger.warning("%s is not writable; falling back to local dir %s", DATA_DIR, _LOCAL_FALLBACK_DIR) DATA_DIR = _LOCAL_FALLBACK_DIR RUNS_DIR = os.path.join(DATA_DIR, "runs") os.makedirs(RUNS_DIR, exist_ok=True) def get_data_dir() -> str: """Returns the data directory actually in use (call after ensure_data_dirs(), since it may fall back to a local directory).""" return DATA_DIR def create_run(image_path: str) -> tuple[str, str, str]: """Creates a new run directory and copies the uploaded image into it as the input. Returns (run_id, run_dir, input_path). """ run_id = f"{int(time.time())}-{uuid.uuid4().hex[:8]}" run_dir = os.path.join(RUNS_DIR, run_id) os.makedirs(run_dir, exist_ok=True) _, ext = os.path.splitext(image_path) input_path = os.path.join(run_dir, f"input{ext or '.png'}") shutil.copyfile(image_path, input_path) return run_id, run_dir, input_path def save_output_video(run_dir: str, provider_name: str, video_bytes: bytes) -> str: output_path = os.path.join(run_dir, f"{provider_name}.mp4") with open(output_path, "wb") as video_file: video_file.write(video_bytes) return output_path def save_run(run: Run) -> None: """Writes a run's metadata to its own run.json inside its run directory (data/runs/{run_id}/run.json), instead of a shared index file, so concurrent runs never race on read-modify-write of the same file.""" run_path = os.path.join(RUNS_DIR, run.run_id, "run.json") try: with open(run_path, "w") as run_file: json.dump(asdict(run), run_file, indent=2) except OSError: logger.exception("Failed to save run %s to %s", run.run_id, run_path) raise def load_runs() -> list[Run]: """Loads every run by scanning each run directory for its run.json.""" runs = [] if not os.path.isdir(RUNS_DIR): return runs for run_id in sorted(os.listdir(RUNS_DIR)): run_path = os.path.join(RUNS_DIR, run_id, "run.json") try: with open(run_path, "r") as run_file: entry = json.load(run_file) except FileNotFoundError: # Legacy/partial run dirs can exist without metadata; skip quietly. continue except (NotADirectoryError, json.JSONDecodeError, OSError) as exc: logger.warning("Skipping run %s: could not read %s (%s)", run_id, run_path, exc) continue results = [RunResult(**result) for result in entry.get("results", [])] runs.append( Run( run_id=entry["run_id"], timestamp=entry["timestamp"], input_path=entry["input_path"], prompt=entry.get("prompt"), results=results, ) ) runs.sort(key=lambda run: run.timestamp) return runs