vllm-omni/vllm_omni/diffusion/models/soulx_singer/modules/preprocess/rosvot.py

"""ROSVOT note transcription."""

from __future__ import annotations

import math
from pathlib import Path
from typing import Any, ClassVar

import numpy as np
import torch
import torch.nn as nn
from vllm.logger import init_logger

from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery
from vllm_omni.diffusion.models.soulx_singer.modules.mel_transform import spectral_normalize_torch
from vllm_omni.diffusion.models.soulx_singer.modules.note_transcription import (
    BackboneNet,
    ConvBlocks,
    Embedding,
)
from vllm_omni.diffusion.models.soulx_singer.modules.preprocess.utils import (
    boundary2Interval,
    denorm_f0,
    f0_to_coarse,
    load_model_ckpt,
    load_mono_audio,
    load_rosvot_config,
    norm_interp_f0,
    pad_or_cut_xd,
    resample_mono,
)
from vllm_omni.utils.audio import mel_filter_bank as librosa_mel_fn

logger = init_logger(__name__)


def regulate_real_note_itv(
    note_itv: np.ndarray,
    note_bd: np.ndarray,
    word_bd: np.ndarray,
    word_durs: np.ndarray,
    hop_size: int,
    audio_sample_rate: int,
) -> tuple[np.ndarray, np.ndarray]:
    assert note_itv.shape[0] == np.sum(note_bd) + 1
    assert np.sum(word_bd) <= np.sum(note_bd)
    assert word_durs.shape[0] == np.sum(word_bd) + 1, f"{word_durs.shape[0]} {np.sum(word_bd) + 1}"
    word_bd = np.cumsum(word_bd) * word_bd
    word_itv = np.zeros((word_durs.shape[0], 2))
    word_offsets = np.cumsum(word_durs)
    note2words = np.zeros(note_itv.shape[0], dtype=int)
    for idx in range(len(word_offsets) - 1):
        word_itv[idx, 1] = word_itv[idx + 1, 0] = word_offsets[idx]
    word_itv[-1, 1] = word_offsets[-1]
    note_itv_secs = note_itv * hop_size / audio_sample_rate
    for idx, itv in enumerate(note_itv):
        start_idx, end_idx = itv
        if word_bd[start_idx] > 0:
            word_dur_idx = word_bd[start_idx]
            note_itv_secs[idx, 0] = word_itv[word_dur_idx, 0]
            note2words[idx] = word_dur_idx
        if word_bd[end_idx] > 0:
            word_dur_idx = word_bd[end_idx] - 1
            note_itv_secs[idx, 1] = word_itv[word_dur_idx, 1]
            note2words[idx] = word_dur_idx
    note2words += 1
    return note_itv_secs, note2words


def regulate_ill_slur(
    notes: np.ndarray, note_itv: np.ndarray, note2words: np.ndarray
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    res_note2words: list[int] = []
    res_note_itv: list[list[float]] = []
    res_notes: list[int] = []
    note_idx = 0
    note_idx_end = 0
    while note_idx <= len(notes) - 1:
        while note_idx <= note_idx_end < len(notes) and note2words[note_idx] == note2words[note_idx_end]:
            note_idx_end += 1
        res_note2words.append(note2words[note_idx])
        res_note_itv.append(note_itv[note_idx].tolist())
        res_notes.append(notes[note_idx])
        for idx in range(note_idx + 1, note_idx_end):
            if notes[idx] == notes[idx - 1]:
                res_note_itv[-1][1] = note_itv[idx][1]
            else:
                res_note_itv.append(note_itv[idx].tolist())
                res_note2words.append(note2words[idx])
                res_notes.append(notes[idx])
        note_idx = note_idx_end
    return (
        np.array(res_notes, dtype=notes.dtype),
        np.array(res_note_itv, dtype=note_itv.dtype),
        np.array(res_note2words, dtype=note2words.dtype),
    )


def regulate_boundary(bd_logits, threshold, min_gap=18, ref_bd=None, ref_bd_min_gap=8, non_padding=None):
    # this doesn't preserve gradient
    device = bd_logits.device
    bd_logits = torch.sigmoid(bd_logits).data.cpu()
    # bd_logits[0] = bd_logits[-1] = 1e-5     # avoid itv invalid problem
    bd = (bd_logits > threshold).long()
    bd_res = torch.zeros_like(bd).long()
    for i in range(bd.shape[0]):
        bd_i = bd[i]
        last_bd_idx = -1
        start = -1
        for j in range(bd_i.shape[0]):
            if bd_i[j] == 1:
                if 0 <= start < j:
                    continue
                elif start < 0:
                    start = j
            else:
                if 0 <= start < j:
                    if j - 1 > start:
                        bd_idx = start + int(torch.argmax(bd_logits[i, start:j]).item())
                    else:
                        bd_idx = start
                    if bd_idx - last_bd_idx < min_gap and last_bd_idx > 0:
                        bd_idx = round((bd_idx + last_bd_idx) / 2)
                        bd_res[i, last_bd_idx] = 0
                    bd_res[i, bd_idx] = 1
                    last_bd_idx = bd_idx
                    start = -1

    # assert ref_bd_min_gap <= min_gap // 2
    if ref_bd is not None and ref_bd_min_gap > 0:
        ref = ref_bd.data.cpu()
        for i in range(bd_res.shape[0]):
            ref_bd_i = ref[i]
            ref_bd_i_js = []
            for j in range(ref_bd_i.shape[0]):
                if ref_bd_i[j] == 1:
                    ref_bd_i_js.append(j)
                    seg_sum = torch.sum(bd_res[i, max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap])
                    if seg_sum == 0:
                        bd_res[i, j] = 1
                    elif seg_sum == 1 and bd_res[i, j] != 1:
                        bd_res[i, max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap] = ref_bd_i[
                            max(0, j - ref_bd_min_gap) : j + ref_bd_min_gap
                        ]
                    elif seg_sum > 1:
                        for k in range(1, ref_bd_min_gap + 1):
                            if bd_res[i, max(0, j - k)] == 1 and ref_bd_i[max(0, j - k)] != 1:
                                bd_res[i, max(0, j - k)] = 0
                                break
                            if (
                                bd_res[i, min(bd_res.shape[1] - 1, j + k)] == 1
                                and ref_bd_i[min(bd_res.shape[1] - 1, j + k)] != 1
                            ):
                                bd_res[i, min(bd_res.shape[1] - 1, j + k)] = 0
                                break
                        bd_res[i, j] = 1
            # final check
            assert torch.sum(bd_res[i, ref_bd_i_js]) == len(ref_bd_i_js), (
                f"{torch.sum(bd_res[i, ref_bd_i_js])} {len(ref_bd_i_js)}"
            )

    bd_res = bd_res.to(device)

    # force valid begin and end
    bd_res[:, 0] = 0
    if non_padding is not None:
        for i in range(bd_res.shape[0]):
            bd_res[i, sum(non_padding[i]) - 1 :] = 0
    else:
        bd_res[:, -1] = 0

    return bd_res


def _align_word(word_durs: list[float], mel_len: int, hop_size: int, audio_sample_rate: int) -> np.ndarray:
    mel2word = np.zeros([mel_len], int)
    start_time = 0.0
    for i_word, wd in enumerate(word_durs):
        start_frame = int(start_time * audio_sample_rate / hop_size + 0.5)
        end_frame = int((start_time + wd) * audio_sample_rate / hop_size + 0.5)
        mel2word[start_frame:end_frame] = i_word + 1
        start_time += wd
    return mel2word


class PitchDecoder(nn.Module):
    def __init__(self, hparams):
        super().__init__()
        self.hidden_size = hidden_size = hparams["hidden_size"]
        self.dropout = hparams.get("dropout", 0.0)
        self.note_bd_out = nn.Linear(hidden_size, 1)
        self.note_bd_temperature = max(1e-7, hparams.get("note_bd_temperature", 1.0))

        # note prediction
        self.pitch_attn_num_head = hparams.get("pitch_attn_num_head", 1)
        self.multihead_dot_attn = nn.Linear(hidden_size, self.pitch_attn_num_head)
        self.post = ConvBlocks(
            hidden_size,
            out_dims=hidden_size,
            dilations=None,
            kernel_size=3,
            layers_in_block=1,
            c_multiple=1,
            post_net_kernel=3,
        )
        self.pitch_out = nn.Linear(hidden_size, hparams.get("note_num", 100) + 4)
        self.note_num = hparams.get("note_num", 100)
        self.note_start = hparams.get("note_start", 30)
        self.pitch_temperature = max(1e-7, hparams.get("note_pitch_temperature", 1.0))

    def forward(self, feat, note_bd, train=True):
        bsz, T, _ = feat.shape

        attn = torch.sigmoid(self.multihead_dot_attn(feat))  # [B, T, C] -> [B, T, num_head]
        attn_feat = feat.unsqueeze(3) * attn.unsqueeze(2)  # [B, T, C, 1] x [B, T, 1, num_head] -> [B, T, C, num_head]
        attn_feat = torch.mean(attn_feat, dim=-1)  # [B, T, C, num_head] -> [B, T, C]
        mel2note = torch.cumsum(note_bd, 1)
        note_length = torch.max(torch.sum(note_bd, dim=1)).item() + 1  # max length
        note_lengths = torch.sum(note_bd, dim=1) + 1  # [B]

        attn = torch.mean(attn, dim=-1, keepdim=True)  # [B, T, num_head] -> [B, T, 1]
        denom = mel2note.new_zeros(bsz, note_length, dtype=attn.dtype).scatter_add_(
            dim=1, index=mel2note, src=attn.squeeze(-1)
        )  # [B, T] -> [B, note_length] count the note frames of each note (with padding excluded)
        frame2note = mel2note.unsqueeze(-1).repeat(1, 1, self.hidden_size)  # [B, T] -> [B, T, C], with padding included
        note_aggregate = frame2note.new_zeros(bsz, note_length, self.hidden_size, dtype=attn_feat.dtype).scatter_add_(
            dim=1, index=frame2note, src=attn_feat
        )  # [B, T, C] -> [B, note_length, C]
        note_aggregate = note_aggregate / (denom.unsqueeze(-1) + 1e-5)
        note_logits = self.post(note_aggregate)
        note_logits = self.pitch_out(note_logits) / self.pitch_temperature
        # note_logits = torch.clamp(note_logits, min=-16., max=16.)     # don't know need it or not

        note_pred = torch.softmax(note_logits, dim=-1)  # [B, note_length, note_num]
        note_pred = torch.argmax(note_pred, dim=-1)  # [B, note_length]
        # for some reason, note idx maybe 130 (why?)
        note_pred[note_pred > self.note_num] = 0
        note_pred[note_pred < self.note_start] = 0

        return note_lengths, note_logits, note_pred


class MidiExtractor(nn.Module):
    def __init__(self, hparams):
        super().__init__()
        self.hparams = hparams.copy()
        self.hidden_size = hidden_size = hparams["hidden_size"]
        self.note_bd_threshold = hparams.get("note_bd_threshold", 0.5)
        self.note_bd_min_gap = round(
            hparams.get("note_bd_min_gap", 100) * hparams["audio_sample_rate"] / 1000 / hparams["hop_size"]
        )
        self.note_bd_ref_min_gap = round(
            hparams.get("note_bd_ref_min_gap", 50) * hparams["audio_sample_rate"] / 1000 / hparams["hop_size"]
        )

        self.mel_proj = nn.Conv1d(hparams["use_mel_bins"], hidden_size, kernel_size=3, padding=1)
        self.mel_encoder = ConvBlocks(
            hidden_size,
            out_dims=hidden_size,
            dilations=None,
            kernel_size=3,
            layers_in_block=2,
            c_multiple=1,
            post_net_kernel=3,
        )
        self.use_pitch = hparams.get("use_pitch_embed", True)
        if self.use_pitch:
            self.pitch_embed = Embedding(300, hidden_size, 0, "kaiming")
            self.uv_embed = Embedding(3, hidden_size, 0, "kaiming")
        self.use_wbd = hparams.get("use_wbd", True)
        if self.use_wbd:
            self.word_bd_embed = Embedding(3, hidden_size, 0, "kaiming")
        self.cond_encoder = ConvBlocks(
            hidden_size,
            out_dims=hidden_size,
            dilations=None,
            kernel_size=3,
            layers_in_block=1,
            c_multiple=1,
            post_net_kernel=3,
        )

        # backbone
        self.net = BackboneNet(hparams)

        # note bd prediction
        self.note_bd_out = nn.Linear(hidden_size, 1)
        self.note_bd_temperature = max(1e-7, hparams.get("note_bd_temperature", 1.0))

        # note prediction
        self.pitch_decoder = PitchDecoder(hparams)

    def run_encoder(self, mel, word_bd=None, pitch=None, uv=None):
        mel_embed = self.mel_proj(mel.transpose(1, 2)).transpose(1, 2)
        mel_embed = self.mel_encoder(mel_embed)
        pitch_embed = word_bd_embed = 0
        if self.use_pitch and pitch is not None and uv is not None:
            pitch_embed = self.pitch_embed(pitch) + self.uv_embed(uv)  # [B, T, C]
        if self.use_wbd and word_bd is not None:
            word_bd_embed = self.word_bd_embed(word_bd)
        feat = self.cond_encoder(mel_embed + pitch_embed + word_bd_embed)
        return feat

    def forward(self, mel, word_bd=None, note_bd=None, pitch=None, uv=None, non_padding=None, train=True):
        ret = {}
        bsz, T, _ = mel.shape

        feat = self.run_encoder(mel, word_bd, pitch, uv)
        feat = self.net(feat)  # [B, T, C]

        # note bd prediction
        # dropout has been dropped (inference mode)
        note_bd_logits = self.note_bd_out(feat).squeeze(-1) / self.note_bd_temperature
        note_bd_logits = torch.clamp(note_bd_logits, min=-16.0, max=16.0)
        ret["note_bd_logits"] = note_bd_logits  # [B, T]
        if note_bd is None or not train:
            note_bd = regulate_boundary(
                note_bd_logits,
                self.note_bd_threshold,
                self.note_bd_min_gap,
                word_bd,
                self.note_bd_ref_min_gap,
                non_padding,
            )
            ret["note_bd_pred"] = note_bd  # [B, T]

        # note pitch prediction
        note_lengths, note_logits, note_pred = self.pitch_decoder(feat, note_bd, train)
        ret["note_lengths"], ret["note_logits"], ret["note_pred"] = note_lengths, note_logits, note_pred

        return ret


class MelNet(nn.Module):
    def __init__(self, hparams, device="cpu") -> None:
        super().__init__()
        self.n_fft = hparams["fft_size"]
        self.num_mels = hparams["audio_num_mel_bins"]
        self.sampling_rate = hparams["audio_sample_rate"]
        self.hop_size = hparams["hop_size"]
        self.win_size = hparams["win_size"]
        self.fmin = hparams["fmin"]
        self.fmax = hparams["fmax"]
        self.device = device

        mel = librosa_mel_fn(
            sr=self.sampling_rate,
            n_fft=self.n_fft,
            n_mels=self.num_mels,
            fmin=self.fmin,
            fmax=self.fmax,
        )
        self.mel_basis = mel.float().to(self.device)
        self.hann_window = torch.hann_window(self.win_size).to(self.device)

    def to(self, device, **kwagrs):
        super().to(device=device, **kwagrs)
        self.mel_basis = self.mel_basis.to(device)
        self.hann_window = self.hann_window.to(device)
        self.device = device

    def forward(self, y, center=False, complex=False):
        if isinstance(y, np.ndarray):
            y = torch.FloatTensor(y)
        if len(y.shape) == 1:
            y = y.unsqueeze(0)
        y = y.clamp(min=-1.0, max=1.0).to(self.device)

        pad_length = math.ceil(y.shape[1] / self.hop_size) * self.hop_size - y.shape[1]
        y = torch.nn.functional.pad(
            y.unsqueeze(1),
            [int((self.n_fft - self.hop_size) / 2), int((self.n_fft - self.hop_size) / 2 + pad_length)],
            mode="reflect",
        )
        y = y.squeeze(1)

        spec = torch.stft(
            y,
            self.n_fft,
            hop_length=self.hop_size,
            win_length=self.win_size,
            window=self.hann_window,
            center=center,
            pad_mode="reflect",
            normalized=False,
            onesided=True,
            return_complex=True,
        )
        if not complex:
            spec = torch.view_as_real(spec)
            spec = torch.sqrt(spec.pow(2).sum(-1) + (1e-9))  # [B, n_fft, T]
            spec = torch.matmul(self.mel_basis, spec)
            spec = spectral_normalize_torch(spec)
            spec = spec.transpose(1, 2)  # [B, T, n_fft]
        else:
            B, C, T, _ = spec.shape
            spec = spec.transpose(1, 2)  # [B, T, n_fft, 2]
        return spec


class RosvotModel(nn.Module, SupportAudioInput, SupportsComponentDiscovery):
    support_audio_input: ClassVar[bool] = True
    _dit_modules: ClassVar[list[str]] = []
    _encoder_modules: ClassVar[list[str]] = ["."]
    _vae_modules: ClassVar[list[str]] = []
    _resident_modules: ClassVar[list[str]] = []
    _layerwise_offload_blocks_attrs: ClassVar[list[str]] = ["midi.net", "midi.pitch_decoder", "midi.cond_encoder"]

    def __init__(
        self,
        rosvot_ckpt: str | Path,
        *,
        config_path: str | Path = "",
        pe: nn.Module | None = None,
        the: float = 0.85,
        verbose: bool = False,
        rosvot_source_dir: str | Path | None = None,
    ):
        super().__init__()
        self.verbose = verbose
        ckpt = Path(rosvot_ckpt)
        resolved_config = Path(config_path) if config_path else ckpt.with_name("config.yaml")
        self.hparams = load_rosvot_config(resolved_config, hparams_str=f"note_bd_threshold={the}")
        if verbose:
            logger.info("ROSVOT config: %s", resolved_config)

        self.midi = MidiExtractor(self.hparams)
        self.mel_net = MelNet(self.hparams)
        self.pe = pe if pe is not None and self.hparams.get("use_pitch_embed", False) else None
        self._checkpoint_path = str(ckpt)
        self.load_checkpoint(str(ckpt), verbose=verbose)
        self.eval()

    def load_checkpoint(self, checkpoint_path: str | None = None, *, verbose: bool = False) -> None:
        load_model_ckpt(self.midi, checkpoint_path or self._checkpoint_path, verbose=verbose)

    @torch.no_grad()
    def forward(self, wav: torch.Tensor, word_durs: list[float]) -> dict[str, Any]:
        hparams = self.hparams
        if wav.ndim == 1:
            wav = wav.unsqueeze(0)

        mel_len = (wav.shape[-1] + hparams["hop_size"] - 1) // hparams["hop_size"]
        min_word_dur = hparams.get("min_word_dur", 20) / 1000
        wd_raw = list(word_durs)
        word_durs_filtered: list[float] = []
        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 word_durs_filtered:
                    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"])
        if mel2word.size > 0 and mel2word[0] == 0:
            mel2word = mel2word + 1

        real_len = min(mel_len, int(np.sum(mel2word > 0)))
        T = math.ceil(min(real_len, hparams["max_frames"]) / hparams["frames_multiple"]) * hparams["frames_multiple"]
        device = wav.device
        self.mel_net.to(device)
        wav_t = pad_or_cut_xd(wav.float(), T * hparams["hop_size"], 1)

        pitch_coarse = uv_t = None
        if self.pe is not None:
            f0s, uvs = self.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.as_tensor(f0_1d, device=device, dtype=torch.float32), T, 0).unsqueeze(0)
            uv_t = pad_or_cut_xd(torch.as_tensor(uv_1d, device=device, dtype=torch.float32), T, 0).long().unsqueeze(0)
            pitch_coarse = f0_to_coarse(denorm_f0(f0_t, uv_t)).to(device)

        mel = pad_or_cut_xd(self.mel_net(wav_t)[0], T, dim=0).unsqueeze(0)
        mel_nonpadding_mask = torch.zeros(1, T, device=device)
        mel_nonpadding_mask[:, :real_len] = 1.0
        mel = (mel.transpose(1, 2) * mel_nonpadding_mask.unsqueeze(1)).transpose(1, 2)
        mel_nonpadding = mel.abs().sum(-1) > 0

        mel2word_t = pad_or_cut_xd(torch.as_tensor(mel2word, device=device, dtype=torch.long), T, 0)
        word_bd = torch.zeros_like(mel2word_t)
        word_bd[1:] = (mel2word_t[1:] != mel2word_t[:-1]).long()
        word_bd[real_len:] = 0
        word_bd = word_bd.unsqueeze(0)

        outputs = self.midi(
            mel=mel[:, :, : hparams.get("use_mel_bins", 80)],
            word_bd=word_bd,
            pitch=pitch_coarse,
            uv=uv_t,
            non_padding=mel_nonpadding,
        )
        outputs["word_durs_filtered"] = word_durs_filtered
        outputs["real_len"] = real_len
        outputs["word_bd"] = word_bd
        return outputs

    @staticmethod
    def _load_wav(wav_src: str | np.ndarray, sample_rate: int, *, src_sample_rate: int | None = None) -> np.ndarray:
        if isinstance(wav_src, str):
            wav, _ = load_mono_audio(wav_src, target_sr=sample_rate)
            return wav
        wav = np.asarray(wav_src, dtype=np.float32)
        if src_sample_rate is not None and src_sample_rate != sample_rate:
            wav = resample_mono(wav, orig_sr=src_sample_rate, target_sr=sample_rate)
        return wav

    @staticmethod
    def _normalize_note2words(note2words: list[int]) -> list[int]:
        if not note2words:
            return []
        out = [note2words[0]]
        for idx in range(1, len(note2words)):
            out.append(max(note2words[idx], out[-1]))
        return out

    @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):
            ep_types.append(1 if w == "<SP>" else (2 if i != prev else 3))
            prev = i
        return ep_types

    @torch.no_grad()
    def transcribe(
        self,
        item: dict[str, Any],
        *,
        segment_info: dict[str, Any] | None = None,
        verbose: bool = False,
    ) -> dict[str, Any]:
        if "word_durs" not in item:
            raise ValueError('item must contain "word_durs" from lyric transcription')

        if item.get("wav_fn"):
            wav = self._load_wav(item["wav_fn"], self.hparams["audio_sample_rate"])
        elif item.get("wav") is not None:
            wav = self._load_wav(
                item["wav"],
                self.hparams["audio_sample_rate"],
                src_sample_rate=item.get("sample_rate"),
            )
        else:
            raise ValueError('item must contain "wav_fn" or "wav"')

        device = next(self.parameters()).device
        outputs = self(torch.from_numpy(wav).float().to(device), list(item["word_durs"]))
        real_len = int(outputs["real_len"])
        word_durs_filtered = outputs["word_durs_filtered"]
        word_bd = outputs["word_bd"]
        item_name = item.get("item_name", "")

        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]]

        if note_pred.shape == (0,):
            rosvot_out = {"item_name": item_name, "pitches": [], "note_durs": [], "note2words": None}
        else:
            note_itv_pred = boundary2Interval(note_bd_pred)
            word_bd_for_reg = word_bd[0].detach().cpu().numpy()[:real_len]
            hop = self.hparams["hop_size"]
            sr = self.hparams["audio_sample_rate"]

            if self.hparams.get("infer_regulate_real_note_itv", True):
                try:
                    note_itv_pred_secs, note2words = regulate_real_note_itv(
                        note_itv_pred,
                        note_bd_pred,
                        word_bd_for_reg,
                        np.array(word_durs_filtered),
                        hop,
                        sr,
                    )
                    note_pred, note_itv_pred_secs, note2words = regulate_ill_slur(
                        note_pred, note_itv_pred_secs, note2words
                    )
                except Exception:
                    if verbose:
                        logger.exception("ROSVOT postprocess failed")
                    note_itv_pred_secs = note_itv_pred * hop / sr
                    note2words = None
            else:
                note_itv_pred_secs = note_itv_pred * hop / sr
                note2words = None

            rosvot_out = {
                "item_name": item_name,
                "pitches": note_pred.tolist(),
                "note_durs": [float(itv[1] - itv[0]) for itv in note_itv_pred_secs],
                "note2words": note2words.tolist() if note2words is not None else None,
            }

        note2words_raw = rosvot_out.get("note2words") or []
        align_words = [item["words"][idx - 1] for idx in note2words_raw if 0 < idx <= len(item["words"])]
        ep_types = self._build_ep_types(self._normalize_note2words(note2words_raw), align_words) if align_words else []
        seg = segment_info or item

        return {
            "item_name": seg.get("item_name", item_name),
            "wav_fn": seg.get("wav_fn", item.get("wav_fn", "")),
            "origin_wav_fn": seg.get("origin_wav_fn", item.get("origin_wav_fn", "")),
            "start_time_ms": seg.get("start_time_ms", ""),
            "end_time_ms": seg.get("end_time_ms", ""),
            "language": item.get("language", ""),
            "note_text": align_words,
            "note_dur": rosvot_out.get("note_durs", []),
            "note_type": ep_types,
            "note_pitch": rosvot_out.get("pitches", []),
        }