|
import spaces
|
|
import torch
|
|
import librosa
|
|
import torchaudio
|
|
import numpy as np
|
|
from pydub import AudioSegment
|
|
from hf_utils import load_custom_model_from_hf
|
|
|
|
DEFAULT_REPO_ID = "Plachta/Seed-VC"
|
|
DEFAULT_CFM_CHECKPOINT = "v2/cfm_small.pth"
|
|
DEFAULT_AR_CHECKPOINT = "v2/ar_base.pth"
|
|
|
|
DEFAULT_CE_REPO_ID = "Plachta/ASTRAL-quantization"
|
|
DEFAULT_CE_NARROW_CHECKPOINT = "bsq32/bsq32_light.pth"
|
|
DEFAULT_CE_WIDE_CHECKPOINT = "bsq2048/bsq2048_light.pth"
|
|
|
|
DEFAULT_SE_REPO_ID = "funasr/campplus"
|
|
DEFAULT_SE_CHECKPOINT = "campplus_cn_common.bin"
|
|
|
|
class VoiceConversionWrapper(torch.nn.Module):
|
|
def __init__(
|
|
self,
|
|
sr: int,
|
|
hop_size: int,
|
|
mel_fn: callable,
|
|
cfm: torch.nn.Module,
|
|
cfm_length_regulator: torch.nn.Module,
|
|
content_extractor_narrow: torch.nn.Module,
|
|
content_extractor_wide: torch.nn.Module,
|
|
ar_length_regulator: torch.nn.Module,
|
|
ar: torch.nn.Module,
|
|
style_encoder: torch.nn.Module,
|
|
vocoder: torch.nn.Module,
|
|
):
|
|
super(VoiceConversionWrapper, self).__init__()
|
|
self.sr = sr
|
|
self.hop_size = hop_size
|
|
self.mel_fn = mel_fn
|
|
self.cfm = cfm
|
|
self.cfm_length_regulator = cfm_length_regulator
|
|
self.content_extractor_narrow = content_extractor_narrow
|
|
self.content_extractor_wide = content_extractor_wide
|
|
self.vocoder = vocoder
|
|
self.ar_length_regulator = ar_length_regulator
|
|
self.ar = ar
|
|
self.style_encoder = style_encoder
|
|
|
|
self.overlap_frame_len = 16
|
|
self.bitrate = "320k"
|
|
self.compiled_decode_fn = None
|
|
self.dit_compiled = False
|
|
self.dit_max_context_len = 30
|
|
self.compile_len = 87 * self.dit_max_context_len
|
|
|
|
def compile_ar(self):
|
|
"""
|
|
Compile the AR model for inference.
|
|
"""
|
|
self.compiled_decode_fn = torch.compile(
|
|
self.ar.model.forward_generate,
|
|
fullgraph=True,
|
|
backend="inductor" if torch.cuda.is_available() else "aot_eager",
|
|
mode="reduce-overhead" if torch.cuda.is_available() else None,
|
|
)
|
|
|
|
def compile_cfm(self):
|
|
self.cfm.estimator.transformer = torch.compile(
|
|
self.cfm.estimator.transformer,
|
|
fullgraph=True,
|
|
backend="inductor" if torch.cuda.is_available() else "aot_eager",
|
|
mode="reduce-overhead" if torch.cuda.is_available() else None,
|
|
)
|
|
self.dit_compiled = True
|
|
|
|
@staticmethod
|
|
def strip_prefix(state_dict: dict, prefix: str = "module.") -> dict:
|
|
"""
|
|
Strip the prefix from the state_dict keys.
|
|
"""
|
|
new_state_dict = {}
|
|
for k, v in state_dict.items():
|
|
if k.startswith(prefix):
|
|
new_key = k[len(prefix):]
|
|
else:
|
|
new_key = k
|
|
new_state_dict[new_key] = v
|
|
return new_state_dict
|
|
|
|
@staticmethod
|
|
def duration_reduction_func(token_seq, n_gram=1):
|
|
"""
|
|
Args:
|
|
token_seq: (T,)
|
|
Returns:
|
|
reduced_token_seq: (T')
|
|
reduced_token_seq_len: T'
|
|
"""
|
|
n_gram_seq = token_seq.unfold(0, n_gram, 1)
|
|
mask = torch.all(n_gram_seq[1:] != n_gram_seq[:-1], dim=1)
|
|
reduced_token_seq = torch.cat(
|
|
(n_gram_seq[0, :n_gram], n_gram_seq[1:, -1][mask])
|
|
)
|
|
return reduced_token_seq, len(reduced_token_seq)
|
|
|
|
@staticmethod
|
|
def crossfade(chunk1, chunk2, overlap):
|
|
"""Apply crossfade between two audio chunks."""
|
|
fade_out = np.cos(np.linspace(0, np.pi / 2, overlap)) ** 2
|
|
fade_in = np.cos(np.linspace(np.pi / 2, 0, overlap)) ** 2
|
|
if len(chunk2) < overlap:
|
|
chunk2[:overlap] = chunk2[:overlap] * fade_in[:len(chunk2)] + (chunk1[-overlap:] * fade_out)[:len(chunk2)]
|
|
else:
|
|
chunk2[:overlap] = chunk2[:overlap] * fade_in + chunk1[-overlap:] * fade_out
|
|
return chunk2
|
|
|
|
def _stream_wave_chunks(self, vc_wave, processed_frames, vc_mel, overlap_wave_len,
|
|
generated_wave_chunks, previous_chunk, is_last_chunk, stream_output):
|
|
"""
|
|
Helper method to handle streaming wave chunks.
|
|
|
|
Args:
|
|
vc_wave: The current wave chunk
|
|
processed_frames: Number of frames processed so far
|
|
vc_mel: The mel spectrogram
|
|
overlap_wave_len: Length of overlap between chunks
|
|
generated_wave_chunks: List of generated wave chunks
|
|
previous_chunk: Previous wave chunk for crossfading
|
|
is_last_chunk: Whether this is the last chunk
|
|
stream_output: Whether to stream the output
|
|
|
|
Returns:
|
|
Tuple of (processed_frames, previous_chunk, should_break, mp3_bytes, full_audio)
|
|
where should_break indicates if processing should stop
|
|
mp3_bytes is the MP3 bytes if streaming, None otherwise
|
|
full_audio is the full audio if this is the last chunk, None otherwise
|
|
"""
|
|
mp3_bytes = None
|
|
full_audio = None
|
|
|
|
if processed_frames == 0:
|
|
if is_last_chunk:
|
|
output_wave = vc_wave[0].cpu().numpy()
|
|
generated_wave_chunks.append(output_wave)
|
|
|
|
if stream_output:
|
|
output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
|
|
mp3_bytes = AudioSegment(
|
|
output_wave_int16.tobytes(), frame_rate=self.sr,
|
|
sample_width=output_wave_int16.dtype.itemsize, channels=1
|
|
).export(format="mp3", bitrate=self.bitrate).read()
|
|
full_audio = (self.sr, np.concatenate(generated_wave_chunks))
|
|
else:
|
|
return processed_frames, previous_chunk, True, None, np.concatenate(generated_wave_chunks)
|
|
|
|
return processed_frames, previous_chunk, True, mp3_bytes, full_audio
|
|
|
|
output_wave = vc_wave[0, :-overlap_wave_len].cpu().numpy()
|
|
generated_wave_chunks.append(output_wave)
|
|
previous_chunk = vc_wave[0, -overlap_wave_len:]
|
|
processed_frames += vc_mel.size(2) - self.overlap_frame_len
|
|
|
|
if stream_output:
|
|
output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
|
|
mp3_bytes = AudioSegment(
|
|
output_wave_int16.tobytes(), frame_rate=self.sr,
|
|
sample_width=output_wave_int16.dtype.itemsize, channels=1
|
|
).export(format="mp3", bitrate=self.bitrate).read()
|
|
|
|
elif is_last_chunk:
|
|
output_wave = self.crossfade(previous_chunk.cpu().numpy(), vc_wave[0].cpu().numpy(), overlap_wave_len)
|
|
generated_wave_chunks.append(output_wave)
|
|
processed_frames += vc_mel.size(2) - self.overlap_frame_len
|
|
|
|
if stream_output:
|
|
output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
|
|
mp3_bytes = AudioSegment(
|
|
output_wave_int16.tobytes(), frame_rate=self.sr,
|
|
sample_width=output_wave_int16.dtype.itemsize, channels=1
|
|
).export(format="mp3", bitrate=self.bitrate).read()
|
|
full_audio = (self.sr, np.concatenate(generated_wave_chunks))
|
|
else:
|
|
return processed_frames, previous_chunk, True, None, np.concatenate(generated_wave_chunks)
|
|
|
|
return processed_frames, previous_chunk, True, mp3_bytes, full_audio
|
|
|
|
else:
|
|
output_wave = self.crossfade(previous_chunk.cpu().numpy(), vc_wave[0, :-overlap_wave_len].cpu().numpy(), overlap_wave_len)
|
|
generated_wave_chunks.append(output_wave)
|
|
previous_chunk = vc_wave[0, -overlap_wave_len:]
|
|
processed_frames += vc_mel.size(2) - self.overlap_frame_len
|
|
|
|
if stream_output:
|
|
output_wave_int16 = (output_wave * 32768.0).astype(np.int16)
|
|
mp3_bytes = AudioSegment(
|
|
output_wave_int16.tobytes(), frame_rate=self.sr,
|
|
sample_width=output_wave_int16.dtype.itemsize, channels=1
|
|
).export(format="mp3", bitrate=self.bitrate).read()
|
|
|
|
return processed_frames, previous_chunk, False, mp3_bytes, full_audio
|
|
|
|
def load_checkpoints(
|
|
self,
|
|
cfm_checkpoint_path = None,
|
|
ar_checkpoint_path = None,
|
|
):
|
|
if cfm_checkpoint_path is None:
|
|
cfm_checkpoint_path = load_custom_model_from_hf(
|
|
repo_id=DEFAULT_REPO_ID,
|
|
model_filename=DEFAULT_CFM_CHECKPOINT,
|
|
)
|
|
if ar_checkpoint_path is None:
|
|
ar_checkpoint_path = load_custom_model_from_hf(
|
|
repo_id=DEFAULT_REPO_ID,
|
|
model_filename=DEFAULT_AR_CHECKPOINT,
|
|
)
|
|
|
|
cfm_checkpoint = torch.load(cfm_checkpoint_path, map_location="cpu")
|
|
cfm_length_regulator_state_dict = self.strip_prefix(cfm_checkpoint["net"]['length_regulator'], "module.")
|
|
cfm_state_dict = self.strip_prefix(cfm_checkpoint["net"]['cfm'], "module.")
|
|
self.cfm.load_state_dict(cfm_state_dict, strict=False)
|
|
self.cfm_length_regulator.load_state_dict(cfm_length_regulator_state_dict, strict=False)
|
|
|
|
|
|
ar_checkpoint = torch.load(ar_checkpoint_path, map_location="cpu")
|
|
ar_length_regulator_state_dict = self.strip_prefix(ar_checkpoint["net"]['length_regulator'], "module.")
|
|
ar_state_dict = self.strip_prefix(ar_checkpoint["net"]['ar'], "module.")
|
|
self.ar.load_state_dict(ar_state_dict, strict=False)
|
|
self.ar_length_regulator.load_state_dict(ar_length_regulator_state_dict, strict=False)
|
|
|
|
|
|
content_extractor_narrow_checkpoint_path = load_custom_model_from_hf(
|
|
repo_id=DEFAULT_CE_REPO_ID,
|
|
model_filename=DEFAULT_CE_NARROW_CHECKPOINT,
|
|
)
|
|
content_extractor_narrow_checkpoint = torch.load(content_extractor_narrow_checkpoint_path, map_location="cpu")
|
|
self.content_extractor_narrow.load_state_dict(
|
|
content_extractor_narrow_checkpoint, strict=False
|
|
)
|
|
|
|
content_extractor_wide_checkpoint_path = load_custom_model_from_hf(
|
|
repo_id=DEFAULT_CE_REPO_ID,
|
|
model_filename=DEFAULT_CE_WIDE_CHECKPOINT,
|
|
)
|
|
content_extractor_wide_checkpoint = torch.load(content_extractor_wide_checkpoint_path, map_location="cpu")
|
|
self.content_extractor_wide.load_state_dict(
|
|
content_extractor_wide_checkpoint, strict=False
|
|
)
|
|
|
|
|
|
style_encoder_checkpoint_path = load_custom_model_from_hf(DEFAULT_SE_REPO_ID, DEFAULT_SE_CHECKPOINT, config_filename=None)
|
|
style_encoder_checkpoint = torch.load(style_encoder_checkpoint_path, map_location="cpu")
|
|
self.style_encoder.load_state_dict(style_encoder_checkpoint, strict=False)
|
|
|
|
def setup_ar_caches(self, max_batch_size=1, max_seq_len=4096, dtype=torch.float32, device=torch.device("cpu")):
|
|
self.ar.setup_caches(max_batch_size=max_batch_size, max_seq_len=max_seq_len, dtype=dtype, device=device)
|
|
|
|
def compute_style(self, waves_16k: torch.Tensor):
|
|
feat = torchaudio.compliance.kaldi.fbank(waves_16k,
|
|
num_mel_bins=80,
|
|
dither=0,
|
|
sample_frequency=16000)
|
|
feat = feat - feat.mean(dim=0, keepdim=True)
|
|
style = self.style_encoder(feat.unsqueeze(0))
|
|
return style
|
|
|
|
@torch.no_grad()
|
|
@torch.inference_mode()
|
|
def convert_timbre(
|
|
self,
|
|
source_audio_path: str,
|
|
target_audio_path: str,
|
|
diffusion_steps: int = 30,
|
|
length_adjust: float = 1.0,
|
|
inference_cfg_rate: float = 0.5,
|
|
use_sway_sampling: bool = False,
|
|
use_amo_sampling: bool = False,
|
|
device: torch.device = torch.device("cpu"),
|
|
dtype: torch.dtype = torch.float32,
|
|
):
|
|
source_wave = librosa.load(source_audio_path, sr=self.sr)[0]
|
|
target_wave = librosa.load(target_audio_path, sr=self.sr)[0]
|
|
source_wave_tensor = torch.tensor(source_wave).unsqueeze(0).to(device)
|
|
target_wave_tensor = torch.tensor(target_wave).unsqueeze(0).to(device)
|
|
|
|
|
|
source_wave_16k = librosa.resample(source_wave, orig_sr=self.sr, target_sr=16000)
|
|
target_wave_16k = librosa.resample(target_wave, orig_sr=self.sr, target_sr=16000)
|
|
source_wave_16k_tensor = torch.tensor(source_wave_16k).unsqueeze(0).to(device)
|
|
target_wave_16k_tensor = torch.tensor(target_wave_16k).unsqueeze(0).to(device)
|
|
|
|
|
|
source_mel = self.mel_fn(source_wave_tensor)
|
|
target_mel = self.mel_fn(target_wave_tensor)
|
|
source_mel_len = source_mel.size(2)
|
|
target_mel_len = target_mel.size(2)
|
|
|
|
with torch.autocast(device_type=device.type, dtype=dtype):
|
|
|
|
_, source_content_indices, _ = self.content_extractor_wide(source_wave_16k_tensor, [source_wave_16k.size])
|
|
_, target_content_indices, _ = self.content_extractor_wide(target_wave_16k_tensor, [target_wave_16k.size])
|
|
|
|
|
|
target_style = self.compute_style(target_wave_16k_tensor)
|
|
|
|
|
|
cond, _ = self.cfm_length_regulator(source_content_indices, ylens=torch.LongTensor([source_mel_len]).to(device))
|
|
prompt_condition, _, = self.cfm_length_regulator(target_content_indices, ylens=torch.LongTensor([target_mel_len]).to(device))
|
|
|
|
cat_condition = torch.cat([prompt_condition, cond], dim=1)
|
|
|
|
vc_mel = self.cfm.inference(
|
|
cat_condition,
|
|
torch.LongTensor([cat_condition.size(1)]).to(device),
|
|
target_mel, target_style, diffusion_steps,
|
|
inference_cfg_rate=inference_cfg_rate,
|
|
sway_sampling=use_sway_sampling,
|
|
amo_sampling=use_amo_sampling,
|
|
)
|
|
vc_mel = vc_mel[:, :, target_mel_len:]
|
|
vc_wave = self.vocoder(vc_mel.float()).squeeze()[None]
|
|
return vc_wave.cpu().numpy()
|
|
|
|
@torch.no_grad()
|
|
@torch.inference_mode()
|
|
def convert_voice(
|
|
self,
|
|
source_audio_path: str,
|
|
target_audio_path: str,
|
|
diffusion_steps: int = 30,
|
|
length_adjust: float = 1.0,
|
|
inference_cfg_rate: float = 0.5,
|
|
top_p: float = 0.7,
|
|
temperature: float = 0.7,
|
|
repetition_penalty: float = 1.5,
|
|
use_sway_sampling: bool = False,
|
|
use_amo_sampling: bool = False,
|
|
device: torch.device = torch.device("cpu"),
|
|
dtype: torch.dtype = torch.float32,
|
|
):
|
|
source_wave = librosa.load(source_audio_path, sr=self.sr)[0]
|
|
target_wave = librosa.load(target_audio_path, sr=self.sr)[0]
|
|
source_wave_tensor = torch.tensor(source_wave).unsqueeze(0).to(device)
|
|
target_wave_tensor = torch.tensor(target_wave).unsqueeze(0).to(device)
|
|
|
|
|
|
source_wave_16k = librosa.resample(source_wave, orig_sr=self.sr, target_sr=16000)
|
|
target_wave_16k = librosa.resample(target_wave, orig_sr=self.sr, target_sr=16000)
|
|
source_wave_16k_tensor = torch.tensor(source_wave_16k).unsqueeze(0).to(device)
|
|
target_wave_16k_tensor = torch.tensor(target_wave_16k).unsqueeze(0).to(device)
|
|
|
|
|
|
source_mel = self.mel_fn(source_wave_tensor)
|
|
target_mel = self.mel_fn(target_wave_tensor)
|
|
source_mel_len = source_mel.size(2)
|
|
target_mel_len = target_mel.size(2)
|
|
|
|
with torch.autocast(device_type=device.type, dtype=dtype):
|
|
|
|
_, source_content_indices, _ = self.content_extractor_wide(source_wave_16k_tensor, [source_wave_16k.size])
|
|
_, target_content_indices, _ = self.content_extractor_wide(target_wave_16k_tensor, [target_wave_16k.size])
|
|
|
|
_, source_narrow_indices, _ = self.content_extractor_narrow(source_wave_16k_tensor,
|
|
[source_wave_16k.size], ssl_model=self.content_extractor_wide.ssl_model)
|
|
_, target_narrow_indices, _ = self.content_extractor_narrow(target_wave_16k_tensor,
|
|
[target_wave_16k.size], ssl_model=self.content_extractor_wide.ssl_model)
|
|
|
|
src_narrow_reduced, src_narrow_len = self.duration_reduction_func(source_narrow_indices[0], 1)
|
|
tgt_narrow_reduced, tgt_narrow_len = self.duration_reduction_func(target_narrow_indices[0], 1)
|
|
|
|
ar_cond = self.ar_length_regulator(torch.cat([tgt_narrow_reduced, src_narrow_reduced], dim=0)[None])[0]
|
|
|
|
ar_out = self.ar.generate(ar_cond, target_content_indices, top_p=top_p, temperature=temperature, repetition_penalty=repetition_penalty)
|
|
ar_out_mel_len = torch.LongTensor([int(source_mel_len / source_content_indices.size(-1) * ar_out.size(-1) * length_adjust)]).to(device)
|
|
|
|
target_style = self.compute_style(target_wave_16k_tensor)
|
|
|
|
|
|
cond, _ = self.cfm_length_regulator(ar_out, ylens=torch.LongTensor([ar_out_mel_len]).to(device))
|
|
prompt_condition, _, = self.cfm_length_regulator(target_content_indices, ylens=torch.LongTensor([target_mel_len]).to(device))
|
|
|
|
cat_condition = torch.cat([prompt_condition, cond], dim=1)
|
|
|
|
vc_mel = self.cfm.inference(
|
|
cat_condition,
|
|
torch.LongTensor([cat_condition.size(1)]).to(device),
|
|
target_mel, target_style, diffusion_steps,
|
|
inference_cfg_rate=inference_cfg_rate,
|
|
sway_sampling=use_sway_sampling,
|
|
amo_sampling=use_amo_sampling,
|
|
)
|
|
vc_mel = vc_mel[:, :, target_mel_len:]
|
|
vc_wave = self.vocoder(vc_mel.float()).squeeze()[None]
|
|
return vc_wave.cpu().numpy()
|
|
|
|
def _process_content_features(self, audio_16k_tensor, is_narrow=False):
|
|
"""Process audio through Whisper model to extract features."""
|
|
content_extractor_fn = self.content_extractor_narrow if is_narrow else self.content_extractor_wide
|
|
if audio_16k_tensor.size(-1) <= 16000 * 30:
|
|
|
|
_, content_indices, _ = content_extractor_fn(audio_16k_tensor, [audio_16k_tensor.size(-1)], ssl_model=self.content_extractor_wide.ssl_model)
|
|
else:
|
|
|
|
overlapping_time = 5
|
|
features_list = []
|
|
buffer = None
|
|
traversed_time = 0
|
|
while traversed_time < audio_16k_tensor.size(-1):
|
|
if buffer is None:
|
|
chunk = audio_16k_tensor[:, traversed_time:traversed_time + 16000 * 30]
|
|
else:
|
|
chunk = torch.cat([
|
|
buffer,
|
|
audio_16k_tensor[:, traversed_time:traversed_time + 16000 * (30 - overlapping_time)]
|
|
], dim=-1)
|
|
_, chunk_content_indices, _ = content_extractor_fn(chunk, [chunk.size(-1)], ssl_model=self.content_extractor_wide.ssl_model)
|
|
if traversed_time == 0:
|
|
features_list.append(chunk_content_indices)
|
|
else:
|
|
features_list.append(chunk_content_indices[:, 50 * overlapping_time:])
|
|
buffer = chunk[:, -16000 * overlapping_time:]
|
|
traversed_time += 30 * 16000 if traversed_time == 0 else chunk.size(-1) - 16000 * overlapping_time
|
|
content_indices = torch.cat(features_list, dim=1)
|
|
|
|
return content_indices
|
|
|
|
@spaces.GPU
|
|
@torch.no_grad()
|
|
@torch.inference_mode()
|
|
def convert_voice_with_streaming(
|
|
self,
|
|
source_audio_path: str,
|
|
target_audio_path: str,
|
|
diffusion_steps: int = 30,
|
|
length_adjust: float = 1.0,
|
|
intelligebility_cfg_rate: float = 0.7,
|
|
similarity_cfg_rate: float = 0.7,
|
|
top_p: float = 0.7,
|
|
temperature: float = 0.7,
|
|
repetition_penalty: float = 1.5,
|
|
convert_style: bool = False,
|
|
anonymization_only: bool = False,
|
|
device: torch.device = torch.device("cuda"),
|
|
dtype: torch.dtype = torch.float16,
|
|
stream_output: bool = True,
|
|
):
|
|
"""
|
|
Convert voice with streaming support for long audio files.
|
|
|
|
Args:
|
|
source_audio_path: Path to source audio file
|
|
target_audio_path: Path to target audio file
|
|
diffusion_steps: Number of diffusion steps (default: 30)
|
|
length_adjust: Length adjustment factor (default: 1.0)
|
|
intelligebility_cfg_rate: CFG rate for intelligibility (default: 0.7)
|
|
similarity_cfg_rate: CFG rate for similarity (default: 0.7)
|
|
top_p: Top-p sampling parameter (default: 0.7)
|
|
temperature: Temperature for sampling (default: 0.7)
|
|
repetition_penalty: Repetition penalty (default: 1.5)
|
|
device: Device to use (default: cpu)
|
|
dtype: Data type to use (default: float32)
|
|
stream_output: Whether to stream the output (default: True)
|
|
|
|
Returns:
|
|
If stream_output is True, yields (mp3_bytes, full_audio) tuples
|
|
If stream_output is False, returns the full audio as a numpy array
|
|
"""
|
|
|
|
source_wave = librosa.load(source_audio_path, sr=self.sr)[0]
|
|
target_wave = librosa.load(target_audio_path, sr=self.sr)[0]
|
|
|
|
|
|
target_wave = target_wave[:self.sr * (self.dit_max_context_len - 5)]
|
|
|
|
source_wave_tensor = torch.tensor(source_wave).unsqueeze(0).float().to(device)
|
|
target_wave_tensor = torch.tensor(target_wave).unsqueeze(0).float().to(device)
|
|
|
|
|
|
source_wave_16k = librosa.resample(source_wave, orig_sr=self.sr, target_sr=16000)
|
|
target_wave_16k = librosa.resample(target_wave, orig_sr=self.sr, target_sr=16000)
|
|
source_wave_16k_tensor = torch.tensor(source_wave_16k).unsqueeze(0).to(device)
|
|
target_wave_16k_tensor = torch.tensor(target_wave_16k).unsqueeze(0).to(device)
|
|
|
|
|
|
source_mel = self.mel_fn(source_wave_tensor)
|
|
target_mel = self.mel_fn(target_wave_tensor)
|
|
source_mel_len = source_mel.size(2)
|
|
target_mel_len = target_mel.size(2)
|
|
|
|
|
|
max_context_window = self.sr // self.hop_size * self.dit_max_context_len
|
|
overlap_wave_len = self.overlap_frame_len * self.hop_size
|
|
|
|
with torch.autocast(device_type=device.type, dtype=dtype):
|
|
|
|
source_content_indices = self._process_content_features(source_wave_16k_tensor, is_narrow=False)
|
|
target_content_indices = self._process_content_features(target_wave_16k_tensor, is_narrow=False)
|
|
|
|
target_style = self.compute_style(target_wave_16k_tensor)
|
|
prompt_condition, _, = self.cfm_length_regulator(target_content_indices,
|
|
ylens=torch.LongTensor([target_mel_len]).to(device))
|
|
|
|
|
|
generated_wave_chunks = []
|
|
processed_frames = 0
|
|
previous_chunk = None
|
|
if convert_style:
|
|
with torch.autocast(device_type=device.type, dtype=dtype):
|
|
source_narrow_indices = self._process_content_features(source_wave_16k_tensor, is_narrow=True)
|
|
target_narrow_indices = self._process_content_features(target_wave_16k_tensor, is_narrow=True)
|
|
src_narrow_reduced, src_narrow_len = self.duration_reduction_func(source_narrow_indices[0], 1)
|
|
tgt_narrow_reduced, tgt_narrow_len = self.duration_reduction_func(target_narrow_indices[0], 1)
|
|
|
|
max_chunk_size = 1000
|
|
|
|
|
|
for i in range(0, len(src_narrow_reduced), max_chunk_size):
|
|
is_last_chunk = i + max_chunk_size >= len(src_narrow_reduced)
|
|
with torch.autocast(device_type=device.type, dtype=dtype):
|
|
chunk = src_narrow_reduced[i:i + max_chunk_size]
|
|
if anonymization_only:
|
|
chunk_ar_cond = self.ar_length_regulator(chunk[None])[0]
|
|
chunk_ar_out = self.ar.generate(chunk_ar_cond, torch.zeros([1, 0]).long().to(device),
|
|
compiled_decode_fn=self.compiled_decode_fn,
|
|
top_p=top_p, temperature=temperature,
|
|
repetition_penalty=repetition_penalty)
|
|
else:
|
|
|
|
chunk_ar_cond = self.ar_length_regulator(torch.cat([tgt_narrow_reduced, chunk], dim=0)[None])[0]
|
|
chunk_ar_out = self.ar.generate(chunk_ar_cond, target_content_indices, compiled_decode_fn=self.compiled_decode_fn,
|
|
top_p=top_p, temperature=temperature,
|
|
repetition_penalty=repetition_penalty)
|
|
chunkar_out_mel_len = torch.LongTensor([int(source_mel_len / source_content_indices.size(
|
|
-1) * chunk_ar_out.size(-1) * length_adjust)]).to(device)
|
|
|
|
chunk_cond, _ = self.cfm_length_regulator(chunk_ar_out, ylens=torch.LongTensor([chunkar_out_mel_len]).to(device))
|
|
cat_condition = torch.cat([prompt_condition, chunk_cond], dim=1)
|
|
original_len = cat_condition.size(1)
|
|
|
|
if self.dit_compiled:
|
|
cat_condition = torch.nn.functional.pad(cat_condition,
|
|
(0, 0, 0, self.compile_len - cat_condition.size(1),),
|
|
value=0)
|
|
|
|
vc_mel = self.cfm.inference(
|
|
cat_condition,
|
|
torch.LongTensor([original_len]).to(device),
|
|
target_mel, target_style, diffusion_steps,
|
|
inference_cfg_rate=[intelligebility_cfg_rate, similarity_cfg_rate],
|
|
random_voice=anonymization_only,
|
|
)
|
|
vc_mel = vc_mel[:, :, target_mel_len:original_len]
|
|
vc_wave = self.vocoder(vc_mel).squeeze()[None]
|
|
processed_frames, previous_chunk, should_break, mp3_bytes, full_audio = self._stream_wave_chunks(
|
|
vc_wave, processed_frames, vc_mel, overlap_wave_len,
|
|
generated_wave_chunks, previous_chunk, is_last_chunk, stream_output
|
|
)
|
|
|
|
if stream_output and mp3_bytes is not None:
|
|
yield mp3_bytes, full_audio
|
|
|
|
if should_break:
|
|
if not stream_output:
|
|
return full_audio
|
|
break
|
|
else:
|
|
cond, _ = self.cfm_length_regulator(source_content_indices, ylens=torch.LongTensor([source_mel_len]).to(device))
|
|
|
|
|
|
max_source_window = max_context_window - target_mel.size(2)
|
|
|
|
|
|
while processed_frames < cond.size(1):
|
|
chunk_cond = cond[:, processed_frames:processed_frames + max_source_window]
|
|
is_last_chunk = processed_frames + max_source_window >= cond.size(1)
|
|
cat_condition = torch.cat([prompt_condition, chunk_cond], dim=1)
|
|
original_len = cat_condition.size(1)
|
|
|
|
if self.dit_compiled:
|
|
cat_condition = torch.nn.functional.pad(cat_condition,
|
|
(0, 0, 0, self.compile_len - cat_condition.size(1),), value=0)
|
|
with torch.autocast(device_type=device.type, dtype=dtype):
|
|
|
|
vc_mel = self.cfm.inference(
|
|
cat_condition,
|
|
torch.LongTensor([original_len]).to(device),
|
|
target_mel, target_style, diffusion_steps,
|
|
inference_cfg_rate=[intelligebility_cfg_rate, similarity_cfg_rate],
|
|
random_voice=anonymization_only,
|
|
)
|
|
vc_mel = vc_mel[:, :, target_mel_len:original_len]
|
|
vc_wave = self.vocoder(vc_mel).squeeze()[None]
|
|
|
|
processed_frames, previous_chunk, should_break, mp3_bytes, full_audio = self._stream_wave_chunks(
|
|
vc_wave, processed_frames, vc_mel, overlap_wave_len,
|
|
generated_wave_chunks, previous_chunk, is_last_chunk, stream_output
|
|
)
|
|
|
|
if stream_output and mp3_bytes is not None:
|
|
yield mp3_bytes, full_audio
|
|
|
|
if should_break:
|
|
if not stream_output:
|
|
return full_audio
|
|
break
|
|
|
|
|
|
|