deepafx_st_runtime.py

from __future__ import annotations

import hashlib
import os
import shutil
import sys
import tarfile
import tempfile
import threading
import urllib.request
from pathlib import Path

import numpy as np
import pyloudnorm as pyln
import torch
import torchaudio

CACHE_ROOT = Path(
    os.environ.get("DEEPAFX_CACHE_DIR", Path.home() / ".cache" / "deepafx_st")
)

SOURCE_TAG = "0.1.0"
SOURCE_DIR_NAME = f"DeepAFx-ST-{SOURCE_TAG}"
SOURCE_DIR = CACHE_ROOT / SOURCE_DIR_NAME
SOURCE_ARCHIVE = CACHE_ROOT / f"{SOURCE_DIR_NAME}.tar.gz"
SOURCE_URL = (
    f"https://github.com/adobe-research/DeepAFx-ST/"
    f"archive/refs/tags/v{SOURCE_TAG}.tar.gz"
)
SOURCE_SHA256 = (
    "404f9b291043ac070b78fb0b6cae6e601d6b49edcd0874dd33fb828f358ce83b"
)

CHECKPOINT_REPO = "nateraw/deepafx-st-libritts-autodiff"
CHECKPOINT_REVISION = "080b56b9f4573c0a8fe38de45f49f90ab74521f6"
CHECKPOINT_FILE = "lit_model.ckpt"

SAMPLE_RATE = 24000
TARGET_LOUDNESS = -24.0


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


_source_lock = threading.Lock()
_model_lock = threading.Lock()
_system = None


def _download_archive() -> None:
    if SOURCE_ARCHIVE.is_file():
        if _sha256(SOURCE_ARCHIVE) == SOURCE_SHA256:
            return
        SOURCE_ARCHIVE.unlink()

    partial = Path(f"{SOURCE_ARCHIVE}.part")
    partial.unlink(missing_ok=True)
    request = urllib.request.Request(
        SOURCE_URL,
        headers={"User-Agent": "TEAMuP-DeepAFx-ST-Space/1.0"},
    )
    print("Downloading DeepAFx-ST source...", flush=True)
    try:
        with (
            urllib.request.urlopen(request, timeout=120) as response,
            partial.open("wb") as f,
        ):
            shutil.copyfileobj(response, f)

        actual = _sha256(partial)
        if actual != SOURCE_SHA256:
            raise RuntimeError(
                f"SHA-256 mismatch for DeepAFx-ST source archive: "
                f"expected {SOURCE_SHA256}, got {actual}"
            )
        partial.replace(SOURCE_ARCHIVE)
    except Exception:
        partial.unlink(missing_ok=True)
        raise


def _download_source() -> Path:
    if SOURCE_DIR.is_dir() and (SOURCE_DIR / "deepafx_st" / "system.py").is_file():
        return SOURCE_DIR

    with _source_lock:
        if SOURCE_DIR.is_dir() and (SOURCE_DIR / "deepafx_st" / "system.py").is_file():
            return SOURCE_DIR

        CACHE_ROOT.mkdir(parents=True, exist_ok=True)
        _download_archive()

        with tempfile.TemporaryDirectory(
            prefix="deepafx-src-", dir=CACHE_ROOT
        ) as tmp:
            with tarfile.open(SOURCE_ARCHIVE, "r:gz") as tar:
                tar.extractall(tmp, filter="data")
            extracted = Path(tmp) / SOURCE_DIR_NAME
            if not (extracted / "deepafx_st" / "system.py").is_file():
                raise RuntimeError(
                    "DeepAFx-ST source archive is missing expected files."
                )
            if SOURCE_DIR.exists():
                shutil.rmtree(SOURCE_DIR)
            shutil.move(str(extracted), str(SOURCE_DIR))

    return SOURCE_DIR


def _download_checkpoint() -> Path:
    from huggingface_hub import hf_hub_download

    return Path(
        hf_hub_download(
            repo_id=CHECKPOINT_REPO,
            filename=CHECKPOINT_FILE,
            revision=CHECKPOINT_REVISION,
        )
    )


def _load_system():
    global _system
    if _system is not None:
        return _system

    with _model_lock:
        if _system is not None:
            return _system

        source_dir = _download_source()
        source_path = str(source_dir)
        if source_path not in sys.path:
            sys.path.insert(0, source_path)

        os.environ.setdefault("TORCH_FORCE_NO_WEIGHTS_ONLY_LOAD", "1")

        from deepafx_st.system import System

        ckpt_path = _download_checkpoint()
        print(f"Loading DeepAFx-ST from {ckpt_path}...", flush=True)
        _system = System.load_from_checkpoint(
            str(ckpt_path),
            map_location="cpu",
        )
        _system.eval()

    return _system


def _loudness_normalize(
    waveform: torch.Tensor,
    sample_rate: int,
) -> torch.Tensor:
    audio_np = waveform.squeeze().numpy()
    meter = pyln.Meter(sample_rate)
    loudness = meter.integrated_loudness(audio_np)
    if np.isinf(loudness):
        return waveform
    normalized = pyln.normalize.loudness(audio_np, loudness, TARGET_LOUDNESS)
    return torch.from_numpy(normalized).float().unsqueeze(0)


def _prepare_audio(path: str) -> torch.Tensor:
    waveform, sr = torchaudio.load(path)
    waveform = waveform.mean(dim=0, keepdim=True)
    if sr != SAMPLE_RATE:
        waveform = torchaudio.functional.resample(waveform, sr, SAMPLE_RATE)
    waveform = _loudness_normalize(waveform, SAMPLE_RATE)
    return waveform.unsqueeze(0)


def style_transfer(
    input_path: str,
    reference_path: str,
) -> tuple[torch.Tensor, int, dict]:
    system = _load_system()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    system.to(device)

    x = _prepare_audio(input_path).to(device)
    r = _prepare_audio(reference_path).to(device)

    with torch.inference_mode():
        y_hat, p, _ = system(x, r)

    audio_out = y_hat.squeeze().cpu()

    p_list = p.squeeze().cpu().tolist()
    params = {"raw_parameters": [round(v, 4) for v in p_list]}

    return audio_out, SAMPLE_RATE, params