runtime_support.py

from __future__ import annotations

import hashlib
import math
import os
import shutil
import subprocess
import tempfile
import time
import urllib.request
import uuid
import wave
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, Callable, Mapping

ZERO_GPU_SIZE = "large"
ZERO_GPU_MAX_AUDIO_SECONDS = 10 * 60
MIN_ZERO_GPU_DURATION_SECONDS = 60
MAX_ZERO_GPU_DURATION_SECONDS = 600
OUTPUT_MAX_AGE_SECONDS = 12 * 60 * 60

ASSET_SOURCE_REPO = "TheStinger/UVR5_UI"
ASSET_SOURCE_REVISION = "4790d084e368856b420270939498481f844bd59d"
ASSET_MANIFEST: tuple[dict[str, str | int], ...] = (
    {
        "path": "ilariaaisuite.png",
        "sha256": "60831754632678b333f6301cddb6c96234cfec9750424e60dbed657e0541fcfc",
    },
    {
        "path": "assets/favicon.ico",
        "sha256": "b8001bb2affa855ac0374fa738fc0b257053a8180efde4382a0b94dee411d3f7",
    },
    {
        "path": "test.mp3",
        "size": 296685,
    },
)


class SeparationInputError(ValueError):
    """A user-correctable request validation error safe to display in the UI."""


@dataclass(frozen=True)
class RuntimeInfo:
    mode: str
    device: str
    use_autocast: bool
    gpu_name: str | None

    @property
    def is_zerogpu(self) -> bool:
        return self.mode == "zerogpu"

    @property
    def is_assigned_gpu(self) -> bool:
        return self.mode == "assigned_gpu"

    @property
    def is_cpu(self) -> bool:
        return self.mode == "cpu"


def _env_truthy(name: str) -> bool:
    return os.getenv(name, "").strip().lower() in {"1", "true", "yes", "on"}


def detect_runtime(torch_module: Any) -> RuntimeInfo:
    """Keep ZeroGPU, assigned CUDA, and CPU as separate runtime evidence."""
    is_zerogpu = _env_truthy("SPACES_ZERO_GPU") or _env_truthy("ZEROGPU_V2")
    cuda_available = bool(torch_module.cuda.is_available())

    if is_zerogpu:
        return RuntimeInfo(
            mode="zerogpu",
            device="cuda",
            use_autocast=True,
            gpu_name=f"ZeroGPU {ZERO_GPU_SIZE}",
        )

    if cuda_available:
        try:
            gpu_name = str(torch_module.cuda.get_device_name(torch_module.cuda.current_device()))
        except Exception:
            gpu_name = "CUDA device"
        return RuntimeInfo(
            mode="assigned_gpu",
            device="cuda",
            use_autocast=True,
            gpu_name=gpu_name,
        )

    return RuntimeInfo(mode="cpu", device="cpu", use_autocast=False, gpu_name=None)


def runtime_summary(runtime: RuntimeInfo) -> str:
    return f"UVR5 runtime: {asdict(runtime)}"


def backend_capability_summary(torch_module: Any, ort_module: Any | None = None) -> dict[str, Any]:
    """Report discoverable backends without treating availability as execution proof."""
    try:
        torch_cuda_available = bool(torch_module.cuda.is_available())
    except Exception:
        torch_cuda_available = False

    providers: list[str] = []
    provider_error: str | None = None
    if ort_module is not None:
        try:
            providers = [str(provider) for provider in ort_module.get_available_providers()]
        except Exception as exc:
            provider_error = f"{type(exc).__name__}: {exc}"

    return {
        "torch_cuda_available": torch_cuda_available,
        "onnx_available_providers": providers,
        "onnx_provider_query_error": provider_error,
    }


def separator_backend_summary(separator: Any) -> dict[str, Any]:
    """Expose the backend selected by audio-separator for one request."""
    torch_device = getattr(separator, "torch_device", None)
    onnx_provider = getattr(separator, "onnx_execution_provider", None)
    return {
        "torch_device": None if torch_device is None else str(torch_device),
        "onnx_execution_provider": onnx_provider,
        "use_autocast": bool(getattr(separator, "use_autocast", False)),
    }


def runtime_banner_markdown(runtime: RuntimeInfo) -> str:
    if runtime.is_zerogpu:
        return (
            f"**Runtime: ZeroGPU `{ZERO_GPU_SIZE}`** — GPU time is requested only for separation, "
            "using an initial workload-based quota. Uploaded audio is currently limited to 10 minutes."
        )
    if runtime.is_assigned_gpu:
        return (
            f"**Runtime: assigned GPU** — `{runtime.gpu_name or 'CUDA device'}` detected. "
            "This mode does not request ZeroGPU quota."
        )
    return (
        "**Runtime: CPU** — separation remains enabled as a best-effort compatibility path, "
        "but it can be extremely slow and some model backends may not work in CPU Basic."
    )


def _asset_url(relative_path: str) -> str:
    return (
        f"https://huggingface.co/spaces/{ASSET_SOURCE_REPO}/resolve/"
        f"{ASSET_SOURCE_REVISION}/{relative_path}"
    )


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def ensure_runtime_assets(
    root: str | Path,
    *,
    manifest: tuple[Mapping[str, Any], ...] = ASSET_MANIFEST,
    opener: Callable[..., Any] = urllib.request.urlopen,
) -> list[dict[str, str]]:
    """Download missing binary assets atomically; existing files are intentionally skipped."""
    root_path = Path(root).resolve()
    results: list[dict[str, str]] = []

    for item in manifest:
        relative_path = str(item["path"])
        expected_hash = str(item.get("sha256", "")).lower()
        expected_size = int(item["size"]) if "size" in item else None
        destination = root_path / relative_path

        if destination.exists():
            results.append({"path": relative_path, "status": "existing"})
            continue

        destination.parent.mkdir(parents=True, exist_ok=True)
        temporary = destination.with_name(f".{destination.name}.{uuid.uuid4().hex}.part")
        request = urllib.request.Request(
            _asset_url(relative_path),
            headers={"User-Agent": "UVR5-Tri-Runtime-Modernization/0.1"},
        )

        try:
            with opener(request, timeout=60) as response, temporary.open("wb") as target:
                shutil.copyfileobj(response, target)
            if expected_size is not None:
                actual_size = temporary.stat().st_size
                if actual_size != expected_size:
                    raise RuntimeError(
                        f"Size mismatch for {relative_path}: {actual_size} != {expected_size}"
                    )
            if expected_hash:
                actual_hash = _sha256(temporary)
                if actual_hash != expected_hash:
                    raise RuntimeError(
                        f"SHA256 mismatch for {relative_path}: {actual_hash} != {expected_hash}"
                    )
            os.replace(temporary, destination)
            results.append({"path": relative_path, "status": "downloaded"})
        except Exception as exc:
            temporary.unlink(missing_ok=True)
            results.append({"path": relative_path, "status": "failed", "error": str(exc)})
            print(f"Asset bootstrap warning for {relative_path}: {exc}", flush=True)

    return results


def _ffprobe_duration(path: Path) -> float | None:
    try:
        result = subprocess.run(
            [
                "ffprobe",
                "-v",
                "error",
                "-show_entries",
                "format=duration",
                "-of",
                "default=noprint_wrappers=1:nokey=1",
                str(path),
            ],
            check=False,
            capture_output=True,
            text=True,
            timeout=15,
        )
        if result.returncode == 0:
            duration = float(result.stdout.strip())
            if math.isfinite(duration) and duration > 0:
                return duration
    except Exception:
        pass
    return None


def _wave_duration(path: Path) -> float | None:
    try:
        with wave.open(str(path), "rb") as handle:
            rate = handle.getframerate()
            frames = handle.getnframes()
        if rate > 0 and frames > 0:
            return frames / rate
    except Exception:
        pass
    return None


def audio_duration_seconds(audio_path: str | os.PathLike[str] | None) -> float | None:
    if not audio_path:
        return None
    path = Path(audio_path)
    if not path.is_file():
        return None
    return _ffprobe_duration(path) or _wave_duration(path)


def validate_separation_request(
    *,
    audio_path: str | os.PathLike[str] | None,
    model: str | None,
    output_format: str | None,
    runtime: RuntimeInfo,
) -> float | None:
    if not audio_path or not Path(audio_path).is_file():
        raise SeparationInputError("Please upload an audio file.")
    if not model:
        raise SeparationInputError("Please select a model.")
    if not output_format:
        raise SeparationInputError("Please select an output format.")

    duration = audio_duration_seconds(audio_path)
    if runtime.is_zerogpu and duration and duration > ZERO_GPU_MAX_AUDIO_SECONDS:
        raise SeparationInputError(
            f"ZeroGPU currently accepts audio up to 10 minutes; this file is about "
            f"{duration / 60:.1f} minutes."
        )
    return duration


def estimate_zero_gpu_duration(
    audio_path: str | os.PathLike[str] | None,
    *,
    family: str,
    shifts: int | float = 1,
) -> int:
    """Initial conservative quota formula; calibrate with returned ZeroGPU probes."""
    duration = audio_duration_seconds(audio_path)
    if duration is None:
        return 180

    factors = {
        "roformer": 0.72,
        "mdxc": 0.68,
        "mdxnet": 0.46,
        "vrarch": 0.52,
        "demucs": 0.64,
    }
    factor = factors.get(family, 0.65)
    if family == "demucs":
        factor *= max(1.0, min(float(shifts), 10.0) / 2.0)

    seconds = math.ceil(35.0 + duration * factor)
    return max(MIN_ZERO_GPU_DURATION_SECONDS, min(MAX_ZERO_GPU_DURATION_SECONDS, seconds))


def roformer_duration(audio: Any, *_args: Any, **_kwargs: Any) -> int:
    return estimate_zero_gpu_duration(audio, family="roformer")


def mdxc_duration(audio: Any, *_args: Any, **_kwargs: Any) -> int:
    return estimate_zero_gpu_duration(audio, family="mdxc")


def mdxnet_duration(audio: Any, *_args: Any, **_kwargs: Any) -> int:
    return estimate_zero_gpu_duration(audio, family="mdxnet")


def vrarch_duration(audio: Any, *_args: Any, **_kwargs: Any) -> int:
    return estimate_zero_gpu_duration(audio, family="vrarch")


def demucs_duration(
    audio: Any,
    model: Any = None,
    out_format: Any = None,
    shifts: Any = 1,
    *_args: Any,
    **_kwargs: Any,
) -> int:
    del model, out_format
    return estimate_zero_gpu_duration(audio, family="demucs", shifts=shifts or 1)


def cleanup_request_outputs(root: str | Path, *, max_age_seconds: int = OUTPUT_MAX_AGE_SECONDS) -> None:
    root_path = Path(root)
    if not root_path.exists():
        return
    cutoff = time.time() - max_age_seconds
    for child in root_path.iterdir():
        try:
            if child.is_dir() and child.stat().st_mtime < cutoff:
                shutil.rmtree(child, ignore_errors=True)
        except OSError:
            continue


def create_request_output_dir(root: str | Path) -> Path:
    root_path = Path(root).resolve()
    root_path.mkdir(parents=True, exist_ok=True)
    cleanup_request_outputs(root_path)
    return Path(tempfile.mkdtemp(prefix="request-", dir=root_path))


def resolve_output_paths(output_dir: str | Path, names: Any) -> list[str]:
    root = Path(output_dir).resolve()
    resolved: list[str] = []
    for name in list(names or []):
        candidate = Path(str(name))
        resolved.append(str(candidate if candidate.is_absolute() else root / candidate))
    return resolved


def two_stem_result(stems: list[str], single_stem: str | None) -> tuple[str | None, str | None]:
    first = stems[0] if stems else None
    if (single_stem or "").strip():
        return first, None
    second = stems[1] if len(stems) > 1 else None
    return first, second