RVC_HF2

Runtime error

App Files Files Community

RVC_HF2 / infer /lib /rmvpe.py

r3gm

Upload 288 files

7bc29af 9 months ago

raw history blame contribute delete

No virus

25.4 kB

	import pdb, os

	import numpy as np
	import torch
	try:
	#Fix "Torch not compiled with CUDA enabled"
	import intel_extension_for_pytorch as ipex # pylint: disable=import-error, unused-import
	if torch.xpu.is_available():
	from infer.modules.ipex import ipex_init
	ipex_init()
	except Exception:
	pass
	import torch.nn as nn
	import torch.nn.functional as F
	from librosa.util import normalize, pad_center, tiny
	from scipy.signal import get_window

	import logging

	logger = logging.getLogger(__name__)


	###stft codes from https://github.com/pseeth/torch-stft/blob/master/torch_stft/util.py
	def window_sumsquare(
	window,
	n_frames,
	hop_length=200,
	win_length=800,
	n_fft=800,
	dtype=np.float32,
	norm=None,
	):
	"""
	# from librosa 0.6
	Compute the sum-square envelope of a window function at a given hop length.
	This is used to estimate modulation effects induced by windowing
	observations in short-time fourier transforms.
	Parameters
	----------
	window : string, tuple, number, callable, or list-like
	Window specification, as in `get_window`
	n_frames : int > 0
	The number of analysis frames
	hop_length : int > 0
	The number of samples to advance between frames
	win_length : [optional]
	The length of the window function. By default, this matches `n_fft`.
	n_fft : int > 0
	The length of each analysis frame.
	dtype : np.dtype
	The data type of the output
	Returns
	-------
	wss : np.ndarray, shape=`(n_fft + hop_length * (n_frames - 1))`
	The sum-squared envelope of the window function
	"""
	if win_length is None:
	win_length = n_fft

	n = n_fft + hop_length * (n_frames - 1)
	x = np.zeros(n, dtype=dtype)

	# Compute the squared window at the desired length
	win_sq = get_window(window, win_length, fftbins=True)
	win_sq = normalize(win_sq, norm=norm) ** 2
	win_sq = pad_center(win_sq, n_fft)

	# Fill the envelope
	for i in range(n_frames):
	sample = i * hop_length
	x[sample : min(n, sample + n_fft)] += win_sq[: max(0, min(n_fft, n - sample))]
	return x


	class STFT(torch.nn.Module):
	def __init__(
	self, filter_length=1024, hop_length=512, win_length=None, window="hann"
	):
	"""
	This module implements an STFT using 1D convolution and 1D transpose convolutions.
	This is a bit tricky so there are some cases that probably won't work as working
	out the same sizes before and after in all overlap add setups is tough. Right now,
	this code should work with hop lengths that are half the filter length (50% overlap
	between frames).

	Keyword Arguments:
	filter_length {int} -- Length of filters used (default: {1024})
	hop_length {int} -- Hop length of STFT (restrict to 50% overlap between frames) (default: {512})
	win_length {[type]} -- Length of the window function applied to each frame (if not specified, it
	equals the filter length). (default: {None})
	window {str} -- Type of window to use (options are bartlett, hann, hamming, blackman, blackmanharris)
	(default: {'hann'})
	"""
	super(STFT, self).__init__()
	self.filter_length = filter_length
	self.hop_length = hop_length
	self.win_length = win_length if win_length else filter_length
	self.window = window
	self.forward_transform = None
	self.pad_amount = int(self.filter_length / 2)
	scale = self.filter_length / self.hop_length
	fourier_basis = np.fft.fft(np.eye(self.filter_length))

	cutoff = int((self.filter_length / 2 + 1))
	fourier_basis = np.vstack(
	[np.real(fourier_basis[:cutoff, :]), np.imag(fourier_basis[:cutoff, :])]
	)
	forward_basis = torch.FloatTensor(fourier_basis[:, None, :])
	inverse_basis = torch.FloatTensor(
	np.linalg.pinv(scale * fourier_basis).T[:, None, :]
	)

	assert filter_length >= self.win_length
	# get window and zero center pad it to filter_length
	fft_window = get_window(window, self.win_length, fftbins=True)
	fft_window = pad_center(fft_window, size=filter_length)
	fft_window = torch.from_numpy(fft_window).float()

	# window the bases
	forward_basis *= fft_window
	inverse_basis *= fft_window

	self.register_buffer("forward_basis", forward_basis.float())
	self.register_buffer("inverse_basis", inverse_basis.float())

	def transform(self, input_data):
	"""Take input data (audio) to STFT domain.

	Arguments:
	input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)

	Returns:
	magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
	num_frequencies, num_frames)
	phase {tensor} -- Phase of STFT with shape (num_batch,
	num_frequencies, num_frames)
	"""
	num_batches = input_data.shape[0]
	num_samples = input_data.shape[-1]

	self.num_samples = num_samples

	# similar to librosa, reflect-pad the input
	input_data = input_data.view(num_batches, 1, num_samples)
	# print(1234,input_data.shape)
	input_data = F.pad(
	input_data.unsqueeze(1),
	(self.pad_amount, self.pad_amount, 0, 0, 0, 0),
	mode="reflect",
	).squeeze(1)
	# print(2333,input_data.shape,self.forward_basis.shape,self.hop_length)
	# pdb.set_trace()
	forward_transform = F.conv1d(
	input_data, self.forward_basis, stride=self.hop_length, padding=0
	)

	cutoff = int((self.filter_length / 2) + 1)
	real_part = forward_transform[:, :cutoff, :]
	imag_part = forward_transform[:, cutoff:, :]

	magnitude = torch.sqrt(real_part2 + imag_part2)
	# phase = torch.atan2(imag_part.data, real_part.data)

	return magnitude # , phase

	def inverse(self, magnitude, phase):
	"""Call the inverse STFT (iSTFT), given magnitude and phase tensors produced
	by the ```transform``` function.

	Arguments:
	magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
	num_frequencies, num_frames)
	phase {tensor} -- Phase of STFT with shape (num_batch,
	num_frequencies, num_frames)

	Returns:
	inverse_transform {tensor} -- Reconstructed audio given magnitude and phase. Of
	shape (num_batch, num_samples)
	"""
	recombine_magnitude_phase = torch.cat(
	[magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1
	)

	inverse_transform = F.conv_transpose1d(
	recombine_magnitude_phase,
	self.inverse_basis,
	stride=self.hop_length,
	padding=0,
	)

	if self.window is not None:
	window_sum = window_sumsquare(
	self.window,
	magnitude.size(-1),
	hop_length=self.hop_length,
	win_length=self.win_length,
	n_fft=self.filter_length,
	dtype=np.float32,
	)
	# remove modulation effects
	approx_nonzero_indices = torch.from_numpy(
	np.where(window_sum > tiny(window_sum))[0]
	)
	window_sum = torch.from_numpy(window_sum).to(inverse_transform.device)
	inverse_transform[:, :, approx_nonzero_indices] /= window_sum[
	approx_nonzero_indices
	]

	# scale by hop ratio
	inverse_transform *= float(self.filter_length) / self.hop_length

	inverse_transform = inverse_transform[..., self.pad_amount :]
	inverse_transform = inverse_transform[..., : self.num_samples]
	inverse_transform = inverse_transform.squeeze(1)

	return inverse_transform

	def forward(self, input_data):
	"""Take input data (audio) to STFT domain and then back to audio.

	Arguments:
	input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)

	Returns:
	reconstruction {tensor} -- Reconstructed audio given magnitude and phase. Of
	shape (num_batch, num_samples)
	"""
	self.magnitude, self.phase = self.transform(input_data)
	reconstruction = self.inverse(self.magnitude, self.phase)
	return reconstruction


	from time import time as ttime


	class BiGRU(nn.Module):
	def __init__(self, input_features, hidden_features, num_layers):
	super(BiGRU, self).__init__()
	self.gru = nn.GRU(
	input_features,
	hidden_features,
	num_layers=num_layers,
	batch_first=True,
	bidirectional=True,
	)

	def forward(self, x):
	return self.gru(x)[0]


	class ConvBlockRes(nn.Module):
	def __init__(self, in_channels, out_channels, momentum=0.01):
	super(ConvBlockRes, self).__init__()
	self.conv = nn.Sequential(
	nn.Conv2d(
	in_channels=in_channels,
	out_channels=out_channels,
	kernel_size=(3, 3),
	stride=(1, 1),
	padding=(1, 1),
	bias=False,
	),
	nn.BatchNorm2d(out_channels, momentum=momentum),
	nn.ReLU(),
	nn.Conv2d(
	in_channels=out_channels,
	out_channels=out_channels,
	kernel_size=(3, 3),
	stride=(1, 1),
	padding=(1, 1),
	bias=False,
	),
	nn.BatchNorm2d(out_channels, momentum=momentum),
	nn.ReLU(),
	)
	if in_channels != out_channels:
	self.shortcut = nn.Conv2d(in_channels, out_channels, (1, 1))
	self.is_shortcut = True
	else:
	self.is_shortcut = False

	def forward(self, x):
	if self.is_shortcut:
	return self.conv(x) + self.shortcut(x)
	else:
	return self.conv(x) + x


	class Encoder(nn.Module):
	def __init__(
	self,
	in_channels,
	in_size,
	n_encoders,
	kernel_size,
	n_blocks,
	out_channels=16,
	momentum=0.01,
	):
	super(Encoder, self).__init__()
	self.n_encoders = n_encoders
	self.bn = nn.BatchNorm2d(in_channels, momentum=momentum)
	self.layers = nn.ModuleList()
	self.latent_channels = []
	for i in range(self.n_encoders):
	self.layers.append(
	ResEncoderBlock(
	in_channels, out_channels, kernel_size, n_blocks, momentum=momentum
	)
	)
	self.latent_channels.append([out_channels, in_size])
	in_channels = out_channels
	out_channels *= 2
	in_size //= 2
	self.out_size = in_size
	self.out_channel = out_channels

	def forward(self, x):
	concat_tensors = []
	x = self.bn(x)
	for i in range(self.n_encoders):
	_, x = self.layers[i](x)
	concat_tensors.append(_)
	return x, concat_tensors


	class ResEncoderBlock(nn.Module):
	def __init__(
	self, in_channels, out_channels, kernel_size, n_blocks=1, momentum=0.01
	):
	super(ResEncoderBlock, self).__init__()
	self.n_blocks = n_blocks
	self.conv = nn.ModuleList()
	self.conv.append(ConvBlockRes(in_channels, out_channels, momentum))
	for i in range(n_blocks - 1):
	self.conv.append(ConvBlockRes(out_channels, out_channels, momentum))
	self.kernel_size = kernel_size
	if self.kernel_size is not None:
	self.pool = nn.AvgPool2d(kernel_size=kernel_size)

	def forward(self, x):
	for i in range(self.n_blocks):
	x = self.conv[i](x)
	if self.kernel_size is not None:
	return x, self.pool(x)
	else:
	return x


	class Intermediate(nn.Module): #
	def __init__(self, in_channels, out_channels, n_inters, n_blocks, momentum=0.01):
	super(Intermediate, self).__init__()
	self.n_inters = n_inters
	self.layers = nn.ModuleList()
	self.layers.append(
	ResEncoderBlock(in_channels, out_channels, None, n_blocks, momentum)
	)
	for i in range(self.n_inters - 1):
	self.layers.append(
	ResEncoderBlock(out_channels, out_channels, None, n_blocks, momentum)
	)

	def forward(self, x):
	for i in range(self.n_inters):
	x = self.layers[i](x)
	return x


	class ResDecoderBlock(nn.Module):
	def __init__(self, in_channels, out_channels, stride, n_blocks=1, momentum=0.01):
	super(ResDecoderBlock, self).__init__()
	out_padding = (0, 1) if stride == (1, 2) else (1, 1)
	self.n_blocks = n_blocks
	self.conv1 = nn.Sequential(
	nn.ConvTranspose2d(
	in_channels=in_channels,
	out_channels=out_channels,
	kernel_size=(3, 3),
	stride=stride,
	padding=(1, 1),
	output_padding=out_padding,
	bias=False,
	),
	nn.BatchNorm2d(out_channels, momentum=momentum),
	nn.ReLU(),
	)
	self.conv2 = nn.ModuleList()
	self.conv2.append(ConvBlockRes(out_channels * 2, out_channels, momentum))
	for i in range(n_blocks - 1):
	self.conv2.append(ConvBlockRes(out_channels, out_channels, momentum))

	def forward(self, x, concat_tensor):
	x = self.conv1(x)
	x = torch.cat((x, concat_tensor), dim=1)
	for i in range(self.n_blocks):
	x = self.conv2[i](x)
	return x


	class Decoder(nn.Module):
	def __init__(self, in_channels, n_decoders, stride, n_blocks, momentum=0.01):
	super(Decoder, self).__init__()
	self.layers = nn.ModuleList()
	self.n_decoders = n_decoders
	for i in range(self.n_decoders):
	out_channels = in_channels // 2
	self.layers.append(
	ResDecoderBlock(in_channels, out_channels, stride, n_blocks, momentum)
	)
	in_channels = out_channels

	def forward(self, x, concat_tensors):
	for i in range(self.n_decoders):
	x = self.layers[i](x, concat_tensors[-1 - i])
	return x


	class DeepUnet(nn.Module):
	def __init__(
	self,
	kernel_size,
	n_blocks,
	en_de_layers=5,
	inter_layers=4,
	in_channels=1,
	en_out_channels=16,
	):
	super(DeepUnet, self).__init__()
	self.encoder = Encoder(
	in_channels, 128, en_de_layers, kernel_size, n_blocks, en_out_channels
	)
	self.intermediate = Intermediate(
	self.encoder.out_channel // 2,
	self.encoder.out_channel,
	inter_layers,
	n_blocks,
	)
	self.decoder = Decoder(
	self.encoder.out_channel, en_de_layers, kernel_size, n_blocks
	)

	def forward(self, x):
	x, concat_tensors = self.encoder(x)
	x = self.intermediate(x)
	x = self.decoder(x, concat_tensors)
	return x


	class E2E(nn.Module):
	def __init__(
	self,
	n_blocks,
	n_gru,
	kernel_size,
	en_de_layers=5,
	inter_layers=4,
	in_channels=1,
	en_out_channels=16,
	):
	super(E2E, self).__init__()
	self.unet = DeepUnet(
	kernel_size,
	n_blocks,
	en_de_layers,
	inter_layers,
	in_channels,
	en_out_channels,
	)
	self.cnn = nn.Conv2d(en_out_channels, 3, (3, 3), padding=(1, 1))
	if n_gru:
	self.fc = nn.Sequential(
	BiGRU(3 * 128, 256, n_gru),
	nn.Linear(512, 360),
	nn.Dropout(0.25),
	nn.Sigmoid(),
	)
	else:
	self.fc = nn.Sequential(
	nn.Linear(3 * nn.N_MELS, nn.N_CLASS), nn.Dropout(0.25), nn.Sigmoid()
	)

	def forward(self, mel):
	# print(mel.shape)
	mel = mel.transpose(-1, -2).unsqueeze(1)
	x = self.cnn(self.unet(mel)).transpose(1, 2).flatten(-2)
	x = self.fc(x)
	# print(x.shape)
	return x


	from librosa.filters import mel


	class MelSpectrogram(torch.nn.Module):
	def __init__(
	self,
	is_half,
	n_mel_channels,
	sampling_rate,
	win_length,
	hop_length,
	n_fft=None,
	mel_fmin=0,
	mel_fmax=None,
	clamp=1e-5,
	):
	super().__init__()
	n_fft = win_length if n_fft is None else n_fft
	self.hann_window = {}
	mel_basis = mel(
	sr=sampling_rate,
	n_fft=n_fft,
	n_mels=n_mel_channels,
	fmin=mel_fmin,
	fmax=mel_fmax,
	htk=True,
	)
	mel_basis = torch.from_numpy(mel_basis).float()
	self.register_buffer("mel_basis", mel_basis)
	self.n_fft = win_length if n_fft is None else n_fft
	self.hop_length = hop_length
	self.win_length = win_length
	self.sampling_rate = sampling_rate
	self.n_mel_channels = n_mel_channels
	self.clamp = clamp
	self.is_half = is_half

	def forward(self, audio, keyshift=0, speed=1, center=True):
	factor = 2 ** (keyshift / 12)
	n_fft_new = int(np.round(self.n_fft * factor))
	win_length_new = int(np.round(self.win_length * factor))
	hop_length_new = int(np.round(self.hop_length * speed))
	keyshift_key = str(keyshift) + "_" + str(audio.device)
	if keyshift_key not in self.hann_window:
	self.hann_window[keyshift_key] = torch.hann_window(win_length_new).to(
	# "cpu"if(audio.device.type=="privateuseone") else audio.device
	audio.device
	)
	# fft = torch.stft(#doesn't support pytorch_dml
	# # audio.cpu() if(audio.device.type=="privateuseone")else audio,
	# audio,
	# n_fft=n_fft_new,
	# hop_length=hop_length_new,
	# win_length=win_length_new,
	# window=self.hann_window[keyshift_key],
	# center=center,
	# return_complex=True,
	# )
	# magnitude = torch.sqrt(fft.real.pow(2) + fft.imag.pow(2))
	# print(1111111111)
	# print(222222222222222,audio.device,self.is_half)
	if hasattr(self, "stft") == False:
	# print(n_fft_new,hop_length_new,win_length_new,audio.shape)
	self.stft = STFT(
	filter_length=n_fft_new,
	hop_length=hop_length_new,
	win_length=win_length_new,
	window="hann",
	).to(audio.device)
	magnitude = self.stft.transform(audio) # phase
	# if (audio.device.type == "privateuseone"):
	# magnitude=magnitude.to(audio.device)
	if keyshift != 0:
	size = self.n_fft // 2 + 1
	resize = magnitude.size(1)
	if resize < size:
	magnitude = F.pad(magnitude, (0, 0, 0, size - resize))
	magnitude = magnitude[:, :size, :] * self.win_length / win_length_new
	mel_output = torch.matmul(self.mel_basis, magnitude)
	if self.is_half == True:
	mel_output = mel_output.half()
	log_mel_spec = torch.log(torch.clamp(mel_output, min=self.clamp))
	# print(log_mel_spec.device.type)
	return log_mel_spec


	class RMVPE:
	def __init__(self, model_path, is_half, device=None):
	self.resample_kernel = {}
	self.resample_kernel = {}
	self.is_half = is_half
	if device is None:
	device = "cuda" if torch.cuda.is_available() else "cpu"
	self.device = device
	self.mel_extractor = MelSpectrogram(
	is_half, 128, 16000, 1024, 160, None, 30, 8000
	).to(device)
	if "privateuseone" in str(device):
	import onnxruntime as ort

	ort_session = ort.InferenceSession(
	"%s/rmvpe.onnx" % os.environ["rmvpe_root"],
	providers=["DmlExecutionProvider"],
	)
	self.model = ort_session
	else:
	model = E2E(4, 1, (2, 2))
	ckpt = torch.load(model_path, map_location="cpu")
	model.load_state_dict(ckpt)
	model.eval()
	if is_half == True:
	model = model.half()
	self.model = model
	self.model = self.model.to(device)
	cents_mapping = 20 * np.arange(360) + 1997.3794084376191
	self.cents_mapping = np.pad(cents_mapping, (4, 4)) # 368

	def mel2hidden(self, mel):
	with torch.no_grad():
	n_frames = mel.shape[-1]
	mel = F.pad(
	mel, (0, 32 * ((n_frames - 1) // 32 + 1) - n_frames), mode="constant"
	)
	if "privateuseone" in str(self.device):
	onnx_input_name = self.model.get_inputs()[0].name
	onnx_outputs_names = self.model.get_outputs()[0].name
	hidden = self.model.run(
	[onnx_outputs_names],
	input_feed={onnx_input_name: mel.cpu().numpy()},
	)[0]
	else:
	hidden = self.model(mel)
	return hidden[:, :n_frames]

	def decode(self, hidden, thred=0.03):
	cents_pred = self.to_local_average_cents(hidden, thred=thred)
	f0 = 10 * (2 ** (cents_pred / 1200))
	f0[f0 == 10] = 0
	# f0 = np.array([10 * (2 ** (cent_pred / 1200)) if cent_pred else 0 for cent_pred in cents_pred])
	return f0

	def infer_from_audio(self, audio, thred=0.03):
	# torch.cuda.synchronize()
	t0 = ttime()
	mel = self.mel_extractor(
	torch.from_numpy(audio).float().to(self.device).unsqueeze(0), center=True
	)
	# print(123123123,mel.device.type)
	# torch.cuda.synchronize()
	t1 = ttime()
	hidden = self.mel2hidden(mel)
	# torch.cuda.synchronize()
	t2 = ttime()
	# print(234234,hidden.device.type)
	if "privateuseone" not in str(self.device):
	hidden = hidden.squeeze(0).cpu().numpy()
	else:
	hidden = hidden[0]
	if self.is_half == True:
	hidden = hidden.astype("float32")

	f0 = self.decode(hidden, thred=thred)
	# torch.cuda.synchronize()
	t3 = ttime()
	# print("hmvpe:%s\t%s\t%s\t%s"%(t1-t0,t2-t1,t3-t2,t3-t0))
	return f0

	def infer_from_audio_with_pitch(self, audio, thred=0.03, f0_min=50, f0_max=1100):
	audio = torch.from_numpy(audio).float().to(self.device).unsqueeze(0)
	mel = self.mel_extractor(audio, center=True)
	hidden = self.mel2hidden(mel)
	hidden = hidden.squeeze(0).cpu().numpy()
	if self.is_half == True:
	hidden = hidden.astype("float32")
	f0 = self.decode(hidden, thred=thred)
	f0[(f0 < f0_min) \| (f0 > f0_max)] = 0
	return f0

	def to_local_average_cents(self, salience, thred=0.05):
	# t0 = ttime()
	center = np.argmax(salience, axis=1) # 帧长#index
	salience = np.pad(salience, ((0, 0), (4, 4))) # 帧长,368
	# t1 = ttime()
	center += 4
	todo_salience = []
	todo_cents_mapping = []
	starts = center - 4
	ends = center + 5
	for idx in range(salience.shape[0]):
	todo_salience.append(salience[:, starts[idx] : ends[idx]][idx])
	todo_cents_mapping.append(self.cents_mapping[starts[idx] : ends[idx]])
	# t2 = ttime()
	todo_salience = np.array(todo_salience) # 帧长，9
	todo_cents_mapping = np.array(todo_cents_mapping) # 帧长，9
	product_sum = np.sum(todo_salience * todo_cents_mapping, 1)
	weight_sum = np.sum(todo_salience, 1) # 帧长
	devided = product_sum / weight_sum # 帧长
	# t3 = ttime()
	maxx = np.max(salience, axis=1) # 帧长
	devided[maxx <= thred] = 0
	# t4 = ttime()
	# print("decode:%s\t%s\t%s\t%s" % (t1 - t0, t2 - t1, t3 - t2, t4 - t3))
	return devided


	if __name__ == "__main__":
	import librosa
	import soundfile as sf

	audio, sampling_rate = sf.read(r"C:\Users\liujing04\Desktop\Z\冬之花clip1.wav")
	if len(audio.shape) > 1:
	audio = librosa.to_mono(audio.transpose(1, 0))
	audio_bak = audio.copy()
	if sampling_rate != 16000:
	audio = librosa.resample(audio, orig_sr=sampling_rate, target_sr=16000)
	model_path = r"D:\BaiduNetdiskDownload\RVC-beta-v2-0727AMD_realtime\rmvpe.pt"
	thred = 0.03 # 0.01
	device = "cuda" if torch.cuda.is_available() else "cpu"
	rmvpe = RMVPE(model_path, is_half=False, device=device)
	t0 = ttime()
	f0 = rmvpe.infer_from_audio(audio, thred=thred)
	# f0 = rmvpe.infer_from_audio(audio, thred=thred)
	# f0 = rmvpe.infer_from_audio(audio, thred=thred)
	# f0 = rmvpe.infer_from_audio(audio, thred=thred)
	# f0 = rmvpe.infer_from_audio(audio, thred=thred)
	t1 = ttime()
	logger.info("%s %.2f", f0.shape, t1 - t0)