Add CUDA Graph inference acceleration support

This commit is contained in:
RVC-Boss
2026-07-20 22:34:58 +08:00
parent f748ba4c48
commit b3f54df1be
26 changed files with 735 additions and 181 deletions

View File

@@ -10,6 +10,8 @@ from tools.file_io import read_text
import torch
import logging
from tools.cuda_graph import configure_cuda_graph
logger = logging.getLogger(__name__)
@@ -126,6 +128,12 @@ if infer_device.type != "cuda":
)
# Run a real capture/replay probe on the selected inference device. Both
# application entry points import this module, so downstream inference code
# receives one consistent 0/1 switch without duplicating device checks.
CUDA_GRAPH_AVAILABLE = configure_cuda_graph(infer_device)
CONFIGS_DIR = Path(__file__).resolve().parent
MODEL_CONFIG_FILES = (
"v1/32k.json",
@@ -152,6 +160,7 @@ class Config:
self.device = str(infer_device)
self.dtype = infer_dtype
self.is_half = infer_dtype == torch.float16
self.cuda_graph = CUDA_GRAPH_AVAILABLE
self.n_cpu = 0
self.gpu_name = None
self.json_config = self.load_config_json()

View File

@@ -257,5 +257,7 @@
"[索引训练] 外部索引链接已存在:%s": "[Index training] External index link already exists: %s",
"[索引训练][跳过] added索引已存在%s": "[Index training][Skipped] added index already exists: %s",
"[索引训练][跳过] trained索引已存在%s": "[Index training][Skipped] trained index already exists: %s",
"当前设备:%s | 推理精度:%s": "Current device: %s | Inference precision: %s"
"当前设备:%s | 推理精度:%s": "Current device: %s | Inference precision: %s",
"正在预热CUDA Graph": "Warming up CUDA Graph",
"CUDA Graph预热完成": "CUDA Graph warm-up complete"
}

View File

@@ -257,5 +257,7 @@
"[索引训练] 外部索引链接已存在:%s": "[Entrenamiento de índice] El enlace externo ya existe: %s",
"[索引训练][跳过] added索引已存在%s": "[Entrenamiento de índice][Omitido] El índice added ya existe: %s",
"[索引训练][跳过] trained索引已存在%s": "[Entrenamiento de índice][Omitido] El índice trained ya existe: %s",
"当前设备:%s | 推理精度:%s": "Dispositivo actual: %s | Precisión de inferencia: %s"
"当前设备:%s | 推理精度:%s": "Dispositivo actual: %s | Precisión de inferencia: %s",
"正在预热CUDA Graph": "Preparando CUDA Graph",
"CUDA Graph预热完成": "Preparación de CUDA Graph completada"
}

View File

@@ -257,5 +257,7 @@
"[索引训练] 外部索引链接已存在:%s": "[Entraînement de lindex] Le lien externe existe déjà : %s",
"[索引训练][跳过] added索引已存在%s": "[Entraînement de lindex][Ignoré] Lindex added existe déjà : %s",
"[索引训练][跳过] trained索引已存在%s": "[Entraînement de lindex][Ignoré] Lindex trained existe déjà : %s",
"当前设备:%s | 推理精度:%s": "Périphérique actuel : %s | Précision dinférence : %s"
"当前设备:%s | 推理精度:%s": "Périphérique actuel : %s | Précision dinférence : %s",
"正在预热CUDA Graph": "Préchauffage de CUDA Graph",
"CUDA Graph预热完成": "Préchauffage de CUDA Graph terminé"
}

View File

@@ -257,5 +257,7 @@
"[索引训练] 外部索引链接已存在:%s": "[Addestramento indice] Il collegamento esterno esiste già: %s",
"[索引训练][跳过] added索引已存在%s": "[Addestramento indice][Saltato] Lindice added esiste già: %s",
"[索引训练][跳过] trained索引已存在%s": "[Addestramento indice][Saltato] Lindice trained esiste già: %s",
"当前设备:%s | 推理精度:%s": "Dispositivo corrente: %s | Precisione di inferenza: %s"
"当前设备:%s | 推理精度:%s": "Dispositivo corrente: %s | Precisione di inferenza: %s",
"正在预热CUDA Graph": "Riscaldamento di CUDA Graph",
"CUDA Graph预热完成": "Riscaldamento di CUDA Graph completato"
}

View File

@@ -257,5 +257,7 @@
"[索引训练] 外部索引链接已存在:%s": "[インデックス学習] 外部インデックスリンクは既に存在します:%s",
"[索引训练][跳过] added索引已存在%s": "[インデックス学習][スキップ] addedインデックスは既に存在します%s",
"[索引训练][跳过] trained索引已存在%s": "[インデックス学習][スキップ] trainedインデックスは既に存在します%s",
"当前设备:%s | 推理精度:%s": "現在のデバイス:%s | 推論精度:%s"
"当前设备:%s | 推理精度:%s": "現在のデバイス:%s | 推論精度:%s",
"正在预热CUDA Graph": "CUDA Graphをウォームアップしています",
"CUDA Graph预热完成": "CUDA Graphのウォームアップが完了しました"
}

View File

@@ -257,5 +257,7 @@
"[索引训练] 外部索引链接已存在:%s": "[인덱스 학습] 외부 인덱스 링크가 이미 있습니다: %s",
"[索引训练][跳过] added索引已存在%s": "[인덱스 학습][건너뜀] added 인덱스가 이미 있습니다: %s",
"[索引训练][跳过] trained索引已存在%s": "[인덱스 학습][건너뜀] trained 인덱스가 이미 있습니다: %s",
"当前设备:%s | 推理精度:%s": "현재 장치: %s | 추론 정밀도: %s"
"当前设备:%s | 推理精度:%s": "현재 장치: %s | 추론 정밀도: %s",
"正在预热CUDA Graph": "CUDA Graph 워밍업 중",
"CUDA Graph预热完成": "CUDA Graph 워밍업 완료"
}

View File

@@ -257,5 +257,7 @@
"[索引训练] 外部索引链接已存在:%s": "[Treinamento de índice] O link externo já existe: %s",
"[索引训练][跳过] added索引已存在%s": "[Treinamento de índice][Ignorado] O índice added já existe: %s",
"[索引训练][跳过] trained索引已存在%s": "[Treinamento de índice][Ignorado] O índice trained já existe: %s",
"当前设备:%s | 推理精度:%s": "Dispositivo atual: %s | Precisão de inferência: %s"
"当前设备:%s | 推理精度:%s": "Dispositivo atual: %s | Precisão de inferência: %s",
"正在预热CUDA Graph": "Aquecendo o CUDA Graph",
"CUDA Graph预热完成": "Aquecimento do CUDA Graph concluído"
}

View File

@@ -257,5 +257,7 @@
"[索引训练] 外部索引链接已存在:%s": "[Обучение индекса] Внешняя ссылка уже существует: %s",
"[索引训练][跳过] added索引已存在%s": "[Обучение индекса][Пропущено] Индекс added уже существует: %s",
"[索引训练][跳过] trained索引已存在%s": "[Обучение индекса][Пропущено] Индекс trained уже существует: %s",
"当前设备:%s | 推理精度:%s": "Текущее устройство: %s | Точность вывода: %s"
"当前设备:%s | 推理精度:%s": "Текущее устройство: %s | Точность вывода: %s",
"正在预热CUDA Graph": "Прогрев CUDA Graph",
"CUDA Graph预热完成": "Прогрев CUDA Graph завершён"
}

View File

@@ -257,5 +257,7 @@
"[索引训练] 外部索引链接已存在:%s": "[Dizin eğitimi] Harici dizin bağlantısı zaten var: %s",
"[索引训练][跳过] added索引已存在%s": "[Dizin eğitimi][Atlandı] added dizini zaten var: %s",
"[索引训练][跳过] trained索引已存在%s": "[Dizin eğitimi][Atlandı] trained dizini zaten var: %s",
"当前设备:%s | 推理精度:%s": "Geçerli cihaz: %s | Çıkarım hassasiyeti: %s"
"当前设备:%s | 推理精度:%s": "Geçerli cihaz: %s | Çıkarım hassasiyeti: %s",
"正在预热CUDA Graph": "CUDA Graph ısınıyor",
"CUDA Graph预热完成": "CUDA Graph ısınması tamamlandı"
}

View File

@@ -257,5 +257,7 @@
"[索引训练] 外部索引链接已存在:%s": "[索引训练] 外部索引链接已存在:%s",
"[索引训练][跳过] added索引已存在%s": "[索引训练][跳过] added索引已存在%s",
"[索引训练][跳过] trained索引已存在%s": "[索引训练][跳过] trained索引已存在%s",
"当前设备:%s | 推理精度:%s": "当前设备:%s | 推理精度:%s"
"当前设备:%s | 推理精度:%s": "当前设备:%s | 推理精度:%s",
"正在预热CUDA Graph": "正在预热CUDA Graph",
"CUDA Graph预热完成": "CUDA Graph预热完成"
}

View File

@@ -257,5 +257,7 @@
"[索引训练] 外部索引链接已存在:%s": "[索引訓練] 外部索引連結已存在:%s",
"[索引训练][跳过] added索引已存在%s": "[索引訓練][跳過] added索引已存在%s",
"[索引训练][跳过] trained索引已存在%s": "[索引訓練][跳過] trained索引已存在%s",
"当前设备:%s | 推理精度:%s": "目前裝置:%s | 推理精度:%s"
"当前设备:%s | 推理精度:%s": "目前裝置:%s | 推理精度:%s",
"正在预热CUDA Graph": "正在預熱 CUDA Graph",
"CUDA Graph预热完成": "CUDA Graph 預熱完成"
}

View File

@@ -257,5 +257,7 @@
"[索引训练] 外部索引链接已存在:%s": "[索引训练] 外部索引链接已存在:%s",
"[索引训练][跳过] added索引已存在%s": "[索引训练][跳过] added索引已存在%s",
"[索引训练][跳过] trained索引已存在%s": "[索引训练][跳过] trained索引已存在%s",
"当前设备:%s | 推理精度:%s": "当前设备:%s | 推理精度:%s"
"当前设备:%s | 推理精度:%s": "当前设备:%s | 推理精度:%s",
"正在预热CUDA Graph": "正在预热CUDA Graph",
"CUDA Graph预热完成": "CUDA Graph预热完成"
}

View File

@@ -257,5 +257,7 @@
"[索引训练] 外部索引链接已存在:%s": "[索引訓練] 外部索引連結已存在:%s",
"[索引训练][跳过] added索引已存在%s": "[索引訓練][跳過] added索引已存在%s",
"[索引训练][跳过] trained索引已存在%s": "[索引訓練][跳過] trained索引已存在%s",
"当前设备:%s | 推理精度:%s": "目前裝置:%s | 推論精度:%s"
"当前设备:%s | 推理精度:%s": "目前裝置:%s | 推論精度:%s",
"正在预热CUDA Graph": "正在預熱 CUDA Graph",
"CUDA Graph预热完成": "CUDA Graph 預熱完成"
}

View File

@@ -1,5 +1,7 @@
import torch
from tools.cuda_graph import cuda_graph_enabled, run_cuda_graph
def _is_directml_device(device):
"""Return whether *device* is the PrivateUse1 device registered by DirectML."""
@@ -35,6 +37,46 @@ class FCPEInfer:
self.infer_model.model.to(device).eval()
else:
self.infer_model = spawn_bundled_infer_model(device)
if getattr(device, "type", None) == "cuda" or str(device).startswith(
"cuda"
):
# torchfcpe creates this tensor on CPU and copies it to CUDA in
# every local-argmax decode. Host-to-device copies are forbidden
# during CUDA Graph capture, so keep the immutable offsets on the
# target device and use the graph-safe equivalent decoder below.
self.local_offsets = torch.arange(
9, device=device, dtype=torch.long
).view(1, 1, 9)
def _graphable_model_infer(self, mel, decoder_mode, threshold):
"""Run FCPE network and an exactly equivalent capture-safe decoder."""
model = self.infer_model.model
latent = model(mel)
batch, frames, _ = latent.shape
cents = model.cent_table[None, None, :].expand(batch, frames, -1)
if decoder_mode == "argmax":
confidence = torch.max(latent, dim=-1, keepdim=True).values
decoded = torch.sum(cents * latent, dim=-1, keepdim=True) / torch.sum(
latent, dim=-1, keepdim=True
)
elif decoder_mode == "local_argmax":
confidence, max_index = torch.max(latent, dim=-1, keepdim=True)
local_index = self.local_offsets + (max_index - 4)
local_index = local_index.clamp(0, model.out_dims - 1)
local_cents = torch.gather(cents, -1, local_index)
local_latent = torch.gather(latent, -1, local_index)
decoded = torch.sum(
local_cents * local_latent, dim=-1, keepdim=True
) / torch.sum(local_latent, dim=-1, keepdim=True)
else:
raise ValueError("Unknown FCPE decoder mode: %s" % decoder_mode)
# Match torchfcpe's masking and cent-to-Hz formulas operation-for-operation.
confidence_mask = torch.ones_like(confidence)
confidence_mask.masked_fill_(confidence <= threshold, float("-inf"))
decoded = decoded * confidence_mask
return 10.0 * torch.pow(2.0, decoded / 1200.0)
def _decode_on_cpu(self, latent, decoder_mode, threshold):
"""Decode DML network logits on CPU with torchfcpe's exact formulas.
@@ -79,11 +121,33 @@ class FCPEInfer:
threshold=0.006,
):
if not self.is_directml:
return self.infer_model.infer(
wav,
wav = wav.to(self.device)
if cuda_graph_enabled(wav.device):
# Wav2MelModule contains tensor-dependent Python conditionals and
# cannot be captured. Running it eagerly also creates/caches its
# STFT window before the graph. The much larger FCPE neural net
# and decoder form the stable-shape CUDA Graph boundary.
mel = self.infer_model.wav2mel(wav, sr)
return run_cuda_graph(
self.infer_model.model,
"fcpe-core-%s-%s" % (decoder_mode, threshold),
lambda input_mel: self._graphable_model_infer(
input_mel,
decoder_mode,
threshold,
),
mel,
)
return run_cuda_graph(
self.infer_model,
"fcpe-%s-%s-%s" % (sr, decoder_mode, threshold),
lambda input_wav: self.infer_model.infer(
input_wav,
sr=sr,
decoder_mode=decoder_mode,
threshold=threshold,
),
wav,
)
wav_cpu = wav.detach().to(device="cpu", dtype=torch.float32)

View File

@@ -6,6 +6,8 @@ import torch
from torch import nn
from transformers import AutoFeatureExtractor, HubertModel
from tools.cuda_graph import run_cuda_graph
logger = logging.getLogger(__name__)
@@ -78,19 +80,48 @@ def extract_hubert_features(model, source, version, padding_mask=None):
attention_mask = (~padding_mask.bool()).long()
if version == "v1":
if attention_mask is None:
def forward(input_values):
outputs = model(
input_values=source,
attention_mask=attention_mask,
input_values=input_values,
attention_mask=None,
output_hidden_states=True,
return_dict=True,
)
features = outputs.hidden_states[9]
return model.final_proj(features)
return model.final_proj(outputs.hidden_states[9])
return run_cuda_graph(model, "hubert-v1-no-mask", forward, source)
def forward(input_values, mask):
outputs = model(
input_values=source,
attention_mask=attention_mask,
output_hidden_states=False,
input_values=input_values,
attention_mask=mask,
output_hidden_states=True,
return_dict=True,
)
return outputs.last_hidden_state
return model.final_proj(outputs.hidden_states[9])
return run_cuda_graph(
model, "hubert-v1-mask", forward, source, attention_mask
)
if attention_mask is None:
def forward(input_values):
return model(
input_values=input_values,
attention_mask=None,
output_hidden_states=False,
return_dict=True,
).last_hidden_state
return run_cuda_graph(model, "hubert-v2-no-mask", forward, source)
def forward(input_values, mask):
return model(
input_values=input_values,
attention_mask=mask,
output_hidden_states=False,
return_dict=True,
).last_hidden_state
return run_cuda_graph(model, "hubert-v2-mask", forward, source, attention_mask)

View File

@@ -67,8 +67,11 @@ class TextEncoder(nn.Module):
)
x = self.encoder(x * x_mask, x_mask)
if skip_head is not None:
assert isinstance(skip_head, torch.Tensor)
head = int(skip_head.item())
head = (
int(skip_head.item())
if isinstance(skip_head, torch.Tensor)
else int(skip_head)
)
x = x[:, :, head:]
x_mask = x_mask[:, :, head:]
stats = self.proj(x) * x_mask
@@ -231,8 +234,7 @@ class Generator(torch.nn.Module):
n_res = None,
):
if n_res is not None:
assert isinstance(n_res, torch.Tensor)
n = int(n_res.item())
n = int(n_res.item()) if isinstance(n_res, torch.Tensor) else int(n_res)
if n != x.shape[-1]:
x = F.interpolate(x, size=n, mode="linear")
x = self.conv_pre(x)
@@ -501,8 +503,7 @@ class GeneratorNSF(torch.nn.Module):
har_source, noi_source, uv = self.m_source(f0, self.upp)
har_source = har_source.transpose(1, 2)
if n_res is not None:
assert isinstance(n_res, torch.Tensor)
n = int(n_res.item())
n = int(n_res.item()) if isinstance(n_res, torch.Tensor) else int(n_res)
if n * self.upp != har_source.shape[-1]:
har_source = F.interpolate(har_source, size=n * self.upp, mode="linear")
if n != x.shape[-1]:
@@ -674,12 +675,14 @@ class SynthesizerTrnMs256NSFsid(nn.Module):
):
g = self.emb_g(sid).unsqueeze(-1)
if skip_head is not None and return_length is not None:
assert isinstance(skip_head, torch.Tensor)
assert isinstance(return_length, torch.Tensor)
head = int(skip_head.item())
length = int(return_length.item())
flow_head = torch.clamp(skip_head - 24, min=0)
dec_head = head - int(flow_head.item())
head = int(skip_head.item()) if isinstance(skip_head, torch.Tensor) else int(skip_head)
length = (
int(return_length.item())
if isinstance(return_length, torch.Tensor)
else int(return_length)
)
flow_head = max(head - 24, 0)
dec_head = head - flow_head
m_p, logs_p, x_mask = self.enc_p(phone, pitch, phone_lengths, flow_head)
z_p = (m_p + torch.exp(logs_p) * torch.randn_like(m_p) * 0.66666) * x_mask
z = self.flow(z_p, x_mask, g=g, reverse=True)
@@ -862,12 +865,14 @@ class SynthesizerTrnMs256NSFsid_nono(nn.Module):
):
g = self.emb_g(sid).unsqueeze(-1)
if skip_head is not None and return_length is not None:
assert isinstance(skip_head, torch.Tensor)
assert isinstance(return_length, torch.Tensor)
head = int(skip_head.item())
length = int(return_length.item())
flow_head = torch.clamp(skip_head - 24, min=0)
dec_head = head - int(flow_head.item())
head = int(skip_head.item()) if isinstance(skip_head, torch.Tensor) else int(skip_head)
length = (
int(return_length.item())
if isinstance(return_length, torch.Tensor)
else int(return_length)
)
flow_head = max(head - 24, 0)
dec_head = head - flow_head
m_p, logs_p, x_mask = self.enc_p(phone, None, phone_lengths, flow_head)
z_p = (m_p + torch.exp(logs_p) * torch.randn_like(m_p) * 0.66666) * x_mask
z = self.flow(z_p, x_mask, g=g, reverse=True)

View File

@@ -8,6 +8,8 @@ import torch.nn.functional as F
from librosa.util import normalize, pad_center, tiny
from scipy.signal import get_window
from tools.cuda_graph import run_cuda_graph
import logging
logger = logging.getLogger(__name__)
@@ -551,9 +553,29 @@ class RMVPE:
)[0]
else:
mel = mel.half() if self.is_half else mel.float()
hidden = self.model(mel)
hidden = run_cuda_graph(
self.model,
"rmvpe-network",
lambda input_mel: self.model(input_mel),
mel,
)
return hidden[:, :n_frames]
def extract_mel(self, audio, center=True):
if not torch.is_tensor(audio):
audio = torch.from_numpy(audio)
audio = audio.float().to(self.device)
if audio.dim() == 1:
audio = audio.unsqueeze(0)
if "privateuseone" in str(self.device):
return self.mel_extractor(audio, center=center)
return run_cuda_graph(
self.mel_extractor,
"rmvpe-mel-center-%s" % int(bool(center)),
lambda input_audio: self.mel_extractor(input_audio, center=center),
audio,
)
def decode(self, hidden, thred=0.03):
cents_pred = self.to_local_average_cents(hidden, thred=thred)
f0 = 10 * (2 ** (cents_pred / 1200))
@@ -564,11 +586,7 @@ class RMVPE:
def infer_from_audio(self, audio, thred=0.03):
# torch.cuda.synchronize()
# t0 = ttime()
if not torch.is_tensor(audio):
audio = torch.from_numpy(audio)
mel = self.mel_extractor(
audio.float().to(self.device).unsqueeze(0), center=True
)
mel = self.extract_mel(audio, center=True)
# print(123123123,mel.device.type)
# torch.cuda.synchronize()
# t1 = ttime()

View File

@@ -10,6 +10,7 @@ from torchaudio.transforms import Resample
from infer.hubert import extract_hubert_features, load_hubert_model
from i18n.i18n import I18nAuto
from tools.cuda_graph import run_cuda_graph
i18n = I18nAuto()
@@ -93,6 +94,7 @@ class RVC:
self.cache_pitchf = torch.zeros(
1024, device=self.device, dtype=torch.float32
)
self.infer_count = 0
self.resample_kernel = {}
@@ -179,11 +181,9 @@ class RVC:
if len(f0) < p_len:
f0 = np.pad(f0, (0, p_len - len(f0)))
f0 = f0[:p_len]
try:
uv = f0 == 0
if np.any(~uv):
f0[uv] = np.interp(np.where(uv)[0], np.where(~uv)[0], f0[~uv])
except Exception:
traceback.print_exc()
f0 *= pow(2, f0_up_key / 12)
return self.get_f0_post(f0)
@@ -198,11 +198,9 @@ class RVC:
device=self.device,
)
f0 = self.model_rmvpe.infer_from_audio(x, thred=0.03)
try:
uv = f0 == 0
if np.any(~uv):
f0[uv] = np.interp(np.where(uv)[0], np.where(~uv)[0], f0[~uv])
except Exception:
traceback.print_exc()
f0 *= pow(2, f0_up_key / 12)
return self.get_f0_post(f0)
@@ -218,11 +216,9 @@ class RVC:
decoder_mode="local_argmax",
threshold=0.006,
).squeeze().detach().cpu().numpy()
try:
uv = f0 == 0
if np.any(~uv):
f0[uv] = np.interp(np.where(uv)[0], np.where(~uv)[0], f0[~uv])
except Exception:
traceback.print_exc()
f0 *= pow(2, f0_up_key / 12)
return self.get_f0_post(f0)
@@ -234,6 +230,8 @@ class RVC:
return_length,
f0method,
) :
report_status = self.infer_count < 3 or self.infer_count % 100 == 0
self.infer_count += 1
t1 = ttime()
with torch.no_grad():
if self.config.is_half:
@@ -271,6 +269,7 @@ class RVC:
i18n("索引无效必须使用added_xxxx.index不能使用trained_xxxx.index")
)
else:
if report_status:
printt(i18n("索引检索失败或未启用"))
except Exception:
traceback.print_exc()
@@ -298,26 +297,49 @@ class RVC:
t4 = ttime()
feats = F.interpolate(feats.permute(0, 2, 1), scale_factor=2).permute(0, 2, 1)
feats = feats[:, :p_len, :]
p_len = torch.LongTensor([p_len]).to(self.device)
p_len_tensor = torch.LongTensor([p_len]).to(self.device)
sid = torch.LongTensor([0]).to(self.device)
skip_head = torch.LongTensor([skip_head])
return_length2 = torch.LongTensor([return_length2])
return_length = torch.LongTensor([return_length])
skip_head_value = int(skip_head)
return_length_value = int(return_length)
return_length2_value = int(return_length2)
with torch.no_grad():
if self.if_f0 == 1:
infered_audio, _, _ = self.net_g.infer(
infered_audio = run_cuda_graph(
self.net_g,
"rvc-realtime-f0-%s-%s-%s"
% (skip_head_value, return_length_value, return_length2_value),
lambda phone, lengths, coarse, continuous, speaker: self.net_g.infer(
phone,
lengths,
coarse,
continuous,
speaker,
skip_head_value,
return_length_value,
return_length2_value,
)[0],
feats,
p_len,
p_len_tensor,
cache_pitch,
cache_pitchf,
sid,
skip_head,
return_length,
return_length2,
)
else:
infered_audio, _, _ = self.net_g.infer(
feats, p_len, sid, skip_head, return_length, return_length2
infered_audio = run_cuda_graph(
self.net_g,
"rvc-realtime-no-f0-%s-%s-%s"
% (skip_head_value, return_length_value, return_length2_value),
lambda phone, lengths, speaker: self.net_g.infer(
phone,
lengths,
speaker,
skip_head_value,
return_length_value,
return_length2_value,
)[0],
feats,
p_len_tensor,
sid,
)
infered_audio = infered_audio.squeeze(1).float()
upp_res = int(np.floor(factor * self.tgt_sr // 100))
@@ -332,6 +354,7 @@ class RVC:
infered_audio[:, : return_length * upp_res]
)
t5 = ttime()
if report_status:
printt(
i18n("耗时:特征=%.3f秒,索引=%.3f秒,音高=%.3f秒,模型=%.3f"),
t2 - t1,

View File

@@ -19,6 +19,7 @@ from infer.vc.pipeline import Pipeline
from infer.vc.utils import *
from i18n.i18n import I18nAuto
from tools.progress import batch_status, should_report
from tools.cuda_graph import clear_cuda_graph_cache
i18n = I18nAuto()
@@ -68,6 +69,8 @@ class VC:
self.hubert_model is not None
): # 考虑到轮询, 需要加个判断看是否 sid 是由有模型切换到无模型的
logger.info(i18n("清理模型缓存"))
clear_cuda_graph_cache(self.net_g)
clear_cuda_graph_cache(self.hubert_model)
del (self.net_g, self.n_spk, self.hubert_model, self.tgt_sr) # ,cpt
self.hubert_model = self.net_g = self.n_spk = self.hubert_model = (
self.tgt_sr
@@ -112,6 +115,9 @@ class VC:
person = f'{os.getenv("weight_root")}/{sid}'
logger.info("%s: %s", i18n("正在加载模型"), person)
if self.net_g is not None:
clear_cuda_graph_cache(self.net_g)
self.cpt = torch.load(person, map_location="cpu")
self.tgt_sr = self.cpt["config"][-1]
self.cpt["config"][-3] = self.cpt["weight"]["emb_g.weight"].shape[0] # n_spk

View File

@@ -15,6 +15,7 @@ import torch.nn.functional as F
from scipy import signal
from infer.hubert import extract_hubert_features
from tools.cuda_graph import cuda_graph_enabled, run_cuda_graph
bh, ah = signal.butter(N=5, Wn=48, btype="high", fs=16000)
@@ -217,11 +218,34 @@ class Pipeline(object):
p_len = torch.tensor([p_len], device=self.device).long()
with torch.no_grad():
hasp = pitch is not None and pitchf is not None
arg = (feats, p_len, pitch, pitchf, sid) if hasp else (feats, p_len, sid)
audio1 = (net_g.infer(*arg)[0][0, 0]).data.cpu().float().numpy()
del hasp, arg
if hasp:
synthesized = run_cuda_graph(
net_g,
"rvc-synth-f0",
lambda phone, lengths, coarse, continuous, speaker: net_g.infer(
phone, lengths, coarse, continuous, speaker
)[0],
feats,
p_len,
pitch,
pitchf,
sid,
)
else:
synthesized = run_cuda_graph(
net_g,
"rvc-synth-no-f0",
lambda phone, lengths, speaker: net_g.infer(
phone, lengths, speaker
)[0],
feats,
p_len,
sid,
)
audio1 = synthesized[0, 0].data.cpu().float().numpy()
del hasp, synthesized
del feats, p_len, padding_mask
if torch.cuda.is_available():
if torch.cuda.is_available() and not cuda_graph_enabled(self.device):
torch.cuda.empty_cache()
t2 = ttime()
times[0] += t1 - t0
@@ -381,6 +405,6 @@ class Pipeline(object):
max_int16 /= audio_max
audio_opt = (audio_opt * max_int16).astype(np.int16)
del pitch, pitchf, sid
if torch.cuda.is_available():
if torch.cuda.is_available() and not cuda_graph_enabled(self.device):
torch.cuda.empty_cache()
return audio_opt

View File

@@ -36,9 +36,10 @@ if __name__ == "__main__":
import torch.nn.functional as F
import torchaudio.transforms as tat
from configs.config import Config
from infer import rtrvc as rvc_for_realtime
from i18n.i18n import I18nAuto
from configs.config import Config
from tools.cuda_graph import cuda_graph_enabled, run_cuda_graph
i18n = I18nAuto()
@@ -67,6 +68,7 @@ if __name__ == "__main__":
def __init__(self) :
self.gui_config = GUIConfig()
self.config = Config()
printt("RVC_CUDA_GRAPH=%s", os.environ.get("RVC_CUDA_GRAPH", "0"))
self.function = "vc"
self.delay_time = 0
self.hostapis = None
@@ -632,6 +634,13 @@ if __name__ == "__main__":
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.nr_buffer = self.sola_buffer.clone()
self.output_buffer = self.input_wav.clone()
self.skip_head = self.extra_frame // self.zc
@@ -669,8 +678,81 @@ if __name__ == "__main__":
self.tg = TorchGate(
sr=self.gui_config.samplerate, n_fft=4 * self.zc, prop_decrease=0.9
).to(self.config.device)
self.prewarm_cuda_graph()
self.start_stream()
def prewarm_cuda_graph(self):
if not cuda_graph_enabled(self.config.device):
return
try:
printt(i18n("正在预热CUDA Graph"))
samples = self.input_wav_res.shape[0]
phase = torch.arange(
samples, device=self.config.device, dtype=torch.float32
)
probe = 0.05 * torch.sin(2 * np.pi * 220.0 * phase / 16000.0)
self.input_wav_res.copy_(probe)
if self.gui_config.I_noise_reduce:
short = self.input_wav[
-self.sola_buffer_frame - self.block_frame :
].unsqueeze(0)
run_cuda_graph(
self.tg,
"realtime-input-noise-reduction",
lambda short_audio, full_audio: self.tg(
short_audio, full_audio
),
short,
self.input_wav.unsqueeze(0),
)
resample_input = self.input_wav[-self.block_frame - 2 * self.zc :]
run_cuda_graph(
self.resampler,
"realtime-input-resample",
lambda audio: self.resampler(audio),
resample_input,
)
inferred = self.rvc.infer(
self.input_wav_res,
self.block_frame_16k,
self.skip_head,
self.return_length,
self.gui_config.f0method,
)
if self.resampler2 is not None:
inferred = run_cuda_graph(
self.resampler2,
"realtime-output-resample",
lambda audio: self.resampler2(audio),
inferred,
)
if self.gui_config.O_noise_reduce:
run_cuda_graph(
self.tg,
"realtime-output-noise-reduction",
lambda short_audio, full_audio: self.tg(
short_audio, full_audio
),
inferred.unsqueeze(0),
self.output_buffer.unsqueeze(0),
)
torch.cuda.synchronize(self.config.device)
printt(i18n("CUDA Graph预热完成"))
except Exception:
printt(traceback.format_exc())
finally:
self.input_wav.zero_()
self.input_wav_denoise.zero_()
self.input_wav_res.zero_()
self.output_buffer.zero_()
self.sola_buffer.zero_()
self.nr_buffer.zero_()
self.rvc.cache_pitch.zero_()
self.rvc.cache_pitchf.zero_()
def start_stream(self):
global flag_vc
if not flag_vc:
@@ -739,8 +821,12 @@ if __name__ == "__main__":
self.block_frame :
].clone()
input_wav = self.input_wav[-self.sola_buffer_frame - self.block_frame :]
input_wav = self.tg(
input_wav.unsqueeze(0), self.input_wav.unsqueeze(0)
input_wav = run_cuda_graph(
self.tg,
"realtime-input-noise-reduction",
lambda short, full: self.tg(short, full),
input_wav.unsqueeze(0),
self.input_wav.unsqueeze(0),
).squeeze(0)
input_wav[: self.sola_buffer_frame] *= self.fade_in_window
input_wav[: self.sola_buffer_frame] += (
@@ -750,15 +836,23 @@ if __name__ == "__main__":
: self.block_frame
]
self.nr_buffer[:] = input_wav[self.block_frame :]
self.input_wav_res[-self.block_frame_16k - 160 :] = self.resampler(
self.input_wav_denoise[-self.block_frame - 2 * self.zc :]
resample_input = self.input_wav_denoise[
-self.block_frame - 2 * self.zc :
]
self.input_wav_res[-self.block_frame_16k - 160 :] = run_cuda_graph(
self.resampler,
"realtime-input-resample",
lambda audio: self.resampler(audio),
resample_input,
)[160:]
else:
self.input_wav_res[-160 * (indata.shape[0] // self.zc + 1) :] = (
self.resampler(self.input_wav[-indata.shape[0] - 2 * self.zc :])[
160:
]
)
resample_input = self.input_wav[-indata.shape[0] - 2 * self.zc :]
self.input_wav_res[-160 * (indata.shape[0] // self.zc + 1) :] = run_cuda_graph(
self.resampler,
"realtime-input-resample",
lambda audio: self.resampler(audio),
resample_input,
)[160:]
# infer
if self.function == "vc":
infer_wav = self.rvc.infer(
@@ -769,7 +863,12 @@ if __name__ == "__main__":
self.gui_config.f0method,
)
if self.resampler2 is not None:
infer_wav = self.resampler2(infer_wav)
infer_wav = run_cuda_graph(
self.resampler2,
"realtime-output-resample",
lambda audio: self.resampler2(audio),
infer_wav,
)
elif self.gui_config.I_noise_reduce:
infer_wav = self.input_wav_denoise[self.extra_frame :].clone()
else:
@@ -780,8 +879,12 @@ if __name__ == "__main__":
self.block_frame :
].clone()
self.output_buffer[-self.block_frame :] = infer_wav[-self.block_frame :]
infer_wav = self.tg(
infer_wav.unsqueeze(0), self.output_buffer.unsqueeze(0)
infer_wav = run_cuda_graph(
self.tg,
"realtime-output-noise-reduction",
lambda short, full: self.tg(short, full),
infer_wav.unsqueeze(0),
self.output_buffer.unsqueeze(0),
).squeeze(0)
# volume envelop mixing
if self.gui_config.rms_mix_rate < 1 and self.function == "vc":
@@ -815,7 +918,7 @@ if __name__ == "__main__":
)[0, 0, :-1]
rms2 = torch.max(rms2, torch.zeros_like(rms2) + 1e-3)
infer_wav *= torch.pow(
rms1 / rms2, torch.tensor(1 - self.gui_config.rms_mix_rate)
rms1 / rms2, 1.0 - self.gui_config.rms_mix_rate
)
# SOLA algorithm from https://github.com/yxlllc/DDSP-SVC
conv_input = infer_wav[
@@ -825,7 +928,7 @@ if __name__ == "__main__":
cor_den = torch.sqrt(
F.conv1d(
conv_input**2,
torch.ones(1, 1, self.sola_buffer_frame, device=self.config.device),
self.sola_den_kernel,
)
+ 1e-8
)

227
tools/cuda_graph.py Normal file
View File

@@ -0,0 +1,227 @@
import logging
import os
import threading
import time
from collections import OrderedDict
import torch
logger = logging.getLogger(__name__)
ENV_NAME = "RVC_CUDA_GRAPH"
MAX_CACHE_ENV = "RVC_CUDA_GRAPH_MAX_CACHE"
_probe_lock = threading.Lock()
_probe_result = None
def _device_type(device):
if isinstance(device, torch.device):
return device.type
return str(device).split(":", 1)[0].lower()
def _cuda_device(device):
parsed = device if isinstance(device, torch.device) else torch.device(device)
if parsed.index is None:
parsed = torch.device("cuda", torch.cuda.current_device())
return parsed
def _clone_output(value):
if torch.is_tensor(value):
return value.clone()
if isinstance(value, tuple):
return tuple(_clone_output(item) for item in value)
if isinstance(value, list):
return [_clone_output(item) for item in value]
if isinstance(value, dict):
return {key: _clone_output(item) for key, item in value.items()}
return value
def detect_cuda_graph_support(device):
if _device_type(device) != "cuda" or not torch.cuda.is_available():
return False
if not hasattr(torch.cuda, "CUDAGraph") or not hasattr(torch.cuda, "graph"):
return False
cuda_device = _cuda_device(device)
try:
with torch.cuda.device(cuda_device):
current = torch.cuda.current_stream(cuda_device)
warmup = torch.cuda.Stream(device=cuda_device)
warmup.wait_stream(current)
with torch.cuda.stream(warmup):
probe = torch.arange(32, device=cuda_device, dtype=torch.float32)
for _ in range(3):
expected = probe.square().add_(1)
current.wait_stream(warmup)
torch.cuda.synchronize(cuda_device)
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
captured = probe.square() + 1
probe.copy_(torch.arange(32, device=cuda_device, dtype=torch.float32))
graph.replay()
torch.cuda.synchronize(cuda_device)
valid = torch.equal(
captured.cpu(), torch.arange(32, dtype=torch.float32).square() + 1
)
del captured, expected, graph, probe
return bool(valid)
except Exception:
logger.exception("CUDA Graph support probe failed on %s", cuda_device)
return False
def configure_cuda_graph(device):
global _probe_result
explicit = os.environ.get(ENV_NAME)
if explicit in {"0", "1"}:
if explicit == "0":
return False
if _device_type(device) != "cuda":
os.environ[ENV_NAME] = "0"
return False
with _probe_lock:
if _probe_result is None:
_probe_result = detect_cuda_graph_support(device)
os.environ[ENV_NAME] = "1" if _probe_result else "0"
return bool(_probe_result)
def cuda_graph_enabled(device):
return (
os.environ.get(ENV_NAME) == "1"
and _device_type(device) == "cuda"
and torch.cuda.is_available()
)
def _tensor_signature(tensor):
return (
tuple(tensor.shape),
tuple(tensor.stride()),
str(tensor.dtype),
str(tensor.device),
bool(tensor.requires_grad),
)
class _CapturedCall:
def __init__(self, function, inputs):
started = time.perf_counter()
self.lock = threading.RLock()
self.inputs = tuple(torch.empty_like(value) for value in inputs)
for static, value in zip(self.inputs, inputs):
static.copy_(value)
device = self.inputs[0].device
current = torch.cuda.current_stream(device)
warmup = torch.cuda.Stream(device=device)
warmup.wait_stream(current)
with torch.cuda.stream(warmup), torch.no_grad():
for _ in range(3):
output = function(*self.inputs)
current.wait_stream(warmup)
torch.cuda.synchronize(device)
self.graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(self.graph), torch.no_grad():
self.output = function(*self.inputs)
self.capture_ms = (time.perf_counter() - started) * 1000.0
self.done_event = None
del output
def replay(self, inputs):
with self.lock:
stream = torch.cuda.current_stream(self.inputs[0].device)
if self.done_event is not None:
stream.wait_event(self.done_event)
for static, value in zip(self.inputs, inputs):
static.copy_(value, non_blocking=True)
self.graph.replay()
output = _clone_output(self.output)
self.done_event = torch.cuda.Event(blocking=False)
self.done_event.record(stream)
return output
class _GraphCache:
def __init__(self):
self.entries = OrderedDict()
self.failures = set()
self.lock = threading.RLock()
self.capture_count = 0
self.replay_count = 0
self.fallback_count = 0
self.eviction_count = 0
self.capture_ms = 0.0
def run(self, key, function, inputs):
signature = key + tuple(_tensor_signature(value) for value in inputs)
with self.lock:
if signature in self.failures:
self.fallback_count += 1
return function(*inputs)
entry = self.entries.get(signature)
if entry is None:
try:
entry = _CapturedCall(function, inputs)
self.entries[signature] = entry
self.capture_count += 1
self.capture_ms += entry.capture_ms
max_entries = max(1, int(os.environ.get(MAX_CACHE_ENV, "8")))
while len(self.entries) > max_entries:
self.entries.popitem(last=False)
self.eviction_count += 1
except Exception:
self.failures.add(signature)
self.fallback_count += 1
logger.exception("CUDA Graph capture failed for %s; using eager", key)
return function(*inputs)
else:
self.entries.move_to_end(signature)
output = entry.replay(inputs)
with self.lock:
self.replay_count += 1
return output
def run_cuda_graph(owner, namespace, function, *inputs):
if not inputs or not cuda_graph_enabled(inputs[0].device):
return function(*inputs)
cache = getattr(owner, "_rvc_cuda_graph_cache", None)
if cache is None:
cache = _GraphCache()
setattr(owner, "_rvc_cuda_graph_cache", cache)
return cache.run((str(namespace),), function, tuple(inputs))
def clear_cuda_graph_cache(owner):
cache = getattr(owner, "_rvc_cuda_graph_cache", None)
if cache is not None:
cache.entries.clear()
cache.failures.clear()
delattr(owner, "_rvc_cuda_graph_cache")
def get_cuda_graph_stats(owner):
cache = getattr(owner, "_rvc_cuda_graph_cache", None)
if cache is None:
return {
"entries": 0,
"failures": 0,
"captures": 0,
"replays": 0,
"fallbacks": 0,
"evictions": 0,
"capture_ms": 0.0,
}
with cache.lock:
return {
"entries": len(cache.entries),
"failures": len(cache.failures),
"captures": cache.capture_count,
"replays": cache.replay_count,
"fallbacks": cache.fallback_count,
"evictions": cache.eviction_count,
"capture_ms": cache.capture_ms,
}

View File

@@ -9,6 +9,8 @@ import soundfile as sf
import torch
import torch.nn as nn
import yaml
from tools.cuda_graph import run_cuda_graph
from tqdm import tqdm
from tools.file_io import read_text
from i18n.i18n import I18nAuto
@@ -178,7 +180,12 @@ class Roformer_Loader:
if len(batch_data) >= batch_size or (i >= mix.shape[1]):
arr = torch.stack(batch_data, dim=0)
# print(23333333,arr.dtype)
x = model(arr)
x = run_cuda_graph(
model,
"uvr-bsroformer",
lambda audio: model(audio),
arr,
)
window = window_middle
if i - step == 0: # First audio chunk, no fadein

View File

@@ -1,5 +1,6 @@
import numpy as np
import torch
from tools.cuda_graph import run_cuda_graph
from tqdm import tqdm
@@ -34,7 +35,12 @@ def inference(X_spec, device, model, aggressiveness, data):
X_mag_window = X_mag_window.half()
X_mag_window = X_mag_window.to(device)
pred = model.predict(X_mag_window, aggressiveness)
pred = run_cuda_graph(
model,
"uvr-vr-%s" % repr(aggressiveness),
lambda window: model.predict(window, aggressiveness),
X_mag_window,
)
pred = pred.detach().cpu().numpy()
preds.append(pred[0])

View File

@@ -25,6 +25,7 @@ for name in os.listdir(tmp):
except Exception as error:
print(str(error))
from configs.config import Config, GPU_INDEX, GPU_INFOS, GPU_MEMORY, IS_GPU
from infer.vc.modules import VC
from tools.uvr5.webui import uvr
from tools.file_io import read_text
@@ -35,7 +36,6 @@ from train.process_ckpt import (
show_info,
)
from i18n.i18n import I18nAuto
from configs.config import Config, GPU_INDEX, GPU_INFOS, GPU_MEMORY, IS_GPU
import torch, platform
import numpy as np
import gradio as gr
@@ -143,6 +143,7 @@ print(
i18n("当前设备:%s | 推理精度:%s") % (config.device, config.dtype),
flush=True,
)
logger.info("RVC_CUDA_GRAPH=%s", os.environ.get("RVC_CUDA_GRAPH", "0"))
# GPU filtering and precision rules are shared with inference/extraction/training.
gpu_infos = list(GPU_INFOS)
gpu_indices = sorted(GPU_INDEX)
@@ -394,6 +395,10 @@ def train_task_stopped(state):
def start_train_process(state, cmd):
kwargs = {"shell": True, "cwd": now_dir}
if "train/train.py" in cmd.replace("\\", "/"):
training_env = os.environ.copy()
training_env["RVC_CUDA_GRAPH"] = "0"
kwargs["env"] = training_env
if platform.system() == "Windows":
kwargs["creationflags"] = subprocess.CREATE_NEW_PROCESS_GROUP
else: