Spaces:
Running
Running
| """SA3 pre-encoding job runner — Phase 6. | |
| Encodes every audio clip in a Dataset Workbench project into SA3 latents | |
| once, ahead of training, so the training subprocess can skip the SAME | |
| autoencoder pass per step. Mirrors the shape of `_project_annotate_jobs` | |
| in app.py (background thread, per-project state, cooperative cancel). | |
| Latents land in `<project>/.latents/` — a hidden subdirectory inside the | |
| project folder. Disk layout matches SA3's `pre_encode_dataset.py`: | |
| <project>/.latents/ | |
| 000000000000.npy # latent tensor (shape (256, T_lat)) | |
| 000000000000.json # {"prompt": "...", "padding_mask": [...], ...} | |
| 000001000000.npy | |
| 000001000000.json | |
| ... | |
| silence.npy # padding latent (auto-generated) | |
| _meta.json # Fragmenta-specific: AE used, source clip count | |
| SA3's `train_lora.py --encoded_dir <project>/.latents` consumes this layout | |
| directly. `SA3Trainer._stage_dataset` auto-detects the directory and feeds | |
| `--encoded_dir` to the subprocess when latents are present. | |
| Cache invalidation lives in projects.py — any project mutation that could | |
| desync the latents (commit, delete_clip, slice_clip) wipes the directory. | |
| """ | |
| from __future__ import annotations | |
| import json | |
| import os | |
| import re | |
| import subprocess | |
| import sys | |
| import threading | |
| import time | |
| from pathlib import Path | |
| from typing import Any, Dict, Optional | |
| from app.backend.data.projects import project_path | |
| from app.core.config import get_config | |
| from utils.logger import get_logger | |
| from utils.process_control import graceful_stop | |
| logger = get_logger("PreEncoder") | |
| # --- Per-project job registry ---------------------------------------------- | |
| _pre_encode_jobs: Dict[str, Dict[str, Any]] = {} | |
| _pre_encode_jobs_lock = threading.Lock() | |
| _pre_encode_processes: Dict[str, subprocess.Popen] = {} | |
| def get_pre_encode_job(project_name: str) -> Dict[str, Any]: | |
| """Snapshot of the current job state for a project. Always returns a | |
| well-formed dict so the frontend can render against it without guards.""" | |
| with _pre_encode_jobs_lock: | |
| job = _pre_encode_jobs.get(project_name) | |
| if job is None: | |
| return _idle_job() | |
| return dict(job) | |
| def count_running_pre_encode_jobs() -> int: | |
| """How many pre-encode jobs are queued or running across all projects. | |
| Used by the optional global concurrent-job cap in web deployments. | |
| """ | |
| with _pre_encode_jobs_lock: | |
| return sum(1 for j in _pre_encode_jobs.values() | |
| if j.get("state") in ("queued", "running")) | |
| def _idle_job() -> Dict[str, Any]: | |
| return { | |
| "state": "idle", # idle | queued | running | complete | failed | cancelled | |
| "current": 0, # batch index (0-based) | |
| "total": 0, # total batches (derived from clip count) | |
| "current_file": "", | |
| "error": None, | |
| "started_at": None, | |
| "finished_at": None, | |
| "autoencoder": None, | |
| } | |
| # --- Autoencoder selection ------------------------------------------------- | |
| DEFAULT_AUTOENCODER = "same-s" | |
| DEFAULT_SAMPLE_SIZE = 12_582_912 | |
| # --- Job lifecycle --------------------------------------------------------- | |
| def latents_dir(project_name: str) -> Path: | |
| return project_path(project_name) / ".latents" | |
| def latents_count(project_name: str) -> int: | |
| d = latents_dir(project_name) | |
| if not d.exists(): | |
| return 0 | |
| return sum( | |
| 1 for p in d.glob("*.npy") | |
| if p.name != "silence.npy" | |
| ) | |
| def latents_meta(project_name: str) -> Optional[Dict[str, Any]]: | |
| """Read the manifest we drop alongside the .npy files.""" | |
| p = latents_dir(project_name) / "_meta.json" | |
| if not p.exists(): | |
| return None | |
| try: | |
| return json.loads(p.read_text(encoding="utf-8")) | |
| except Exception: | |
| return None | |
| def latents_match_base(project_name: str, base_model: str) -> bool: | |
| """Whether the cached latents are compatible with the chosen base. | |
| same-s ↔ small-music / small-sfx (and their *-base variants). | |
| same-l ↔ medium (and medium-base). | |
| """ | |
| meta = latents_meta(project_name) | |
| if not meta: | |
| return False | |
| ae = meta.get("autoencoder") | |
| if ae == "same-s": | |
| return base_model in ("sa3-small-music", "sa3-small-music-base", | |
| "sa3-small-sfx", "sa3-small-sfx-base") | |
| if ae == "same-l": | |
| return base_model in ("sa3-medium", "sa3-medium-base") | |
| return False | |
| def cancel_pre_encode(project_name: str) -> bool: | |
| """Stop an in-flight job. Returns True if it was actually cancelled. | |
| Order matters: signal the subprocess FIRST, flip the user-visible state | |
| only once that succeeded. The old code marked the job "cancelled" before | |
| signalling, so a failed signal (the exact Windows failure mode — | |
| send_signal(SIGINT) raises ValueError there) reported success while the | |
| encoder kept running. The `cancelled` flag is set up front though: the | |
| worker thread uses it to label the child's exit, and it would race us if | |
| set only after the process dies. | |
| """ | |
| with _pre_encode_jobs_lock: | |
| job = _pre_encode_jobs.get(project_name) | |
| if not job or job.get("state") not in ("queued", "running"): | |
| return False | |
| job["cancelled"] = True | |
| proc = _pre_encode_processes.get(project_name) | |
| if proc is not None and proc.poll() is None: | |
| try: | |
| graceful_stop(proc, wait_timeout=5, kill_timeout=3) | |
| except Exception as exc: | |
| logger.warning("Failed to signal pre-encode subprocess: %s", exc) | |
| with _pre_encode_jobs_lock: | |
| live = _pre_encode_jobs.get(project_name) | |
| if live is not None: | |
| live["cancelled"] = False | |
| return False | |
| with _pre_encode_jobs_lock: | |
| job = _pre_encode_jobs.get(project_name) | |
| if job is not None and job.get("state") in ("queued", "running"): | |
| job["state"] = "cancelled" | |
| job["finished_at"] = time.time() | |
| return True | |
| def start_pre_encode( | |
| project_name: str, | |
| autoencoder: Optional[str] = None, | |
| sample_size: Optional[int] = None, | |
| ) -> Dict[str, Any]: | |
| """Spawn the pre-encode subprocess in a background thread. Returns the | |
| job state — frontend polls /pre-encode/status thereafter. | |
| """ | |
| proj_dir = project_path(project_name) | |
| if not proj_dir.exists(): | |
| raise FileNotFoundError(f"project not found: {project_name}") | |
| ae = autoencoder or DEFAULT_AUTOENCODER | |
| if ae not in ("same-s", "same-l"): | |
| raise ValueError(f"autoencoder must be 'same-s' or 'same-l'; got {ae!r}") | |
| with _pre_encode_jobs_lock: | |
| existing = _pre_encode_jobs.get(project_name) | |
| if existing and existing.get("state") in ("queued", "running"): | |
| return dict(existing) | |
| # Count source clips (sidecars committed) so we know the denominator. | |
| sidecars = list(proj_dir.glob("*.txt")) | |
| clip_count = sum( | |
| 1 for p in sidecars | |
| if p.read_text(encoding="utf-8").strip() | |
| and p.with_suffix(".wav").exists() # cheap & accurate enough | |
| ) | |
| job: Dict[str, Any] = { | |
| "state": "queued", | |
| "current": 0, | |
| "total": clip_count, | |
| "current_file": "", | |
| "error": None, | |
| "started_at": time.time(), | |
| "finished_at": None, | |
| "autoencoder": ae, | |
| "cancelled": False, | |
| } | |
| _pre_encode_jobs[project_name] = job | |
| thread = threading.Thread( | |
| target=_run_pre_encode, | |
| args=(project_name, ae, sample_size or DEFAULT_SAMPLE_SIZE), | |
| daemon=True, | |
| name=f"sa3-pre-encode:{project_name}", | |
| ) | |
| thread.start() | |
| return get_pre_encode_job(project_name) | |
| # --- Worker ---------------------------------------------------------------- | |
| def _update_job(project_name: str, **fields: Any) -> None: | |
| with _pre_encode_jobs_lock: | |
| job = _pre_encode_jobs.get(project_name) | |
| if job is None: | |
| return | |
| job.update(fields) | |
| def _run_pre_encode(project_name: str, ae: str, sample_size: int) -> None: | |
| """Background-thread target. Spawns the SA3 pre_encode_dataset.py script, | |
| streams stdout for progress, writes a _meta.json manifest on success.""" | |
| cfg = get_config() | |
| proj_dir = project_path(project_name) | |
| out_dir = latents_dir(project_name) | |
| out_dir.mkdir(parents=True, exist_ok=True) | |
| sa3_vendor = cfg.get_path("stable_audio_3") | |
| venv_python = sys.executable | |
| cmd = [ | |
| venv_python, | |
| str(sa3_vendor / "scripts" / "pre_encode_dataset.py"), | |
| "--model", ae, | |
| "--data_dir", str(proj_dir), | |
| "--output_path", str(out_dir), | |
| "--batch_size", "1", | |
| "--sample_size", str(int(sample_size)), | |
| ] | |
| env = os.environ.copy() | |
| pp = env.get("PYTHONPATH", "") | |
| env["PYTHONPATH"] = ( | |
| f"{sa3_vendor}{os.pathsep}{pp}" if pp else str(sa3_vendor) | |
| ) | |
| hub_dir = cfg.get_path("models_pretrained") / "sa3" / "hub" | |
| env["HF_HUB_CACHE"] = str(hub_dir) | |
| env["HUGGINGFACE_HUB_CACHE"] = str(hub_dir) | |
| env["TRANSFORMERS_CACHE"] = str(hub_dir) | |
| env["HF_HUB_DISABLE_PROGRESS_BARS"] = "1" | |
| env["HF_HUB_OFFLINE"] = "1" | |
| env["TRANSFORMERS_OFFLINE"] = "1" | |
| env["PYTHONIOENCODING"] = "utf-8" | |
| env["PYTHONUTF8"] = "1" | |
| env["PYTHONUNBUFFERED"] = "1" | |
| snapshot = get_pre_encode_job(project_name) | |
| if snapshot and snapshot.get("cancelled"): | |
| _update_job(project_name, state="cancelled", finished_at=time.time()) | |
| return | |
| _update_job(project_name, state="running") | |
| logger.info( | |
| "Pre-encoding started · project=%s · autoencoder=%s · clips=%d · sample_size=%d", | |
| project_name, ae, get_pre_encode_job(project_name)["total"], sample_size, | |
| ) | |
| batch_pat = re.compile(r"Processing batch (\d+)") | |
| process: Optional[subprocess.Popen] = None | |
| try: | |
| process = subprocess.Popen( | |
| cmd, | |
| cwd=str(cfg.project_root), | |
| env=env, | |
| stdout=subprocess.PIPE, | |
| stderr=subprocess.STDOUT, | |
| text=True, | |
| encoding="utf-8", | |
| errors="replace", | |
| bufsize=1, | |
| ) | |
| _pre_encode_processes[project_name] = process | |
| if process.stdout is not None: | |
| for line in process.stdout: | |
| line = line.rstrip() | |
| m = batch_pat.search(line) | |
| if m: | |
| _update_job(project_name, current=int(m.group(1)) + 1) | |
| rc = process.wait() if process else 1 | |
| snapshot = get_pre_encode_job(project_name) | |
| if snapshot.get("cancelled"): | |
| _update_job( | |
| project_name, | |
| state="cancelled", | |
| finished_at=time.time(), | |
| ) | |
| logger.info("Pre-encoding cancelled · project=%s", project_name) | |
| return | |
| if rc != 0: | |
| _update_job( | |
| project_name, | |
| state="failed", | |
| error=f"pre_encode_dataset.py exited with code {rc}", | |
| finished_at=time.time(), | |
| ) | |
| logger.error( | |
| "Pre-encoding failed (exit %s) · project=%s", | |
| rc, project_name, | |
| ) | |
| return | |
| manifest = { | |
| "autoencoder": ae, | |
| "sample_size": sample_size, | |
| "created_at": time.time(), | |
| "source_clip_count": snapshot.get("total", 0), | |
| "encoded_count": latents_count(project_name), | |
| } | |
| try: | |
| (out_dir / "_meta.json").write_text( | |
| json.dumps(manifest, indent=2), encoding="utf-8", | |
| ) | |
| except Exception as exc: | |
| logger.warning("Failed to write latents manifest: %s", exc) | |
| _update_job( | |
| project_name, | |
| state="complete", | |
| current=manifest["encoded_count"], | |
| total=manifest["encoded_count"] or snapshot.get("total", 0), | |
| finished_at=time.time(), | |
| ) | |
| logger.info( | |
| "Pre-encoding complete · project=%s · %d latent(s) · ae=%s", | |
| project_name, manifest["encoded_count"], ae, | |
| ) | |
| except Exception as exc: | |
| _update_job( | |
| project_name, | |
| state="failed", | |
| error=str(exc), | |
| finished_at=time.time(), | |
| ) | |
| logger.exception("Pre-encoding crashed for project=%s", project_name) | |
| finally: | |
| _pre_encode_processes.pop(project_name, None) | |