app.py

"""PyMSS 三板斧人声分离 + RVC 变声翻唱 —— HuggingFace ZeroGPU / Gradio Space.

流程（与 PyMss Studio 的工作流一致）：
    原始音频
      └─ 步骤1 big_beta6x                       -> other(伴奏) / vocals
           └─ 步骤2 mel_band_karaoke_fusion     -> Instrumental(和声) / Vocals
                └─ 步骤3 dereverb ...           -> reverb(混响) / noreverb
     noreverb --RVC 变声--> 变声后的主人声
     Instrumental --(可选)RVC 变声--> 和声
     输出：other / 变声主人声 / 和声 / reverb / 合并成品
"""

from __future__ import annotations

import logging
import os
import shutil
import sys
import tempfile
import threading
import time
from pathlib import Path

import gradio as gr
import librosa
import numpy as np
import soundfile as sf
import spaces
import torch
import yaml
from huggingface_hub import hf_hub_download, snapshot_download
from huggingface_hub.utils import disable_progress_bars

# --------------------------------------------------------------------------------------
# 基础设置
# --------------------------------------------------------------------------------------
ROOT = Path(__file__).resolve().parent
JOB_TTL_SECONDS = 3600

# infer_rvc_python 是 vendored 在仓库里的（原因见 requirements.txt 顶部）：
# 必须先把它加进 sys.path，pip 才不会去解析它那条死锁的依赖元数据。
VENDOR_DIR = ROOT / "vendor"
if VENDOR_DIR.is_dir():
    sys.path.insert(0, str(VENDOR_DIR))

logging.basicConfig(level=logging.INFO)
logging.getLogger("infer_rvc_python").setLevel(logging.WARNING)
disable_progress_bars()

with open(ROOT / "config.yaml", "r", encoding="utf-8") as fp:
    CFG = yaml.safe_load(fp) or {}

SEP_CFG = CFG.get("separation") or {}
RT_CFG = CFG.get("runtime") or {}
VOICES_CFG = [v for v in (CFG.get("voices") or []) if v.get("repo_id")]

OUT_SR = int(RT_CFG.get("output_sr", 44100))
GPU_DECLARED_SECONDS = int(RT_CFG.get("gpu_declared_seconds", 60))
# ZeroGPU 的 SDK 会把声明值乘上硬件系数再交给调度器（spaces/zero/configs.json 里
# RTX PRO 6000 Blackwell 是 1.5），这个系数同时也是函数真正能跑的时间。
GPU_DURATION_FACTOR = float(RT_CFG.get("gpu_duration_factor", 1.5))
MIX_PEAK = float(RT_CFG.get("mix_peak", 0.97))
PITCH_ALGO = str(RT_CFG.get("pitch_algo", "rmvpe+"))
HF_TOKEN = os.environ.get("HF_TOKEN") or None

# Space 仓库的 Git/LFS 上限约 1 GB，所以模型一律下载到容器的临时盘，不进仓库
WORK_DIR = Path(os.environ.get("WORK_DIR") or RT_CFG.get("work_dir") or "/tmp/rvc_cover")
MODEL_CACHE_DIR = Path(
    os.environ.get("MODEL_CACHE_DIR") or RT_CFG.get("model_cache_dir") or "/tmp/pymss_models"
)
WORK_DIR.mkdir(parents=True, exist_ok=True)
MODEL_CACHE_DIR.mkdir(parents=True, exist_ok=True)

# ZeroGPU：模块作用域下 torch.cuda.is_available() 被 spaces 打了补丁恒为 True，
# 此时 .to("cuda") 只是登记权重，真正计算发生在 @spaces.GPU 内部。
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

GPU_LOCK = threading.Lock()
STARTUP_ERRORS: list[str] = []

# --------------------------------------------------------------------------------------
# 模块级加载模型（ZeroGPU 要求：模型必须在模块加载时就放到 cuda）
# --------------------------------------------------------------------------------------
from pymss import MSSeparator  # noqa: E402

SEPARATOR_KEYS = ("vocal_split", "karaoke_split", "dereverb_split")
SEP_REPO = SEP_CFG.get("repo_id", "baicai1145/pymss")
SEPARATORS: dict[str, MSSeparator] = {}


def _download_separation_models() -> str:
    """从 pymss 官方仓库拉取三板斧模型（公开模型，无需 token），返回快照目录。"""
    patterns: list[str] = []
    for key in SEPARATOR_KEYS:
        section = SEP_CFG.get(key) or {}
        patterns += [section.get("ckpt"), section.get("config")]
    patterns = [p for p in patterns if p]
    print(f"[startup] 正在从 {SEP_REPO} 下载/校验 {len(patterns)} 个模型文件（约 3.3 GB，首次较慢）…")
    return snapshot_download(
        repo_id=SEP_REPO,
        allow_patterns=patterns,
        cache_dir=str(MODEL_CACHE_DIR),
        token=HF_TOKEN,
    )


def _build_separator(section: dict, snapshot_dir: str) -> MSSeparator:
    return MSSeparator(
        model_type=section["type"],
        model_path=str(Path(snapshot_dir) / section["ckpt"]),
        config_path=str(Path(snapshot_dir) / section["config"]),
        device=DEVICE,
        device_ids=[0],
        output_format="wav",
        store_dirs=str(WORK_DIR / "_unused"),
        inference_params={
            "batch_size": section.get("batch_size"),
            "chunk_size": section.get("chunk_size"),
            "overlap_size": section.get("overlap_size"),
            "normalize": False,
        },
    )


_missing_sections = [k for k in SEPARATOR_KEYS if not SEP_CFG.get(k)]
if _missing_sections:
    STARTUP_ERRORS.append("config.yaml 缺少 separation 配置: " + ", ".join(_missing_sections))
else:
    try:
        _snapshot_dir = _download_separation_models()
    except Exception as _exc:  # noqa: BLE001
        STARTUP_ERRORS.append(f"分离模型下载失败: {_exc}")
        print(f"[startup] 分离模型下载失败: {_exc}")
    else:
        for _key in SEPARATOR_KEYS:
            try:
                SEPARATORS[_key] = _build_separator(SEP_CFG[_key], _snapshot_dir)
                print(f"[startup] 已加载分离模型: {_key} -> {SEP_CFG[_key]['ckpt']}")
            except Exception as _exc:  # noqa: BLE001
                STARTUP_ERRORS.append(f"分离模型 {_key} 加载失败: {_exc}")
                print(f"[startup] 分离模型 {_key} 加载失败: {_exc}")

# ---- RVC 音色（每个音色一个 BaseLoader，切换时不互相踢掉已加载的模型）----
CONVERTERS: dict[str, object] = {}
VOICE_ASSETS: dict[str, dict] = {}
VOICE_LABELS: dict[str, str] = {}
APPLIED_PARAMS: dict[str, dict] = {}

# 键名必须和 _apply_voice_conf 的形参名一致；到上游 apply_conf 的参数名
# （index_influence / respiration_median_filtering / ...）由该函数内部映射
DEFAULT_VOICE_PARAMS = {
    "index_rate": 0.66,
    "protect": 0.33,
    "filter_radius": 3,
    "rms_mix_rate": 0.25,
}


def _fetch_voice_assets(voice: dict) -> tuple[str, str]:
    """从私有 HF 仓库拉取 .pth / .index。返回 (model_path, index_path)。"""
    target = MODEL_CACHE_DIR / "rvc" / voice["id"]
    target.mkdir(parents=True, exist_ok=True)
    kwargs = {"token": HF_TOKEN} if HF_TOKEN else {}

    model_path = hf_hub_download(
        repo_id=voice["repo_id"], filename=voice["model_file"], local_dir=str(target), **kwargs
    )
    index_path = ""
    if voice.get("index_file"):
        try:
            index_path = hf_hub_download(
                repo_id=voice["repo_id"], filename=voice["index_file"], local_dir=str(target), **kwargs
            )
        except Exception as exc:  # noqa: BLE001
            print(f"[startup] 音色 {voice['id']} 未取到 index 文件，将不使用 index: {exc}")
    return str(model_path), str(index_path)


def _apply_voice_conf(voice_id: str, pitch: int, index_rate: float, protect: float,
                      filter_radius: int, rms_mix_rate: float) -> None:
    """把参数写进 BaseLoader；参数没变就不重复 apply_conf（避免反复重载模型）。"""
    params = {
        "pitch_algo": PITCH_ALGO,
        "pitch_lvl": int(pitch),
        "index_influence": float(index_rate),
        "respiration_median_filtering": int(filter_radius),
        "envelope_ratio": float(rms_mix_rate),
        "consonant_breath_protection": float(protect),
    }
    if APPLIED_PARAMS.get(voice_id) == params:
        return
    assets = VOICE_ASSETS[voice_id]
    CONVERTERS[voice_id].apply_conf(
        tag=voice_id,
        file_model=assets["model"],
        file_index=assets["index"],
        **params,
    )
    APPLIED_PARAMS[voice_id] = dict(params)


if VOICES_CFG:
    try:
        from infer_rvc_python import BaseLoader

        _rmvpe_path = hf_hub_download(repo_id="r3gm/sonitranslate_voice_models", filename="rmvpe.pt")
        for _voice in VOICES_CFG:
            try:
                _model_path, _index_path = _fetch_voice_assets(_voice)
                _converter = BaseLoader(
                    only_cpu=False,
                    hubert_path="r3gm/hubert_base",
                    rmvpe_path=_rmvpe_path,
                    preload_models=True,
                )
                CONVERTERS[_voice["id"]] = _converter
                VOICE_ASSETS[_voice["id"]] = {"model": _model_path, "index": _index_path}
                VOICE_LABELS[_voice["id"]] = _voice.get("label") or _voice["id"]
                _apply_voice_conf(_voice["id"], pitch=0, **DEFAULT_VOICE_PARAMS)
                print(f"[startup] 已加载 RVC 音色: {_voice['id']} -> {_voice['repo_id']}")
            except Exception as _exc:  # noqa: BLE001
                # 失败时回滚，避免出现「下拉框里有这个音色但告警说它没加载成功」的矛盾状态
                CONVERTERS.pop(_voice["id"], None)
                VOICE_ASSETS.pop(_voice["id"], None)
                VOICE_LABELS.pop(_voice["id"], None)
                APPLIED_PARAMS.pop(_voice["id"], None)
                STARTUP_ERRORS.append(f"RVC 音色 {_voice.get('id')} 加载失败: {_exc}")
                print(f"[startup] RVC 音色 {_voice.get('id')} 加载失败: {_exc}")
    except Exception as _exc:  # noqa: BLE001
        STARTUP_ERRORS.append(f"infer_rvc_python 初始化失败: {_exc}")
        print(f"[startup] infer_rvc_python 初始化失败: {_exc}")
else:
    STARTUP_ERRORS.append("config.yaml 里还没有填写任何 RVC 音色的 repo_id")


# --------------------------------------------------------------------------------------
# 音频工具（内部统一使用 (channels, samples) 的 float32）
# --------------------------------------------------------------------------------------
def _load_mix(path: str, sr: int) -> np.ndarray:
    audio, _ = librosa.load(path, sr=sr, mono=False)
    if audio.ndim == 1:
        audio = np.stack([audio, audio], axis=0)
    elif audio.shape[0] > 2:
        audio = audio[:2]
    return np.ascontiguousarray(audio, dtype=np.float32)


def _fit_length(audio: np.ndarray, length: int) -> np.ndarray:
    if audio.shape[1] >= length:
        return audio[:, :length]
    return np.pad(audio, ((0, 0), (0, length - audio.shape[1])))


def _write_wav(path, audio_ch_first: np.ndarray, sr: int) -> None:
    """落盘。内部一律用 (channels, samples)，而 soundfile 要的是 (frames, channels)。"""
    sf.write(str(path), np.ascontiguousarray(audio_ch_first.T, dtype=np.float32), sr)


def _cleanup_old_jobs() -> None:
    now = time.time()
    try:
        for path in WORK_DIR.glob("job_*"):
            if path.is_dir() and now - path.stat().st_mtime > JOB_TTL_SECONDS:
                shutil.rmtree(path, ignore_errors=True)
    except Exception:  # noqa: BLE001
        pass


def _convert_voice(voice_id: str, mono: np.ndarray, sr: int) -> np.ndarray:
    """RVC 变声：输入单声道 float32，返回等长的单声道 float32（OUT_SR）。"""
    converter = CONVERTERS[voice_id]
    arr, out_sr = converter.generate_from_cache(audio_data=(mono.astype(np.float32), sr), tag=voice_id)
    mono_out = np.asarray(arr, dtype=np.float32)
    if mono_out.ndim > 1:
        mono_out = mono_out.reshape(-1)
    if np.issubdtype(np.asarray(arr).dtype, np.integer):
        mono_out = mono_out / 32768.0
    if int(out_sr) != OUT_SR:
        mono_out = librosa.resample(mono_out, orig_sr=int(out_sr), target_sr=OUT_SR)
    return np.ascontiguousarray(mono_out, dtype=np.float32)


# --------------------------------------------------------------------------------------
# ZeroGPU 时长
# 实测：204s 音频 → 实际 38~40s（约 0.19 s/s），且配额按**实际耗时**扣，
# 但调度门槛按**声明值**预检查 —— 所以声明值固定为一个小常数最省事。
# --------------------------------------------------------------------------------------
def reserved_seconds(declared: int) -> int:
    """声明值换算成调度器实际预占的秒数 —— 就是报错里 "requested" 那个数。"""
    return int(declared * GPU_DURATION_FACTOR)


# --------------------------------------------------------------------------------------
# 主流程
# --------------------------------------------------------------------------------------
@spaces.GPU(duration=GPU_DECLARED_SECONDS)
def generate(audio_path, voice_id, pitch, index_rate, protect, filter_radius,
             rms_mix_rate, remove_harmony, process_harmony):
    missing = [k for k in ("vocal_split", "karaoke_split", "dereverb_split") if k not in SEPARATORS]
    if missing:
        raise gr.Error("分离模型没加载成功，请检查 Space 的 models/ 目录和启动日志：" + "；".join(STARTUP_ERRORS))
    if voice_id not in CONVERTERS:
        raise gr.Error("RVC 音色没加载成功，请检查 config.yaml 的 voices 和 HF_TOKEN：" + "；".join(STARTUP_ERRORS))
    if not audio_path:
        raise gr.Error("请先上传一首歌")

    _cleanup_old_jobs()
    started = time.time()
    with GPU_LOCK:
        job_dir = Path(tempfile.mkdtemp(prefix="job_", dir=str(WORK_DIR)))
        try:
            # 1) 读入原曲
            mix = _load_mix(audio_path, OUT_SR)
            total = mix.shape[1]

            # 2) 步骤1：原曲 -> other / vocals
            step1 = SEPARATORS["vocal_split"].separate(
                mix, pbar=False, stems=["vocals", "other"]
            )
            other = np.ascontiguousarray(step1["other"].T, dtype=np.float32)      # 伴奏
            vocals1 = np.ascontiguousarray(step1["vocals"].T, dtype=np.float32)

            # 3) 步骤2：vocals -> Vocals / Instrumental(和声)
            step2_stems = ["Vocals"] if remove_harmony else ["Vocals", "Instrumental"]
            step2 = SEPARATORS["karaoke_split"].separate(
                vocals1, pbar=False, stems=step2_stems
            )
            lead2 = np.ascontiguousarray(step2["Vocals"].T, dtype=np.float32)
            harmony = None
            if not remove_harmony:
                harmony = np.ascontiguousarray(step2["Instrumental"].T, dtype=np.float32)

            # 4) 步骤3：Vocals -> noreverb / reverb
            step3 = SEPARATORS["dereverb_split"].separate(
                lead2, pbar=False, stems=["noreverb", "reverb"]
            )
            noreverb = np.ascontiguousarray(step3["noreverb"].T, dtype=np.float32)  # (n, ch)
            reverb = np.ascontiguousarray(step3["reverb"].T, dtype=np.float32)

            # 5) RVC 变声：主人声（noreverb），内部统一 (ch, n)，RVC 只吃单声道
            _apply_voice_conf(voice_id, pitch, index_rate, protect, filter_radius, rms_mix_rate)
            lead_mono = _convert_voice(voice_id, noreverb.mean(axis=0), OUT_SR)
            lead = np.stack([lead_mono, lead_mono], axis=0)                          # (2, n)

            # 6) 可选：和声也过 RVC
            harmony_out = None
            if harmony is not None:
                if process_harmony:
                    harmony_mono = _convert_voice(voice_id, harmony.mean(axis=0), OUT_SR)
                    harmony_out = np.stack([harmony_mono, harmony_mono], axis=0)
                else:
                    harmony_out = harmony

            # 7) 合并：伴奏 + 变声主人声 + (可选)和声 + 混响，对齐时长后叠加 + 峰值归一化
            tracks = [other, lead]
            if harmony_out is not None:
                tracks.append(harmony_out)
            tracks.append(reverb)

            length = max(t.shape[1] for t in tracks)
            merged = np.zeros((2, length), dtype=np.float32)
            for track in tracks:
                merged += _fit_length(track, length)
            peak = float(np.max(np.abs(merged))) if merged.size else 0.0
            if peak > 0:
                merged *= MIX_PEAK / peak

            # 8) 落盘
            outputs = {
                "other": job_dir / "01_other.wav",
                "vocal": job_dir / "02_vocal_rvc.wav",
                "harmony": job_dir / "03_harmony.wav",
                "reverb": job_dir / "04_reverb.wav",
                "mix": job_dir / "05_mix.wav",
            }
            _write_wav(outputs["other"], other, OUT_SR)
            _write_wav(outputs["vocal"], lead, OUT_SR)
            if harmony_out is not None:
                _write_wav(outputs["harmony"], harmony_out, OUT_SR)
            _write_wav(outputs["reverb"], reverb, OUT_SR)
            _write_wav(outputs["mix"], merged, OUT_SR)

            elapsed = time.time() - started
            voice_label = VOICE_LABELS.get(voice_id, voice_id)
            notes = [
                f"音色：{voice_label}（变调 {int(pitch):+d} 半音）",
                f"时长：{total / OUT_SR:.1f}s 音频，耗时 {elapsed:.1f}s",
            ]
            notes.append(
                f"ZeroGPU：声明 {GPU_DECLARED_SECONDS}s → 预占配额 "
                f"{reserved_seconds(GPU_DECLARED_SECONDS)}s"
            )
            if remove_harmony:
                notes.append("已勾选「移除和声」：和声不参与流程与合成")
            elif process_harmony:
                notes.append("已勾选「和声处理」：和声也经过 RVC 变声")
            return (str(outputs["other"]), str(outputs["vocal"]),
                    str(outputs["harmony"]) if harmony_out is not None else None,
                    str(outputs["reverb"]), str(outputs["mix"]),
                    "\n\n".join(notes))
        except gr.Error:
            raise
        except Exception as exc:  # noqa: BLE001
            raise gr.Error(f"处理失败：{type(exc).__name__}: {exc}") from exc


# --------------------------------------------------------------------------------------
# 界面
# --------------------------------------------------------------------------------------
WORKFLOW_DOC = """
**工作流（三板斧）**

| 步骤 | 模型 | 输入 | 输出 |
| --- | --- | --- | --- |
| 1 | `big_beta6x` | 原始输入 | `vocals` / `other` |
| 2 | `mel_band_karaoke_fusion_standard` | 步骤1 `vocals` | `Vocals` / `Instrumental`(和声) |
| 3 | `dereverb_mel_band_roformer_anvuew` | 步骤2 `Vocals` | `noreverb` / `reverb` |

`noreverb` 作为主人声送入 RVC 变声；`Instrumental` 可作为和声一并混入。
"""

with gr.Blocks(title="RVC 翻唱工作站") as demo:
    gr.Markdown("# RVC 翻唱工作站\nPyMSS 三板斧分离 → RVC 变声 → 成品混音（ZeroGPU 加速）")

    if STARTUP_ERRORS:
        gr.Markdown("### ⚠️ 启动告警\n" + "\n".join(f"- {e}" for e in STARTUP_ERRORS))

    with gr.Row():
        with gr.Column(scale=1):
            audio_in = gr.Audio(label="上传歌曲", type="filepath", sources=["upload"])
            voice_dd = gr.Dropdown(
                label="RVC 音色",
                choices=[(VOICE_LABELS[k], k) for k in CONVERTERS],
                value=(list(CONVERTERS)[0] if CONVERTERS else None),
                interactive=bool(CONVERTERS),
            )
            pitch_sl = gr.Slider(-24, 24, value=0, step=1, label="变调（半音）")

            with gr.Accordion("RVC 常用参数", open=False):
                index_rate_sl = gr.Slider(0.0, 1.0, value=0.66, step=0.01, label="index 相似度 index_rate")
                protect_sl = gr.Slider(0.0, 0.5, value=0.33, step=0.01, label="气息/辅音保护 protect")
                filter_sl = gr.Slider(0, 7, value=3, step=1, label="音高中值滤波 filter_radius")
                rms_sl = gr.Slider(0.0, 1.0, value=0.25, step=0.01, label="音量包络混合 rms_mix_rate")

            with gr.Accordion("和声选项", open=True):
                remove_harmony_cb = gr.Checkbox(value=False, label="移除和声（和声不参与流程与合成）")
                process_harmony_cb = gr.Checkbox(value=False, label="和声处理（把和声也过 RVC 变声）")

            run_btn = gr.Button("开始分离 + 变声", variant="primary")
            status_md = gr.Markdown("")

        with gr.Column(scale=1):
            gr.Markdown(WORKFLOW_DOC)
            other_out = gr.Audio(label="① 伴奏（other）", type="filepath", interactive=False)
            vocal_out = gr.Audio(label="② 变声后的主人声（noreverb）", type="filepath", interactive=False)
            harmony_out = gr.Audio(label="③ 和声（Instrumental）", type="filepath", interactive=False)
            reverb_out = gr.Audio(label="④ 混响（reverb）", type="filepath", interactive=False)
            mix_out = gr.Audio(label="⑤ 合并成品", type="filepath", interactive=False)

    def _toggle_harmony(remove_flag):
        return gr.update(interactive=not remove_flag)

    remove_harmony_cb.change(_toggle_harmony, inputs=remove_harmony_cb, outputs=process_harmony_cb)

    run_btn.click(
        fn=generate,
        inputs=[audio_in, voice_dd, pitch_sl, index_rate_sl, protect_sl,
                filter_sl, rms_sl, remove_harmony_cb, process_harmony_cb],
        outputs=[other_out, vocal_out, harmony_out, reverb_out, mix_out, status_md],
    )

demo.queue().launch()