mirror of
https://github.com/RVC-Project/Retrieval-based-Voice-Conversion-WebUI.git
synced 2026-08-29 01:59:23 +02:00
346 lines
14 KiB
Python
346 lines
14 KiB
Python
"""RVC inference worker for the RVC Realtime VST2/VST3 plugin."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import ctypes
|
|
import json
|
|
import mmap
|
|
import os
|
|
import struct
|
|
import sys
|
|
import time
|
|
import traceback
|
|
from pathlib import Path
|
|
|
|
MAGIC = 0x50564352
|
|
PROTOCOL_VERSION = 1
|
|
HEADER_BYTES = 4096
|
|
MAX_FRAMES = 131072
|
|
MAP_BYTES = HEADER_BYTES + MAX_FRAMES * 4 * 2
|
|
INPUT_OFFSET = HEADER_BYTES
|
|
OUTPUT_OFFSET = HEADER_BYTES + MAX_FRAMES * 4
|
|
STATUS_TEXT_OFFSET = 128
|
|
STATUS_TEXT_BYTES = 512
|
|
|
|
STATUS_STARTING = 1
|
|
STATUS_LOADING = 2
|
|
STATUS_READY = 3
|
|
STATUS_ERROR = -1
|
|
STATUS_STOP = -2
|
|
|
|
SYNCHRONIZE = 0x00100000
|
|
EVENT_MODIFY_STATE = 0x0002
|
|
WAIT_OBJECT_0 = 0
|
|
WAIT_TIMEOUT = 258
|
|
|
|
|
|
def write_value(shared: mmap.mmap, offset: int, fmt: str, value) -> None:
|
|
struct.pack_into("<" + fmt, shared, offset, value)
|
|
|
|
|
|
def read_value(shared: mmap.mmap, offset: int, fmt: str):
|
|
return struct.unpack_from("<" + fmt, shared, offset)[0]
|
|
|
|
|
|
def write_status(shared: mmap.mmap, state: int, message: str) -> None:
|
|
write_value(shared, 8, "i", state)
|
|
encoded = message.encode("utf-8", errors="replace")[: STATUS_TEXT_BYTES - 1]
|
|
shared[STATUS_TEXT_OFFSET : STATUS_TEXT_OFFSET + STATUS_TEXT_BYTES] = b"\0" * STATUS_TEXT_BYTES
|
|
shared[STATUS_TEXT_OFFSET : STATUS_TEXT_OFFSET + len(encoded)] = encoded
|
|
shared.flush()
|
|
|
|
|
|
class WinEvent:
|
|
def __init__(self, name: str):
|
|
kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
|
kernel32.OpenEventW.argtypes = [ctypes.c_uint32, ctypes.c_int, ctypes.c_wchar_p]
|
|
kernel32.OpenEventW.restype = ctypes.c_void_p
|
|
self._wait = kernel32.WaitForSingleObject
|
|
self._wait.argtypes = [ctypes.c_void_p, ctypes.c_uint32]
|
|
self._wait.restype = ctypes.c_uint32
|
|
self._set = kernel32.SetEvent
|
|
self._set.argtypes = [ctypes.c_void_p]
|
|
self._set.restype = ctypes.c_int
|
|
self._close = kernel32.CloseHandle
|
|
self._close.argtypes = [ctypes.c_void_p]
|
|
self._close.restype = ctypes.c_int
|
|
self.handle = kernel32.OpenEventW(SYNCHRONIZE | EVENT_MODIFY_STATE, False, name)
|
|
if not self.handle:
|
|
raise OSError(ctypes.get_last_error(), f"OpenEventW failed: {name}")
|
|
|
|
def wait(self, timeout_ms: int) -> int:
|
|
return int(self._wait(self.handle, timeout_ms))
|
|
|
|
def set(self) -> None:
|
|
if not self._set(self.handle):
|
|
raise OSError(ctypes.get_last_error(), "SetEvent failed")
|
|
|
|
def close(self) -> None:
|
|
if self.handle:
|
|
self._close(self.handle)
|
|
self.handle = None
|
|
|
|
|
|
class RVCStreamEngine:
|
|
def __init__(self, cfg: dict):
|
|
self.root = Path(cfg["rvc_root"]).resolve()
|
|
os.chdir(self.root)
|
|
sys.path.insert(0, str(self.root))
|
|
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
|
|
os.environ.setdefault("OMP_NUM_THREADS", "4")
|
|
|
|
import librosa
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
import torchaudio.transforms as tat
|
|
|
|
from configs.config import Config
|
|
from infer import rtrvc
|
|
from tools.cuda_graph import run_cuda_graph
|
|
|
|
self.librosa = librosa
|
|
self.np = np
|
|
self.torch = torch
|
|
self.F = F
|
|
self.tat = tat
|
|
self.run_cuda_graph = run_cuda_graph
|
|
|
|
self.sample_rate = int(cfg["sample_rate"])
|
|
self.block_ms = float(cfg["block_ms"])
|
|
self.crossfade_ms = float(cfg["crossfade_ms"])
|
|
self.extra_ms = float(cfg["extra_ms"])
|
|
self.zc = max(1, self.sample_rate // 100)
|
|
self.block_frame = int(round(self.block_ms / 1000 * self.sample_rate / self.zc) * self.zc)
|
|
self.block_frame_16k = 160 * self.block_frame // self.zc
|
|
self.crossfade_frame = int(round(self.crossfade_ms / 1000 * self.sample_rate / self.zc) * self.zc)
|
|
# Match the source realtime GUI: SOLA overlap is capped at 40 ms and
|
|
# remains independent of the audio callback block length.
|
|
self.sola_buffer_frame = min(self.crossfade_frame, 4 * self.zc)
|
|
self.effective_crossfade_ms = 1000.0 * self.sola_buffer_frame / self.sample_rate
|
|
self.sola_search_frame = self.zc
|
|
self.extra_frame = int(round(self.extra_ms / 1000 * self.sample_rate / self.zc) * self.zc)
|
|
|
|
self.config = Config()
|
|
self.rvc = rtrvc.RVC(
|
|
0.0,
|
|
0.0,
|
|
str(Path(cfg["model_path"]).resolve()),
|
|
str(Path(cfg["index_path"]).resolve()) if cfg.get("index_path") else "",
|
|
0.0,
|
|
self.config,
|
|
None,
|
|
)
|
|
|
|
total_frames = self.extra_frame + self.crossfade_frame + self.sola_search_frame + self.block_frame
|
|
self.input_wav = torch.zeros(total_frames, device=self.config.device, dtype=torch.float32)
|
|
self.input_wav_res = torch.zeros(160 * total_frames // self.zc, device=self.config.device, dtype=torch.float32)
|
|
self.rms_buffer = np.zeros(4 * self.zc, dtype="float32")
|
|
self.sola_buffer = torch.zeros(self.sola_buffer_frame, device=self.config.device, dtype=torch.float32)
|
|
self.sola_den_kernel = torch.ones(1, 1, self.sola_buffer_frame, device=self.config.device, dtype=torch.float32)
|
|
self.skip_head = self.extra_frame // self.zc
|
|
self.return_length = (self.block_frame + self.sola_buffer_frame + self.sola_search_frame) // self.zc
|
|
self.fade_in_window = torch.sin(
|
|
0.5 * np.pi * torch.linspace(0.0, 1.0, steps=self.sola_buffer_frame, device=self.config.device)
|
|
) ** 2
|
|
self.fade_out_window = 1 - self.fade_in_window
|
|
self.resampler = tat.Resample(orig_freq=self.sample_rate, new_freq=16000, dtype=torch.float32).to(self.config.device)
|
|
self.resampler2 = None
|
|
if self.rvc.tgt_sr != self.sample_rate:
|
|
self.resampler2 = tat.Resample(orig_freq=self.rvc.tgt_sr, new_freq=self.sample_rate, dtype=torch.float32).to(self.config.device)
|
|
self.last_pitch = None
|
|
self.last_formant = None
|
|
self.last_index_rate = None
|
|
|
|
def prewarm(self) -> None:
|
|
phase = self.torch.arange(self.block_frame, device=self.config.device, dtype=self.torch.float32)
|
|
probe = 0.05 * self.torch.sin(2 * self.np.pi * 220.0 * phase / self.sample_rate)
|
|
self.process(probe.cpu().numpy(), 12.0, 0.0, 0.0, 0.5, -60.0, 0)
|
|
if self.torch.device(self.config.device).type == "cuda":
|
|
self.torch.cuda.synchronize(self.config.device)
|
|
|
|
self.input_wav.zero_()
|
|
self.input_wav_res.zero_()
|
|
self.rms_buffer.fill(0.0)
|
|
self.sola_buffer.zero_()
|
|
if hasattr(self.rvc, "cache_pitch"):
|
|
self.rvc.cache_pitch.zero_()
|
|
if hasattr(self.rvc, "cache_pitchf"):
|
|
self.rvc.cache_pitchf.zero_()
|
|
|
|
def process(self, audio, pitch: float, formant: float, index_rate: float,
|
|
rms_mix: float, threshold: float, f0_method: int):
|
|
np = self.np
|
|
torch = self.torch
|
|
F = self.F
|
|
if len(audio) != self.block_frame:
|
|
raise ValueError(f"block mismatch: worker={self.block_frame}, request={len(audio)}")
|
|
|
|
if self.last_pitch != pitch:
|
|
self.rvc.change_key(pitch)
|
|
self.last_pitch = pitch
|
|
if self.last_formant != formant:
|
|
self.rvc.change_formant(formant)
|
|
self.last_formant = formant
|
|
if self.last_index_rate != index_rate:
|
|
self.rvc.change_index_rate(index_rate)
|
|
self.last_index_rate = index_rate
|
|
|
|
indata = np.asarray(audio, dtype=np.float32).copy()
|
|
if threshold > -60.0:
|
|
gated = np.append(self.rms_buffer, indata)
|
|
rms = self.librosa.feature.rms(y=gated, frame_length=4 * self.zc, hop_length=self.zc)[:, 2:]
|
|
self.rms_buffer[:] = gated[-4 * self.zc :]
|
|
gated = gated[2 * self.zc - self.zc // 2 :]
|
|
below = self.librosa.amplitude_to_db(rms, ref=1.0)[0] < threshold
|
|
for i, mute in enumerate(below):
|
|
if mute:
|
|
gated[i * self.zc : (i + 1) * self.zc] = 0
|
|
indata = gated[self.zc // 2 :]
|
|
|
|
self.input_wav[:-self.block_frame] = self.input_wav[self.block_frame:].clone()
|
|
self.input_wav[-self.block_frame:] = torch.from_numpy(indata).to(self.config.device)
|
|
self.input_wav_res[:-self.block_frame_16k] = self.input_wav_res[self.block_frame_16k:].clone()
|
|
resample_input = self.input_wav[-self.block_frame - 2 * self.zc :]
|
|
resampled = self.run_cuda_graph(
|
|
self.resampler,
|
|
"vst-input-resample",
|
|
lambda value: self.resampler(value),
|
|
resample_input,
|
|
)[160:]
|
|
self.input_wav_res[-self.block_frame_16k:] = resampled[-self.block_frame_16k:]
|
|
|
|
method = ("rmvpe", "fcpe", "pm")[max(0, min(2, int(f0_method)))]
|
|
infer_wav = self.rvc.infer(
|
|
self.input_wav_res,
|
|
self.block_frame_16k,
|
|
self.skip_head,
|
|
self.return_length,
|
|
method,
|
|
)
|
|
if self.resampler2 is not None:
|
|
infer_wav = self.run_cuda_graph(
|
|
self.resampler2,
|
|
"vst-output-resample",
|
|
lambda value: self.resampler2(value),
|
|
infer_wav,
|
|
)
|
|
|
|
if rms_mix < 1.0:
|
|
input_tail = self.input_wav[self.extra_frame :]
|
|
rms1 = self.librosa.feature.rms(
|
|
y=input_tail[: infer_wav.shape[0]].cpu().numpy(), frame_length=4 * self.zc, hop_length=self.zc
|
|
)
|
|
rms1 = torch.from_numpy(rms1).to(self.config.device)
|
|
rms1 = F.interpolate(rms1.unsqueeze(0), size=infer_wav.shape[0] + 1, mode="linear", align_corners=True)[0, 0, :-1]
|
|
rms2 = self.librosa.feature.rms(
|
|
y=infer_wav.cpu().numpy(), frame_length=4 * self.zc, hop_length=self.zc
|
|
)
|
|
rms2 = torch.from_numpy(rms2).to(self.config.device)
|
|
rms2 = F.interpolate(rms2.unsqueeze(0), size=infer_wav.shape[0] + 1, mode="linear", align_corners=True)[0, 0, :-1]
|
|
rms2 = torch.maximum(rms2, torch.full_like(rms2, 1e-3))
|
|
infer_wav *= torch.pow(rms1 / rms2, 1.0 - rms_mix)
|
|
|
|
conv_input = infer_wav[None, None, : self.sola_buffer_frame + self.sola_search_frame]
|
|
cor_nom = F.conv1d(conv_input, self.sola_buffer[None, None, :])
|
|
cor_den = torch.sqrt(F.conv1d(conv_input**2, self.sola_den_kernel) + 1e-8)
|
|
sola_offset = int(torch.argmax(cor_nom[0, 0] / cor_den[0, 0]))
|
|
infer_wav = infer_wav[sola_offset:]
|
|
infer_wav[: self.sola_buffer_frame] *= self.fade_in_window
|
|
infer_wav[: self.sola_buffer_frame] += self.sola_buffer * self.fade_out_window
|
|
self.sola_buffer[:] = infer_wav[self.block_frame : self.block_frame + self.sola_buffer_frame]
|
|
return infer_wav[: self.block_frame].float().cpu().numpy()
|
|
|
|
|
|
def run(args: argparse.Namespace) -> int:
|
|
shared = None
|
|
request_event = None
|
|
response_event = None
|
|
try:
|
|
shared = mmap.mmap(-1, MAP_BYTES, tagname=args.map, access=mmap.ACCESS_WRITE)
|
|
request_event = WinEvent(args.request)
|
|
response_event = WinEvent(args.response)
|
|
if read_value(shared, 0, "I") != MAGIC or read_value(shared, 4, "I") != PROTOCOL_VERSION:
|
|
raise RuntimeError("RVC VST protocol mismatch")
|
|
write_status(shared, STATUS_LOADING, "Loading Python and CUDA")
|
|
with open(args.config, "r", encoding="utf-8") as handle:
|
|
cfg = json.load(handle)
|
|
# RVC's Config parses process-wide CLI flags intended for its WebUI.
|
|
sys.argv = [sys.argv[0]]
|
|
engine = RVCStreamEngine(cfg)
|
|
write_status(shared, STATUS_LOADING, "Prewarming CUDA and F0")
|
|
engine.prewarm()
|
|
write_status(shared, STATUS_READY, f"Ready (actual CF {engine.effective_crossfade_ms:.0f} ms)")
|
|
|
|
import numpy as np
|
|
|
|
input_view = np.ndarray((MAX_FRAMES,), dtype=np.float32, buffer=shared, offset=INPUT_OFFSET)
|
|
output_view = np.ndarray((MAX_FRAMES,), dtype=np.float32, buffer=shared, offset=OUTPUT_OFFSET)
|
|
last_sequence = 0
|
|
while True:
|
|
result = request_event.wait(1000)
|
|
if result == WAIT_TIMEOUT:
|
|
continue
|
|
if result != WAIT_OBJECT_0:
|
|
raise OSError("WaitForSingleObject failed")
|
|
if read_value(shared, 8, "i") == STATUS_STOP:
|
|
break
|
|
sequence = read_value(shared, 12, "I")
|
|
if sequence == last_sequence:
|
|
continue
|
|
frames = read_value(shared, 20, "I")
|
|
if frames <= 0 or frames > MAX_FRAMES:
|
|
raise ValueError(f"invalid frame count: {frames}")
|
|
started = time.perf_counter()
|
|
processed = engine.process(
|
|
input_view[:frames],
|
|
read_value(shared, 32, "f"),
|
|
read_value(shared, 36, "f"),
|
|
read_value(shared, 40, "f"),
|
|
read_value(shared, 44, "f"),
|
|
read_value(shared, 48, "f"),
|
|
read_value(shared, 64, "I"),
|
|
)
|
|
output_view[:frames] = processed
|
|
write_value(shared, 56, "f", (time.perf_counter() - started) * 1000.0)
|
|
write_value(shared, 16, "I", sequence)
|
|
write_value(shared, 8, "i", STATUS_READY)
|
|
last_sequence = sequence
|
|
response_event.set()
|
|
return 0
|
|
except BaseException:
|
|
message = traceback.format_exc()
|
|
try:
|
|
Path(args.config + ".log").write_text(message, encoding="utf-8")
|
|
except Exception:
|
|
pass
|
|
try:
|
|
if shared is not None:
|
|
write_status(shared, STATUS_ERROR, message[-500:])
|
|
if response_event is not None:
|
|
response_event.set()
|
|
except Exception:
|
|
pass
|
|
return 1
|
|
finally:
|
|
if request_event is not None:
|
|
request_event.close()
|
|
if response_event is not None:
|
|
response_event.close()
|
|
if shared is not None:
|
|
shared.close()
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--map", required=True)
|
|
parser.add_argument("--request", required=True)
|
|
parser.add_argument("--response", required=True)
|
|
parser.add_argument("--config", required=True)
|
|
return parser.parse_args()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(run(parse_args()))
|