531 lines
18 KiB
Python
531 lines
18 KiB
Python
# https://github.com/RickyL-2000/ROSVOT
|
|
import math
|
|
import sys
|
|
import traceback
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Dict, Optional
|
|
|
|
import librosa
|
|
import numpy as np
|
|
import torch
|
|
import matplotlib.pyplot as plt
|
|
|
|
from .utils.os_utils import safe_path
|
|
from .utils.commons.hparams import set_hparams
|
|
from .utils.commons.ckpt_utils import load_ckpt
|
|
from .utils.commons.dataset_utils import pad_or_cut_xd
|
|
from .utils.audio.mel import MelNet
|
|
from .utils.audio.pitch_utils import (
|
|
norm_interp_f0,
|
|
denorm_f0,
|
|
f0_to_coarse,
|
|
boundary2Interval,
|
|
save_midi,
|
|
midi_to_hz,
|
|
)
|
|
from .utils.rosvot_utils import (
|
|
get_mel_len,
|
|
align_word,
|
|
regulate_real_note_itv,
|
|
regulate_ill_slur,
|
|
bd_to_durs,
|
|
)
|
|
from .modules.pe.rmvpe import RMVPE
|
|
from .modules.rosvot.rosvot import MidiExtractor, WordbdExtractor
|
|
|
|
|
|
@torch.no_grad()
|
|
def infer_sample(
|
|
item: Dict[str, Any],
|
|
hparams: Dict[str, Any],
|
|
models: Dict[str, Any],
|
|
device: torch.device,
|
|
*,
|
|
save_dir: Optional[str] = None,
|
|
apply_rwbd: Optional[bool] = None,
|
|
# outputs
|
|
save_plot: bool = False,
|
|
no_save_midi: bool = True,
|
|
no_save_npy: bool = True,
|
|
verbose: bool = False,
|
|
) -> Dict[str, Any]:
|
|
if "item_name" not in item or "wav_fn" not in item:
|
|
raise ValueError('item must contain keys: "item_name" and "wav_fn"')
|
|
|
|
item_name = item["item_name"]
|
|
wav_src = item["wav_fn"]
|
|
|
|
# Decide RWBD usage
|
|
if apply_rwbd is None:
|
|
apply_rwbd_ = ("word_durs" not in item)
|
|
else:
|
|
apply_rwbd_ = bool(apply_rwbd)
|
|
|
|
# Models
|
|
model = models["model"]
|
|
mel_net = models["mel_net"]
|
|
pe = models.get("pe")
|
|
wbd_predictor = models.get("wbd_predictor")
|
|
|
|
if wbd_predictor is None and apply_rwbd_:
|
|
raise ValueError("apply_rwbd is True but wbd_predictor model is not provided in models")
|
|
|
|
# ---- Prepare Data ----
|
|
if isinstance(wav_src, str):
|
|
wav, _ = librosa.core.load(wav_src, sr=hparams["audio_sample_rate"])
|
|
else:
|
|
wav = wav_src
|
|
if not isinstance(wav, np.ndarray):
|
|
wav = np.asarray(wav)
|
|
wav = wav.astype(np.float32)
|
|
|
|
# Calculate timestamps and alignment lengths
|
|
wav_len_samples = wav.shape[-1]
|
|
mel_len = get_mel_len(wav_len_samples, hparams["hop_size"])
|
|
|
|
# Word boundary preparation
|
|
mel2word = None
|
|
word_durs_filtered = None
|
|
|
|
if not apply_rwbd_:
|
|
if "word_durs" not in item:
|
|
raise ValueError('apply_rwbd=False but item has no "word_durs"')
|
|
|
|
wd_raw = list(item["word_durs"])
|
|
min_word_dur = hparams.get("min_word_dur", 20) / 1000
|
|
word_durs_filtered = []
|
|
|
|
for i, wd in enumerate(wd_raw):
|
|
if wd < min_word_dur:
|
|
if i == 0 and len(wd_raw) > 1:
|
|
wd_raw[i + 1] += wd
|
|
elif len(word_durs_filtered) > 0:
|
|
word_durs_filtered[-1] += wd
|
|
else:
|
|
word_durs_filtered.append(wd)
|
|
|
|
mel2word, _ = align_word(word_durs_filtered, mel_len, hparams["hop_size"], hparams["audio_sample_rate"])
|
|
mel2word = np.asarray(mel2word)
|
|
if mel2word.size > 0 and mel2word[0] == 0:
|
|
mel2word = mel2word + 1
|
|
|
|
mel2word_len = int(np.sum(mel2word > 0))
|
|
real_len = min(mel_len, mel2word_len)
|
|
else:
|
|
real_len = min(mel_len, hparams["max_frames"])
|
|
|
|
T = math.ceil(min(real_len, hparams["max_frames"]) / hparams["frames_multiple"]) * hparams["frames_multiple"]
|
|
|
|
# ---- Input Tensors & Padding ----
|
|
target_samples = T * hparams["hop_size"]
|
|
wav_t = torch.from_numpy(wav).float().to(device).unsqueeze(0) # [1, L]
|
|
if wav_t.shape[-1] < target_samples:
|
|
wav_t = pad_or_cut_xd(wav_t, target_samples, 1)
|
|
|
|
# ---- Pitch Extraction ----
|
|
if pe is not None:
|
|
f0s, uvs = pe.get_pitch_batch(
|
|
wav_t,
|
|
sample_rate=hparams["audio_sample_rate"],
|
|
hop_size=hparams["hop_size"],
|
|
lengths=[real_len],
|
|
fmax=hparams["f0_max"],
|
|
fmin=hparams["f0_min"],
|
|
)
|
|
f0_1d, uv_1d = norm_interp_f0(f0s[0][:T])
|
|
f0_t = pad_or_cut_xd(torch.FloatTensor(f0_1d).to(device), T, 0).unsqueeze(0)
|
|
uv_t = pad_or_cut_xd(torch.FloatTensor(uv_1d).to(device), T, 0).long().unsqueeze(0)
|
|
pitch_coarse = f0_to_coarse(denorm_f0(f0_t, uv_t)).to(device)
|
|
f0_np = denorm_f0(f0_t, uv_t)[0].detach().cpu().numpy()[:real_len]
|
|
else:
|
|
f0_t = uv_t = pitch_coarse = None
|
|
f0_np = None
|
|
|
|
# ---- Mel Extraction ----
|
|
mel = mel_net(wav_t) # [1, T_padded, C]
|
|
mel = pad_or_cut_xd(mel, T, 1)
|
|
|
|
# Construct non-padding mask
|
|
mel_nonpadding_mask = torch.zeros(1, T, device=device)
|
|
mel_nonpadding_mask[:, :real_len] = 1.0
|
|
|
|
# Apply mask to mel (zero out padding)
|
|
mel = (mel.transpose(1, 2) * mel_nonpadding_mask.unsqueeze(1)).transpose(1, 2)
|
|
# Re-calculate non_padding bool mask
|
|
mel_nonpadding = mel.abs().sum(-1) > 0
|
|
|
|
# ---- Word Boundary ----
|
|
word_durs_used = None
|
|
if apply_rwbd_:
|
|
mel_input = mel[:, :, : hparams.get("wbd_use_mel_bins", 80)]
|
|
wbd_outputs = wbd_predictor(
|
|
mel=mel_input,
|
|
pitch=pitch_coarse,
|
|
uv=uv_t,
|
|
non_padding=mel_nonpadding,
|
|
train=False,
|
|
)
|
|
word_bd = wbd_outputs["word_bd_pred"] # [1, T]
|
|
else:
|
|
# Construct word_bd from provided durs
|
|
mel2word_t = pad_or_cut_xd(torch.LongTensor(mel2word).to(device), T, 0)
|
|
word_bd = torch.zeros_like(mel2word_t)
|
|
# Vectorized check
|
|
word_bd[1:] = (mel2word_t[1:] != mel2word_t[:-1]).long()
|
|
word_bd[real_len:] = 0
|
|
word_bd = word_bd.unsqueeze(0) # [1, T]
|
|
|
|
word_durs_used = np.array(word_durs_filtered)
|
|
|
|
# ---- Main Inference ----
|
|
mel_input = mel[:, :, : hparams.get("use_mel_bins", 80)]
|
|
outputs = model(
|
|
mel=mel_input,
|
|
word_bd=word_bd,
|
|
pitch=pitch_coarse,
|
|
uv=uv_t,
|
|
non_padding=mel_nonpadding,
|
|
train=False,
|
|
)
|
|
|
|
note_lengths = outputs["note_lengths"].detach().cpu().numpy()
|
|
note_bd_pred = outputs["note_bd_pred"][0].detach().cpu().numpy()[:real_len]
|
|
note_pred = outputs["note_pred"][0].detach().cpu().numpy()[: note_lengths[0]]
|
|
note_bd_logits = torch.sigmoid(outputs["note_bd_logits"])[0].detach().cpu().numpy()[:real_len]
|
|
|
|
if note_pred.shape == (0,):
|
|
if verbose:
|
|
print(f"skip {item_name}: no notes detected")
|
|
return {
|
|
"item_name": item_name,
|
|
"pitches": [],
|
|
"note_durs": [],
|
|
"note2words": None,
|
|
}
|
|
|
|
# ---- Post-Processing & Regulation ----
|
|
note_itv_pred = boundary2Interval(note_bd_pred)
|
|
note2words = None
|
|
|
|
if apply_rwbd_:
|
|
word_bd_np = outputs['word_bd_pred'][0].detach().cpu().numpy()[:real_len]
|
|
word_durs_derived = np.array(bd_to_durs(word_bd_np)) * hparams['hop_size'] / hparams['audio_sample_rate']
|
|
word_durs_for_reg = word_durs_derived
|
|
word_bd_for_reg = word_bd_np
|
|
else:
|
|
word_bd_for_reg = word_bd[0].detach().cpu().numpy()[:real_len]
|
|
word_durs_for_reg = word_durs_used
|
|
|
|
should_regulate = hparams.get("infer_regulate_real_note_itv", True) and (not apply_rwbd_)
|
|
|
|
if should_regulate and (word_durs_for_reg is not None):
|
|
try:
|
|
note_itv_pred_secs, note2words = regulate_real_note_itv(
|
|
note_itv_pred,
|
|
note_bd_pred,
|
|
word_bd_for_reg,
|
|
word_durs_for_reg,
|
|
hparams["hop_size"],
|
|
hparams["audio_sample_rate"],
|
|
)
|
|
note_pred, note_itv_pred_secs, note2words = regulate_ill_slur(note_pred, note_itv_pred_secs, note2words)
|
|
except Exception as err:
|
|
if verbose:
|
|
_, exc_value, exc_tb = sys.exc_info()
|
|
tb = traceback.extract_tb(exc_tb)[-1]
|
|
print(f"postprocess failed: {err}: {exc_value} in {tb[0]}:{tb[1]} '{tb[2]}' in {tb[3]}")
|
|
# Fallback
|
|
note_itv_pred_secs = note_itv_pred * hparams["hop_size"] / hparams["audio_sample_rate"]
|
|
note2words = None
|
|
else:
|
|
note_itv_pred_secs = note_itv_pred * hparams["hop_size"] / hparams["audio_sample_rate"]
|
|
|
|
# ---- Output ----
|
|
note_durs = [float((itv[1] - itv[0])) for itv in note_itv_pred_secs]
|
|
|
|
out = {
|
|
"item_name": item_name,
|
|
"pitches": note_pred.tolist(),
|
|
"note_durs": note_durs,
|
|
"note2words": note2words.tolist() if note2words is not None else None,
|
|
}
|
|
|
|
# ---- Saving ----
|
|
if save_dir is not None:
|
|
save_dir_path = Path(save_dir)
|
|
save_dir_path.mkdir(parents=True, exist_ok=True)
|
|
fn = str(item_name)
|
|
|
|
if not no_save_midi:
|
|
save_midi(note_pred, note_itv_pred_secs, safe_path(save_dir_path / "midi" / f"{fn}.mid"))
|
|
|
|
if not no_save_npy:
|
|
np.save(safe_path(save_dir_path / "npy" / f"[note]{fn}.npy"), out, allow_pickle=True)
|
|
|
|
if save_plot:
|
|
fig = plt.figure()
|
|
if f0_np is not None:
|
|
plt.plot(f0_np, color="red", label="f0")
|
|
|
|
midi_pred = np.zeros(note_bd_pred.shape[0], dtype=np.float32)
|
|
itvs = np.round(note_itv_pred_secs * hparams["audio_sample_rate"] / hparams["hop_size"]).astype(int)
|
|
for i, itv in enumerate(itvs):
|
|
midi_pred[itv[0] : itv[1]] = note_pred[i]
|
|
plt.plot(midi_to_hz(midi_pred), color="blue", label="pred midi")
|
|
plt.plot(note_bd_logits * 100, color="green", label="note bd logits x100")
|
|
plt.legend()
|
|
plt.tight_layout()
|
|
plt.savefig(safe_path(save_dir_path / "plot" / f"[MIDI]{fn}.png"), format="png")
|
|
plt.close(fig)
|
|
|
|
return out
|
|
|
|
|
|
def load_rosvot_models(ckpt, config="", wbd_ckpt="", wbd_config="", device="cuda:0", verbose=False, thr=0.85):
|
|
"""
|
|
Load models once to reuse across multiple items.
|
|
"""
|
|
dev = torch.device(device)
|
|
|
|
# 1. Hparams
|
|
config_path = Path(ckpt).with_name("config.yaml") if config == "" else config
|
|
pe_ckpt = Path(ckpt).parent.parent / "rmvpe/model.pt"
|
|
hparams = set_hparams(
|
|
config=config_path,
|
|
print_hparams=verbose,
|
|
hparams_str=f"note_bd_threshold={thr}",
|
|
)
|
|
|
|
# 2. Main Model
|
|
model = MidiExtractor(hparams)
|
|
load_ckpt(model, ckpt, verbose=verbose)
|
|
model.eval().to(dev)
|
|
|
|
# 3. MelNet
|
|
mel_net = MelNet(hparams)
|
|
mel_net.to(dev)
|
|
|
|
# 4. Pitch Extractor
|
|
pe = None
|
|
if hparams.get("use_pitch_embed", False):
|
|
pe = RMVPE(pe_ckpt, device=dev)
|
|
|
|
# 5. Word Boundary Predictor (optional but we load if ckpt provided or needed)
|
|
wbd_predictor = None
|
|
if wbd_ckpt:
|
|
wbd_config_path = Path(wbd_ckpt).with_name("config.yaml") if wbd_config == "" else wbd_config
|
|
wbd_hparams = set_hparams(
|
|
config=wbd_config_path,
|
|
print_hparams=False,
|
|
hparams_str="",
|
|
)
|
|
hparams.update({
|
|
"wbd_use_mel_bins": wbd_hparams["use_mel_bins"],
|
|
"min_word_dur": wbd_hparams["min_word_dur"],
|
|
})
|
|
wbd_predictor = WordbdExtractor(wbd_hparams)
|
|
load_ckpt(wbd_predictor, wbd_ckpt, verbose=verbose)
|
|
wbd_predictor.eval().to(dev)
|
|
|
|
models = {
|
|
"model": model,
|
|
"mel_net": mel_net,
|
|
"pe": pe,
|
|
"wbd_predictor": wbd_predictor
|
|
}
|
|
return hparams, models
|
|
|
|
|
|
class NoteTranscriber:
|
|
"""Note transcription wrapper based on ROSVOT.
|
|
|
|
Loads ROSVOT and optional RWBD models once in ``__init__`` and
|
|
exposes a :py:meth:`process` API that turns an item dict into
|
|
aligned note metadata for downstream SVS.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
rosvot_model_path: str,
|
|
rwbd_model_path: str,
|
|
*,
|
|
rosvot_config_path: str = "",
|
|
rwbd_config_path: str = "",
|
|
device: str = "cuda:0",
|
|
thr: float = 0.85,
|
|
verbose: bool = True,
|
|
):
|
|
"""Initialize the note transcriber.
|
|
|
|
Args:
|
|
ckpt: Path to the main ROSVOT checkpoint.
|
|
config: Optional config YAML path for ROSVOT.
|
|
wbd_ckpt: Optional word-boundary checkpoint path.
|
|
wbd_config: Optional config YAML path for RWBD.
|
|
device: Torch device string, e.g. ``"cuda:0"`` / ``"cpu"``.
|
|
thr: Note boundary threshold.
|
|
verbose: Whether to print verbose logs.
|
|
"""
|
|
self.verbose = verbose
|
|
self.device = torch.device(device)
|
|
self.hparams, self.models = load_rosvot_models(
|
|
ckpt=rosvot_model_path,
|
|
config=rosvot_config_path,
|
|
wbd_ckpt=rwbd_model_path,
|
|
wbd_config=rwbd_config_path,
|
|
device=device,
|
|
verbose=verbose,
|
|
thr=thr,
|
|
)
|
|
|
|
if self.verbose:
|
|
print(
|
|
"[note transcription] init success:",
|
|
f"device={self.device}",
|
|
f"rosvot_model_path={rosvot_model_path}",
|
|
f"rwbd_model_path={rwbd_model_path if rwbd_model_path else 'None'}",
|
|
f"thr={thr}",
|
|
)
|
|
|
|
def process(
|
|
self,
|
|
item: Dict[str, Any],
|
|
*,
|
|
segment_info: Optional[Dict[str, Any]] = None,
|
|
save_dir: Optional[str] = None,
|
|
apply_rwbd: Optional[bool] = None,
|
|
save_plot: bool = False,
|
|
no_save_midi: bool = True,
|
|
no_save_npy: bool = True,
|
|
verbose: Optional[bool] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Run ROSVOT on a single item and post-process outputs.
|
|
|
|
Args:
|
|
item: Input metadata dict with at least ``item_name`` and ``wav_fn``.
|
|
segment_info: Optional segment metadata for sliced audio.
|
|
save_dir: Optional directory for debug artifacts (plots, midis).
|
|
apply_rwbd: Whether to run RWBD-based word boundary refinement.
|
|
save_plot: Whether to save diagnostic plots.
|
|
no_save_midi: If True, skip saving midi.
|
|
no_save_npy: If True, skip saving numpy intermediates.
|
|
verbose: Override instance-level verbose flag for this call.
|
|
|
|
Returns:
|
|
Dict with aligned note information for downstream SVS.
|
|
"""
|
|
v = self.verbose if verbose is None else verbose
|
|
if v:
|
|
item_name = item.get("item_name", "")
|
|
wav_fn = item.get("wav_fn", "")
|
|
print(f"[note transcription] process: start: item_name={item_name} wav_fn={wav_fn}")
|
|
t0 = time.time()
|
|
|
|
rosvot_out = infer_sample(
|
|
item,
|
|
self.hparams,
|
|
self.models,
|
|
device=self.device,
|
|
save_dir=save_dir,
|
|
apply_rwbd=apply_rwbd,
|
|
save_plot=save_plot,
|
|
no_save_midi=no_save_midi,
|
|
no_save_npy=no_save_npy,
|
|
verbose=v,
|
|
)
|
|
|
|
out = self.post_process(
|
|
metadata=item,
|
|
segment_info=segment_info,
|
|
rosvot_out=rosvot_out,
|
|
)
|
|
|
|
if v:
|
|
dt = time.time() - t0
|
|
print(
|
|
"[note transcription] process: done:",
|
|
f"item_name={out.get('item_name','')}",
|
|
f"n_notes={len(out.get('note_pitch', []) or [])}",
|
|
f"time={dt:.3f}s",
|
|
)
|
|
|
|
return out
|
|
|
|
@staticmethod
|
|
def _normalize_note2words(note2words: list[int]) -> list[int]:
|
|
if not note2words:
|
|
return []
|
|
normalized = [note2words[0]]
|
|
for idx in range(1, len(note2words)):
|
|
if note2words[idx] < normalized[-1]:
|
|
normalized.append(normalized[-1])
|
|
else:
|
|
normalized.append(note2words[idx])
|
|
return normalized
|
|
|
|
@staticmethod
|
|
def _build_ep_types(note2words: list[int], align_words: list[str]) -> list[int]:
|
|
ep_types: list[int] = []
|
|
prev = -1
|
|
for i, w in zip(note2words, align_words):
|
|
if w == "<SP>":
|
|
ep_types.append(1)
|
|
else:
|
|
ep_types.append(2 if i != prev else 3)
|
|
prev = i
|
|
return ep_types
|
|
|
|
def post_process(
|
|
self,
|
|
*,
|
|
metadata: Dict[str, Any],
|
|
segment_info: Dict[str, Any],
|
|
rosvot_out: Dict[str, Any],
|
|
) -> Dict[str, Any]:
|
|
"""Build aligned note metadata using ROSVOT outputs."""
|
|
note2words_raw = rosvot_out.get("note2words") or []
|
|
note2words = self._normalize_note2words(note2words_raw)
|
|
align_words = [
|
|
metadata["words"][idx - 1]
|
|
for idx in note2words_raw
|
|
if 0 < idx <= len(metadata["words"])
|
|
]
|
|
ep_types = self._build_ep_types(note2words, align_words) if align_words else []
|
|
|
|
return {
|
|
"item_name": rosvot_out.get("item_name", "") if not segment_info else segment_info["item_name"],
|
|
"wav_fn": metadata.get("wav_fn", "") if not segment_info else segment_info["wav_fn"],
|
|
"origin_wav_fn": metadata.get("origin_wav_fn", "") if not segment_info else segment_info["origin_wav_fn"],
|
|
"start_time_ms": "" if not segment_info else segment_info["start_time_ms"],
|
|
"end_time_ms": "" if not segment_info else segment_info["end_time_ms"],
|
|
"language": metadata.get("language", ""),
|
|
"note_text": align_words,
|
|
"note_dur": rosvot_out.get("note_durs", []),
|
|
"note_type": ep_types,
|
|
"note_pitch": rosvot_out.get("pitches", []),
|
|
}
|
|
|
|
if __name__ == "__main__":
|
|
|
|
item = {
|
|
'item_name': 'vocal_0',
|
|
'wav_fn': 'example/audio/zh_prompt.mp3',
|
|
'start_time_ms': 320,
|
|
'end_time_ms': 10687,
|
|
'origin_wav_fn': 'example/audio/zh_prompt.mp3',
|
|
'duration': 10367,
|
|
'words': ['<SP>', '除', '了', '想', '你', '<SP>', '除', '了', '爱', '你', '<SP>', '我', '什', '么', '什', '么', '都', '愿', '意'],
|
|
'word_durs': [0.21, 0.36, 0.26, 0.7000000000000001, 0.96, 0.3800000000000001, 0.43999999999999995, 0.3799999999999999, 0.6400000000000001, 0.9600000000000002, 1.1199999999999999, 0.28000000000000025, 0.3799999999999999, 0.3199999999999994, 0.3200000000000003, 0.3799999999999999, 0.3200000000000003, 0.5, 1.457981859410431],
|
|
'language': 'Mandarin'
|
|
}
|
|
|
|
m = NoteTranscriber(
|
|
rosvot_model_path="pretrained_models/SoulX-Singer-Preprocess/rosvot/rosvot/model.pt",
|
|
rwbd_model_path="pretrained_models/SoulX-Singer-Preprocess/rosvot/rwbd/model.pt",
|
|
device="cuda"
|
|
)
|
|
out = m.process(item, segment_info=item)
|
|
|
|
print(out) |