Retrieval-based-Voice-Conversion-WebUI/RVCRealtimeVST/worker/rvc_worker.py

"""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()))