|
|
| """
|
| F0 extraction helpers used by the local VC pipeline.
|
| """
|
|
|
| from __future__ import annotations
|
|
|
| from typing import Literal
|
|
|
| import numpy as np
|
| import torch
|
|
|
| from infer.quality_policy import build_conservative_crepe_fill_mask
|
|
|
|
|
| F0Method = Literal["rmvpe", "pm", "harvest", "crepe", "hybrid"]
|
|
|
|
|
| class F0Extractor:
|
| """Base F0 extractor."""
|
|
|
| def __init__(self, sample_rate: int = 16000, hop_length: int = 160):
|
| self.sample_rate = sample_rate
|
| self.hop_length = hop_length
|
| self.f0_min = 50
|
| self.f0_max = 1100
|
|
|
| def extract(self, audio: np.ndarray) -> np.ndarray:
|
| raise NotImplementedError
|
|
|
|
|
| class PMExtractor(F0Extractor):
|
| """Praat/Parselmouth extractor."""
|
|
|
| def extract(self, audio: np.ndarray) -> np.ndarray:
|
| import parselmouth
|
|
|
| time_step = self.hop_length / self.sample_rate
|
| sound = parselmouth.Sound(audio, self.sample_rate)
|
| pitch = sound.to_pitch_ac(
|
| time_step=time_step,
|
| voicing_threshold=0.6,
|
| pitch_floor=self.f0_min,
|
| pitch_ceiling=self.f0_max,
|
| )
|
| f0 = pitch.selected_array["frequency"]
|
| f0[f0 == 0] = np.nan
|
| return f0
|
|
|
|
|
| class HarvestExtractor(F0Extractor):
|
| """PyWorld harvest extractor."""
|
|
|
| def extract(self, audio: np.ndarray) -> np.ndarray:
|
| import pyworld
|
|
|
| audio = audio.astype(np.float64)
|
| f0, _ = pyworld.harvest(
|
| audio,
|
| self.sample_rate,
|
| f0_floor=self.f0_min,
|
| f0_ceil=self.f0_max,
|
| frame_period=self.hop_length / self.sample_rate * 1000,
|
| )
|
| return f0
|
|
|
|
|
| class CrepeExtractor(F0Extractor):
|
| """TorchCrepe extractor."""
|
|
|
| def __init__(
|
| self,
|
| sample_rate: int = 16000,
|
| hop_length: int = 160,
|
| device: str = "cuda",
|
| ):
|
| super().__init__(sample_rate, hop_length)
|
| self.device = device
|
|
|
| def extract(self, audio: np.ndarray) -> np.ndarray:
|
| import torchcrepe
|
|
|
| audio_tensor = torch.from_numpy(audio).float().unsqueeze(0).to(self.device)
|
| f0, _ = torchcrepe.predict(
|
| audio_tensor,
|
| self.sample_rate,
|
| self.hop_length,
|
| self.f0_min,
|
| self.f0_max,
|
| model="full",
|
| batch_size=512,
|
| device=self.device,
|
| return_periodicity=True,
|
| )
|
| return f0.squeeze(0).cpu().numpy()
|
|
|
|
|
| class RMVPEExtractor(F0Extractor):
|
| """RMVPE extractor."""
|
|
|
| def __init__(
|
| self,
|
| model_path: str,
|
| sample_rate: int = 16000,
|
| hop_length: int = 160,
|
| device: str = "cuda",
|
| ):
|
| super().__init__(sample_rate, hop_length)
|
| self.device = device
|
| self.model = None
|
| self.model_path = model_path
|
|
|
| def load_model(self) -> None:
|
| if self.model is not None:
|
| return
|
|
|
| from models.rmvpe import RMVPE
|
|
|
| self.model = RMVPE(self.model_path, device=self.device)
|
| print(f"RMVPE model loaded: {self.device}")
|
|
|
| def extract(self, audio: np.ndarray) -> np.ndarray:
|
| self.load_model()
|
| return self.model.infer_from_audio(audio, thred=0.01)
|
|
|
|
|
| def get_f0_extractor(
|
| method: F0Method,
|
| device: str = "cuda",
|
| rmvpe_path: str | None = None,
|
| crepe_threshold: float = 0.05,
|
| ) -> F0Extractor:
|
| """Create an F0 extractor."""
|
|
|
| if method == "rmvpe":
|
| if rmvpe_path is None:
|
| raise ValueError("RMVPE requires a model path")
|
| return RMVPEExtractor(rmvpe_path, device=device)
|
| if method == "hybrid":
|
| if rmvpe_path is None:
|
| raise ValueError("Hybrid requires an RMVPE model path")
|
| return HybridF0Extractor(
|
| rmvpe_path,
|
| device=device,
|
| crepe_threshold=crepe_threshold,
|
| )
|
| if method == "pm":
|
| return PMExtractor()
|
| if method == "harvest":
|
| return HarvestExtractor()
|
| if method == "crepe":
|
| return CrepeExtractor(device=device)
|
| raise ValueError(f"Unknown F0 method: {method}")
|
|
|
|
|
| class HybridF0Extractor(F0Extractor):
|
| """
|
| Conservative hybrid extractor.
|
|
|
| RMVPE remains the primary estimator. CREPE is only allowed to repair short,
|
| high-confidence dropouts that sit inside already-voiced context.
|
| """
|
|
|
| def __init__(
|
| self,
|
| rmvpe_path: str,
|
| sample_rate: int = 16000,
|
| hop_length: int = 160,
|
| device: str = "cuda",
|
| crepe_threshold: float = 0.05,
|
| max_fill_ratio: float = 0.02,
|
| max_fill_frames: int = 320,
|
| context_radius: int = 6,
|
| ):
|
| super().__init__(sample_rate, hop_length)
|
| self.device = device
|
| self.rmvpe = RMVPEExtractor(rmvpe_path, sample_rate, hop_length, device)
|
| self.crepe = None
|
| self.crepe_threshold = float(crepe_threshold)
|
| self.max_fill_ratio = float(max_fill_ratio)
|
| self.max_fill_frames = int(max_fill_frames)
|
| self.context_radius = int(context_radius)
|
|
|
| def _load_crepe(self) -> None:
|
| if self.crepe is not None:
|
| return
|
|
|
| try:
|
| self.crepe = CrepeExtractor(
|
| self.sample_rate,
|
| self.hop_length,
|
| self.device,
|
| )
|
| except ImportError:
|
| print("Warning: torchcrepe is unavailable, hybrid falls back to RMVPE")
|
| self.crepe = False
|
|
|
| def extract(self, audio: np.ndarray) -> np.ndarray:
|
| f0_rmvpe = self.rmvpe.extract(audio)
|
|
|
| self._load_crepe()
|
| if self.crepe is False:
|
| return f0_rmvpe
|
|
|
| import torchcrepe
|
|
|
| audio_tensor = torch.from_numpy(audio).float().unsqueeze(0).to(self.device)
|
| f0_crepe, confidence = torchcrepe.predict(
|
| audio_tensor,
|
| self.sample_rate,
|
| self.hop_length,
|
| self.f0_min,
|
| self.f0_max,
|
| model="full",
|
| batch_size=512,
|
| device=self.device,
|
| return_periodicity=True,
|
| )
|
| f0_crepe = f0_crepe.squeeze(0).cpu().numpy()
|
| confidence = confidence.squeeze(0).cpu().numpy()
|
|
|
| min_len = min(len(f0_rmvpe), len(f0_crepe), len(confidence))
|
| if min_len <= 0:
|
| return f0_rmvpe
|
|
|
| f0_rmvpe = np.asarray(f0_rmvpe[:min_len], dtype=np.float32)
|
| f0_crepe = np.asarray(f0_crepe[:min_len], dtype=np.float32)
|
| confidence = np.asarray(confidence[:min_len], dtype=np.float32)
|
|
|
| fill_mask = build_conservative_crepe_fill_mask(
|
| f0_rmvpe=f0_rmvpe,
|
| f0_crepe=f0_crepe,
|
| confidence=confidence,
|
| confidence_threshold=self.crepe_threshold,
|
| max_ratio=self.max_fill_ratio,
|
| max_frames=self.max_fill_frames,
|
| context_radius=self.context_radius,
|
| )
|
|
|
| if not np.any(fill_mask):
|
| return f0_rmvpe
|
|
|
| f0_hybrid = f0_rmvpe.copy()
|
| f0_hybrid[fill_mask] = f0_crepe[fill_mask]
|
| return f0_hybrid
|
|
|
|
|
| def shift_f0(f0: np.ndarray, semitones: float) -> np.ndarray:
|
| """Shift F0 by semitones."""
|
|
|
| factor = 2 ** (semitones / 12)
|
| return f0 * factor
|
|
|
