acestep/resume_utils.py

"""
resume_utils.py
================
Save & Resume utility for long-running ACE-Step generations on HuggingFace
ZeroGPU Spaces (120-second hard GPU timeout per call).

DESIGN NOTE (important):
This module is built for the LEGITIMATE resume scenario only:
    - Same user, same Space, same account.
    - Today's free ZeroGPU quota runs out before generation finishes.
    - User comes back tomorrow (quota resets) and continues from where
      they left off, using the SAME IP / SAME browser session.
There is no IP-rotation, identity-spoofing, or quota-evasion logic here.
All this code does is persist intermediate tensors to the Space's
persistent disk (/data) under a task_id, so a later call (today or
tomorrow, same account) can pick the work back up instead of starting
from token 0 / step 0 again.

Storage layout (under persistent_storage_path, e.g. /data/resume_state):
    /data/resume_state/<task_id>/lm_tokens.pt        -> generated LM token ids (CPU tensor)
    /data/resume_state/<task_id>/lm_meta.json        -> phase + params + timestamps
    /data/resume_state/<task_id>/dit_latents.pt       -> latents at last checkpointed step
    /data/resume_state/<task_id>/dit_meta.json        -> step index + DiT params + timestamps
    /data/resume_state/<task_id>/audio_codes.json     -> final LM phase-2 output (string codes)

task_id is short, human-typeable (e.g. "AB3K9F"), so a person can write it
down and paste it into a textbox the next day.
"""

import os
import json
import time
import uuid
import string
import random
from typing import Optional, Dict, Any

import torch
from loguru import logger


# ZeroGPU gives ~120s per call. We checkpoint a safety margin before the
# hard kill so the save itself (disk I/O) has time to complete.
DEFAULT_TIMEOUT_SECONDS = 110

# How long an unfinished task is kept on disk before we consider it stale
# and eligible for cleanup. 7 days comfortably covers "finish tomorrow".
MAX_TASK_AGE_SECONDS = 7 * 24 * 3600


def _resume_root(persistent_storage_path: Optional[str]) -> str:
    """Resolve the root directory for resume state.

    Falls back to a local ./resume_state directory if no persistent path
    is configured (e.g. local dev), so the code never crashes — it just
    won't survive a process restart in that case.
    """
    base = persistent_storage_path or "."
    root = os.path.join(base, "resume_state")
    os.makedirs(root, exist_ok=True)
    return root


def generate_task_id() -> str:
    """Generate a short, human-typeable tracking code, e.g. 'AB3K9F'.

    Avoids ambiguous characters (0/O, 1/I/L) so users can read it back
    off a screen and retype it the next day without errors.
    """
    alphabet = "ABCDEFGHJKMNPQRSTUVWXYZ23456789"
    return "".join(random.choice(alphabet) for _ in range(6))


def _task_dir(persistent_storage_path: Optional[str], task_id: str) -> str:
    safe_id = "".join(c for c in task_id.strip().upper() if c.isalnum())
    if not safe_id:
        raise ValueError("Invalid task_id")
    d = os.path.join(_resume_root(persistent_storage_path), safe_id)
    os.makedirs(d, exist_ok=True)
    return d


def task_exists(persistent_storage_path: Optional[str], task_id: str) -> bool:
    try:
        d = os.path.join(_resume_root(persistent_storage_path), task_id.strip().upper())
        return os.path.isdir(d)
    except Exception:
        return False


def get_task_status(persistent_storage_path: Optional[str], task_id: str) -> Dict[str, Any]:
    """Return a human-readable status dict for a task_id, or {'found': False}."""
    if not task_exists(persistent_storage_path, task_id):
        return {"found": False}

    d = _task_dir(persistent_storage_path, task_id)
    status = {"found": True, "task_id": task_id.strip().upper(), "phase": "unknown"}

    lm_meta_path = os.path.join(d, "lm_meta.json")
    dit_meta_path = os.path.join(d, "dit_meta.json")
    codes_path = os.path.join(d, "audio_codes.json")
    final_audio_meta_path = os.path.join(d, "final_meta.json")

    if os.path.exists(final_audio_meta_path):
        with open(final_audio_meta_path, "r", encoding="utf-8") as f:
            status.update(json.load(f))
        status["phase"] = "complete"
    elif os.path.exists(dit_meta_path):
        with open(dit_meta_path, "r", encoding="utf-8") as f:
            dit_meta = json.load(f)
        status["phase"] = "dit_in_progress"
        status["dit_step"] = dit_meta.get("current_step_idx")
        status["dit_total_steps"] = dit_meta.get("total_steps")
        status["saved_at"] = dit_meta.get("saved_at")
    elif os.path.exists(codes_path):
        status["phase"] = "lm_complete_dit_pending"
    elif os.path.exists(lm_meta_path):
        with open(lm_meta_path, "r", encoding="utf-8") as f:
            lm_meta = json.load(f)
        status["phase"] = "lm_in_progress"
        status["lm_step"] = lm_meta.get("step")
        status["saved_at"] = lm_meta.get("saved_at")

    return status


def cleanup_task(persistent_storage_path: Optional[str], task_id: str) -> None:
    """Remove all state for a finished/abandoned task."""
    import shutil
    try:
        d = os.path.join(_resume_root(persistent_storage_path), task_id.strip().upper())
        if os.path.isdir(d):
            shutil.rmtree(d)
            logger.info(f"[resume_utils] Cleaned up task {task_id}")
    except Exception as e:
        logger.warning(f"[resume_utils] Failed to clean up task {task_id}: {e}")


def cleanup_stale_tasks(persistent_storage_path: Optional[str], max_age_seconds: int = MAX_TASK_AGE_SECONDS) -> int:
    """Delete tasks older than max_age_seconds. Call this occasionally (e.g. on app startup).

    Returns the number of tasks removed.
    """
    import shutil
    root = _resume_root(persistent_storage_path)
    removed = 0
    now = time.time()
    try:
        for name in os.listdir(root):
            d = os.path.join(root, name)
            if not os.path.isdir(d):
                continue
            try:
                mtime = os.path.getmtime(d)
                if now - mtime > max_age_seconds:
                    shutil.rmtree(d)
                    removed += 1
            except Exception:
                continue
    except FileNotFoundError:
        pass
    if removed:
        logger.info(f"[resume_utils] Cleaned up {removed} stale task(s)")
    return removed


# ---------------------------------------------------------------------------
# LM phase (token-by-token generation) save/load
# ---------------------------------------------------------------------------

def save_lm_checkpoint(
    persistent_storage_path: Optional[str],
    task_id: str,
    generated_ids: torch.Tensor,
    step: int,
    extra: Optional[Dict[str, Any]] = None,
) -> None:
    """Persist LM token generation state to disk."""
    d = _task_dir(persistent_storage_path, task_id)
    torch.save(generated_ids.detach().to("cpu"), os.path.join(d, "lm_tokens.pt"))

    meta = {
        "step": step,
        "saved_at": time.time(),
        "shape": list(generated_ids.shape),
    }
    if extra:
        meta.update(extra)

    with open(os.path.join(d, "lm_meta.json"), "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False, indent=2)

    logger.info(f"[resume_utils] LM checkpoint saved for task {task_id} at step {step}")


def load_lm_checkpoint(
    persistent_storage_path: Optional[str],
    task_id: str,
    device: str = "cuda",
) -> Optional[Dict[str, Any]]:
    """Load previously saved LM token state, or None if not found."""
    d = os.path.join(_resume_root(persistent_storage_path), task_id.strip().upper())
    tokens_path = os.path.join(d, "lm_tokens.pt")
    meta_path = os.path.join(d, "lm_meta.json")

    if not (os.path.exists(tokens_path) and os.path.exists(meta_path)):
        return None

    generated_ids = torch.load(tokens_path, map_location=device)
    with open(meta_path, "r", encoding="utf-8") as f:
        meta = json.load(f)

    logger.info(f"[resume_utils] LM checkpoint restored for task {task_id} from step {meta.get('step')}")
    return {"generated_ids": generated_ids, "meta": meta}


def save_lm_audio_codes(
    persistent_storage_path: Optional[str],
    task_id: str,
    metadata: Dict[str, Any],
    audio_codes: str,
) -> None:
    """Persist the final (completed) LM phase output: metadata + audio codes string.

    Once this is saved, the LM phase is fully done — the DiT phase can be
    resumed/started independently using this file, even in a brand-new
    call (today or tomorrow).
    """
    d = _task_dir(persistent_storage_path, task_id)
    with open(os.path.join(d, "audio_codes.json"), "w", encoding="utf-8") as f:
        json.dump({"metadata": metadata, "audio_codes": audio_codes, "saved_at": time.time()}, f, ensure_ascii=False, indent=2)
    logger.info(f"[resume_utils] LM phase output (codes) saved for task {task_id}")


def load_lm_audio_codes(persistent_storage_path: Optional[str], task_id: str) -> Optional[Dict[str, Any]]:
    d = os.path.join(_resume_root(persistent_storage_path), task_id.strip().upper())
    path = os.path.join(d, "audio_codes.json")
    if not os.path.exists(path):
        return None
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


# ---------------------------------------------------------------------------
# DiT phase (diffusion denoise loop) save/load
# ---------------------------------------------------------------------------

def save_dit_checkpoint(
    persistent_storage_path: Optional[str],
    task_id: str,
    latents: torch.Tensor,
    current_step_idx: int,
    total_steps: int,
    extra: Optional[Dict[str, Any]] = None,
) -> None:
    """Persist DiT denoising state (latents + step index) to disk."""
    d = _task_dir(persistent_storage_path, task_id)
    torch.save(latents.detach().to("cpu"), os.path.join(d, "dit_latents.pt"))

    meta = {
        "current_step_idx": current_step_idx,
        "total_steps": total_steps,
        "saved_at": time.time(),
        "shape": list(latents.shape),
    }
    if extra:
        meta.update(extra)

    with open(os.path.join(d, "dit_meta.json"), "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False, indent=2)

    logger.info(f"[resume_utils] DiT checkpoint saved for task {task_id} at step {current_step_idx}/{total_steps}")


def load_dit_checkpoint(
    persistent_storage_path: Optional[str],
    task_id: str,
    device: str = "cuda",
) -> Optional[Dict[str, Any]]:
    """Load previously saved DiT latent state, or None if not found."""
    d = os.path.join(_resume_root(persistent_storage_path), task_id.strip().upper())
    latents_path = os.path.join(d, "dit_latents.pt")
    meta_path = os.path.join(d, "dit_meta.json")

    if not (os.path.exists(latents_path) and os.path.exists(meta_path)):
        return None

    latents = torch.load(latents_path, map_location=device)
    with open(meta_path, "r", encoding="utf-8") as f:
        meta = json.load(f)

    logger.info(f"[resume_utils] DiT checkpoint restored for task {task_id} from step {meta.get('current_step_idx')}")
    return {"latents": latents, "meta": meta}


def save_final_marker(persistent_storage_path: Optional[str], task_id: str, audio_path: Optional[str] = None) -> None:
    """Mark a task as fully complete (audio rendered). Useful for status display."""
    d = _task_dir(persistent_storage_path, task_id)
    with open(os.path.join(d, "final_meta.json"), "w", encoding="utf-8") as f:
        json.dump({"completed_at": time.time(), "audio_path": audio_path}, f, ensure_ascii=False, indent=2)


class GenerationTimer:
    """Small helper to check elapsed time against the ZeroGPU budget.

    Usage:
        timer = GenerationTimer(timeout=110)
        for step in ...:
            if timer.expired():
                # save checkpoint and return early
                break
    """

    def __init__(self, timeout: float = DEFAULT_TIMEOUT_SECONDS):
        self.start_time = time.time()
        self.timeout = timeout

    def elapsed(self) -> float:
        return time.time() - self.start_time

    def expired(self) -> bool:
        return self.elapsed() > self.timeout

    def remaining(self) -> float:
        return max(0.0, self.timeout - self.elapsed())