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

"""Rule-based vocal segmentation from F0 contour (no neural net)."""

from __future__ import annotations

import time
from dataclasses import dataclass
from typing import ClassVar

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

from vllm_omni.diffusion.models.interface import SupportAudioInput, SupportsComponentDiscovery

logger = init_logger(__name__)


@dataclass(frozen=True)
class VocalSegmentConfig:
    hop_ms: int = 20
    smooth_ms: int = 200
    start_ms: int = 120
    end_ms: int = 200
    prepad_ms: int = 80
    postpad_ms: int = 120
    min_len_ms: int = 1000
    max_len_ms: int = 20000
    short_seg_merge_gap_ms: int = 8000
    small_gap_ms: int = 500
    lookback_ms: int = 200
    lookahead_ms: int = 100


def _moving_average(x: np.ndarray, win: int) -> np.ndarray:
    if win <= 1:
        return x
    kernel = np.ones(win, dtype=np.float32) / float(win)
    return np.convolve(x, kernel, mode="same")


def _merge_short_segments(
    segments_ms: list[list[int]],
    *,
    min_len_ms: int,
    max_len_ms: int,
    short_seg_merge_gap_ms: int,
    small_gap_ms: int,
) -> list[list[int]]:
    if not segments_ms:
        return []

    merged: list[list[int]] = []
    cur_start, cur_end = segments_ms[0]
    for next_start, next_end in segments_ms[1:]:
        cur_len = cur_end - cur_start
        gap_ms = next_start - cur_end
        merged_len = next_end - cur_start
        should_merge = (cur_len < min_len_ms and gap_ms < short_seg_merge_gap_ms) or (
            gap_ms < small_gap_ms and merged_len < max_len_ms
        )
        if should_merge:
            cur_end = next_end
            continue
        if (cur_end - cur_start) >= min_len_ms:
            merged.append([cur_start, cur_end])
        cur_start, cur_end = next_start, next_end

    if (cur_end - cur_start) >= min_len_ms:
        merged.append([cur_start, cur_end])

    if not merged:
        return segments_ms

    return merged


def _voiced_to_segments(
    voiced: np.ndarray,
    *,
    hop_ms: int,
    smooth_ms: int,
    start_ms: int,
    end_ms: int,
    prepad_ms: int,
    postpad_ms: int,
    max_len_ms: int,
) -> list[list[int]]:
    smooth_frames = max(1, int(round(smooth_ms / hop_ms)))
    active = _moving_average(voiced.astype(np.float32), smooth_frames) >= 0.5

    segments: list[list[int]] = []
    start_idx = None
    start_frames = max(1, int(round(start_ms / hop_ms)))
    end_frames = max(1, int(round(end_ms / hop_ms)))
    prepad_frames = max(0, int(round(prepad_ms / hop_ms)))
    postpad_frames = max(0, int(round(postpad_ms / hop_ms)))
    active_count = 0
    inactive_count = 0

    for i, flag in enumerate(active):
        if flag:
            active_count += 1
            inactive_count = 0
        else:
            inactive_count += 1
            active_count = 0

        if start_idx is None:
            if active_count >= start_frames:
                start_idx = max(0, i - start_frames + 1 - prepad_frames)
        else:
            if inactive_count >= end_frames:
                end_idx = min(len(active) - 1, i - end_frames + 1 + postpad_frames)
                start_ms_val = start_idx * hop_ms
                end_ms_val = end_idx * hop_ms + hop_ms
                if end_ms_val > start_ms_val:
                    segments.append([int(start_ms_val), int(end_ms_val)])
                start_idx = None

    if start_idx is not None:
        start_ms_val = start_idx * hop_ms
        end_idx = min(len(active) - 1, len(active) - 1 + postpad_frames)
        end_ms_val = end_idx * hop_ms + hop_ms
        if end_ms_val > start_ms_val:
            segments.append([int(start_ms_val), int(end_ms_val)])

    def _split_segment(seg: list[int]) -> list[list[int]]:
        start_ms_val, end_ms_val = seg
        start_frame = int(start_ms_val // hop_ms)
        end_frame = int((end_ms_val - 1) // hop_ms)
        end_frame = max(start_frame, min(end_frame, len(active) - 1))

        best_start = None
        best_len = 0
        cur_start = None
        cur_len = 0
        for idx in range(start_frame, end_frame + 1):
            if not active[idx]:
                if cur_start is None:
                    cur_start = idx
                    cur_len = 1
                else:
                    cur_len += 1
            else:
                if cur_start is not None and cur_len > best_len:
                    best_start, best_len = cur_start, cur_len
                cur_start = None
                cur_len = 0
        if cur_start is not None and cur_len > best_len:
            best_start, best_len = cur_start, cur_len

        if best_start is None:
            split_frame = (start_frame + end_frame) // 2
        else:
            split_frame = best_start + best_len // 2

        split_ms = split_frame * hop_ms
        if split_ms <= start_ms_val:
            split_ms = start_ms_val + hop_ms
        if split_ms >= end_ms_val:
            split_ms = end_ms_val - hop_ms
        if split_ms <= start_ms_val or split_ms >= end_ms_val:
            return [seg]
        return [[start_ms_val, int(split_ms)], [int(split_ms), end_ms_val]]

    queue = segments[:]
    segments = []
    while queue:
        seg = queue.pop(0)
        if (seg[1] - seg[0]) <= max_len_ms:
            segments.append(seg)
            continue
        parts = _split_segment(seg)
        if len(parts) == 1:
            segments.append(seg)
        else:
            queue = parts + queue
    return segments


class VocalSegmenter(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]] = []

    def __init__(self, config: VocalSegmentConfig | None = None, *, verbose: bool = False):
        super().__init__()
        self.config = config or VocalSegmentConfig()
        self.verbose = verbose

    def forward(
        self,
        audio: np.ndarray,
        sample_rate: int,
        f0: np.ndarray,
        *,
        base_name: str = "vocal",
        origin_wav_fn: str = "",
        verbose: bool | None = None,
    ) -> list[dict]:
        verbose = self.verbose if verbose is None else verbose
        if verbose:
            logger.info("[vocal detection] segmenting %s", base_name)
            t0 = time.time()

        cfg = self.config
        y = np.asarray(audio, dtype=np.float32).squeeze()
        sr = int(sample_rate)

        segments_ms = _voiced_to_segments(
            f0 > 0,
            hop_ms=cfg.hop_ms,
            smooth_ms=cfg.smooth_ms,
            start_ms=cfg.start_ms,
            end_ms=cfg.end_ms,
            prepad_ms=cfg.prepad_ms,
            postpad_ms=cfg.postpad_ms,
            max_len_ms=cfg.max_len_ms,
        )
        segments_ms = _merge_short_segments(
            segments_ms,
            min_len_ms=cfg.min_len_ms,
            max_len_ms=cfg.max_len_ms,
            short_seg_merge_gap_ms=cfg.short_seg_merge_gap_ms,
            small_gap_ms=cfg.small_gap_ms,
        )

        adjusted_segments: list[list[int]] = []
        prev_end = 0
        for start_ms, end_ms in segments_ms:
            start_ms = max(0, start_ms - cfg.lookback_ms)
            end_ms = min(end_ms + cfg.lookahead_ms, int(y.shape[0] / sr * 1000))
            if start_ms < prev_end and adjusted_segments:
                adjusted_segments[-1][1] = start_ms
            adjusted_segments.append([start_ms, end_ms])
            prev_end = end_ms

        segment_infos = []
        for idx, (start_ms, end_ms) in enumerate(adjusted_segments):
            if end_ms - start_ms > cfg.max_len_ms:
                start_ms = end_ms - cfg.max_len_ms
            key = f"{base_name}_{idx}"
            start_sample = int(start_ms * sr / 1000)
            end_sample = int(end_ms * sr / 1000)
            segment_infos.append(
                {
                    "item_name": key,
                    "wav": y[start_sample:end_sample],
                    "sample_rate": sr,
                    "wav_fn": key,
                    "start_time_ms": int(start_sample * 1000 / sr),
                    "end_time_ms": int(end_sample * 1000 / sr),
                    "origin_wav_fn": origin_wav_fn,
                    "duration": int((end_sample - start_sample) * 1000 / sr),
                }
            )

        if verbose:
            logger.info(
                "[vocal detection] done n_segments=%s time=%.3fs",
                len(segment_infos),
                time.time() - t0,
            )
        return segment_infos