import gradio as gr import numpy as np import soundfile as sf import noisereduce as nr import spaces import torch import torchaudio import librosa import yaml import tempfile import os from huggingface_hub import hf_hub_download from transformers import AutoFeatureExtractor, WhisperModel from torch.nn.utils import parametrizations from scipy.signal import butter, lfilter from modules.commons import build_model, load_checkpoint, recursive_munch from modules.campplus.DTDNN import CAMPPlus from modules.bigvgan import bigvgan from modules.rmvpe import RMVPE from modules.audio import mel_spectrogram # ---------------------------- # Optimization Settings # ---------------------------- # Set the number of threads to the number of CPU cores torch.set_num_threads(os.cpu_count()) # Enable optimized backends torch.backends.openmp.enabled = True torch.backends.mkldnn.enabled = True torch.backends.cudnn.enabled = False torch.backends.cuda.enabled = False torch.set_grad_enabled(False) # Force CPU usage device = torch.device("cpu") print(f"[DEVICE] | Using device: {device}") # ---------------------------- # Load Models and Configuration # ---------------------------- def load_custom_model_from_hf(repo_id, model_filename="pytorch_model.bin", config_filename="config.yml"): os.makedirs("./checkpoints", exist_ok=True) model_path = hf_hub_download(repo_id=repo_id, filename=model_filename, cache_dir="./checkpoints") if config_filename is None: return model_path config_path = hf_hub_download(repo_id=repo_id, filename=config_filename, cache_dir="./checkpoints") return model_path, config_path # Load DiT model dit_checkpoint_path, dit_config_path = load_custom_model_from_hf("Plachta/Seed-VC", "DiT_seed_v2_uvit_whisper_small_wavenet_bigvgan_pruned.pth", "config_dit_mel_seed_uvit_whisper_small_wavenet.yml") config = yaml.safe_load(open(dit_config_path, 'r')) model_params = recursive_munch(config['model_params']) model = build_model(model_params, stage='DiT') # Debug: Print model keys to identify correct key print(f"[INFO] | Model keys: {model.keys()}") hop_length = config['preprocess_params']['spect_params']['hop_length'] sr = config['preprocess_params']['sr'] # Load DiT checkpoints model, _, _, _ = load_checkpoint(model, None, dit_checkpoint_path, load_only_params=True, ignore_modules=[], is_distributed=False) for key in model: model[key].eval() model[key].to(device) print("[INFO] | DiT model loaded and set to eval mode.") model.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192) # Ensure 'CAMPPlus' is correctly imported and defined try: campplus_model = CAMPPlus(feat_dim=80, embedding_size=192) print("[INFO] | CAMPPlus model instantiated.") except NameError: print("[ERROR] | CAMPPlus is not defined. Please check the import path and ensure CAMPPlus is correctly defined.") raise # Set weights_only=True for security campplus_ckpt_path = load_custom_model_from_hf("funasr/campplus", "campplus_cn_common.bin", config_filename=None) campplus_state = torch.load(campplus_ckpt_path, map_location="cpu", weights_only=True) campplus_model.load_state_dict(campplus_state) campplus_model.eval() campplus_model.to(device) print("[INFO] | CAMPPlus model loaded, set to eval mode, and moved to CPU.") # Load BigVGAN model bigvgan_model = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_22khz_80band_256x', use_cuda_kernel=False) bigvgan_model.remove_weight_norm() bigvgan_model = bigvgan_model.eval().to(device) print("[INFO] | BigVGAN model loaded, weight norm removed, set to eval mode, and moved to CPU.") # Load FAcodec model ckpt_path, config_path = load_custom_model_from_hf("Plachta/FAcodec", 'pytorch_model.bin', 'config.yml') codec_config = yaml.safe_load(open(config_path)) codec_model_params = recursive_munch(codec_config['model_params']) codec_encoder = build_model(codec_model_params, stage="codec") ckpt_params = torch.load(ckpt_path, map_location="cpu", weights_only=True) for key in codec_encoder: codec_encoder[key].load_state_dict(ckpt_params[key], strict=False) codec_encoder = {k: v.eval().to(device) for k, v in codec_encoder.items()} print("[INFO] | FAcodec model loaded, set to eval mode, and moved to CPU.") # Load Whisper model with float32 and compatible size whisper_name = model_params.speech_tokenizer.whisper_name if hasattr(model_params.speech_tokenizer, 'whisper_name') else "openai/whisper-small" whisper_model = WhisperModel.from_pretrained(whisper_name, torch_dtype=torch.float32).to(device) del whisper_model.decoder # Remove decoder as it's not used whisper_feature_extractor = AutoFeatureExtractor.from_pretrained(whisper_name) print(f"[INFO] | Whisper model '{whisper_name}' loaded with dtype {whisper_model.dtype} and moved to CPU.") # Generate mel spectrograms with optimized parameters mel_fn_args = { "n_fft": 1024, "win_size": 1024, "hop_size": 256, "num_mels": 80, "sampling_rate": sr, "fmin": 0, "fmax": None, "center": False } to_mel = lambda x: mel_spectrogram(x, **mel_fn_args) # Load F0 conditioned model dit_checkpoint_path_f0, dit_config_path_f0 = load_custom_model_from_hf("Plachta/Seed-VC", "DiT_seed_v2_uvit_whisper_base_f0_44k_bigvgan_pruned_ft_ema.pth", "config_dit_mel_seed_uvit_whisper_base_f0_44k.yml") config_f0 = yaml.safe_load(open(dit_config_path_f0, 'r')) model_params_f0 = recursive_munch(config_f0['model_params']) model_f0 = build_model(model_params_f0, stage='DiT') hop_length_f0 = config_f0['preprocess_params']['spect_params']['hop_length'] sr_f0 = config_f0['preprocess_params']['sr'] # Load F0 model checkpoints model_f0, _, _, _ = load_checkpoint(model_f0, None, dit_checkpoint_path_f0, load_only_params=True, ignore_modules=[], is_distributed=False) for key in model_f0: model_f0[key].eval() model_f0[key].to(device) print("[INFO] | F0 conditioned DiT model loaded and set to eval mode.") model_f0.cfm.estimator.setup_caches(max_batch_size=1, max_seq_length=8192) # Load F0 extractor model_path = load_custom_model_from_hf("lj1995/VoiceConversionWebUI", "rmvpe.pt", None) rmvpe = RMVPE(model_path, is_half=False, device=device) print("[INFO] | RMVPE model loaded and moved to CPU.") mel_fn_args_f0 = { "n_fft": config_f0['preprocess_params']['spect_params']['n_fft'], "win_size": config_f0['preprocess_params']['spect_params']['win_length'], "hop_size": config_f0['preprocess_params']['spect_params']['hop_length'], "num_mels": 80, # Ensure this matches the primary model "sampling_rate": sr_f0, "fmin": 0, "fmax": None, "center": False } to_mel_f0 = lambda x: mel_spectrogram(x, **mel_fn_args_f0) # Load BigVGAN 44kHz model bigvgan_44k_model = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_44khz_128band_512x', use_cuda_kernel=False) bigvgan_44k_model.remove_weight_norm() bigvgan_44k_model = bigvgan_44k_model.eval().to(device) print("[INFO] | BigVGAN 44kHz model loaded, weight norm removed, set to eval mode, and moved to CPU.") # CSS Styling css = ''' .gradio-container{max-width: 560px !important} h1{text-align:center} footer { visibility: hidden } ''' # ---------------------------- # Functions # ---------------------------- @torch.no_grad() @torch.inference_mode() def voice_conversion(input, reference, steps, guidance, speed): print("[INFO] | Voice conversion started.") inference_module, mel_fn, bigvgan_fn = model, to_mel, bigvgan_model bitrate, sampling_rate, sr_current, hop_length_current = "320k", 16000, 22050, 256 max_context_window, overlap_wave_len = sr_current // hop_length_current * 30, 16 * hop_length_current # Load audio using librosa print("[INFO] | Loading source and reference audio.") source_audio, _ = librosa.load(input, sr=sr_current) ref_audio, _ = librosa.load(reference, sr=sr_current) # Clip reference audio to 25 seconds ref_audio = ref_audio[:sr_current * 25] print(f"[INFO] | Source audio length: {len(source_audio)/sr_current:.2f}s, Reference audio length: {len(ref_audio)/sr_current:.2f}s") # Convert audio to tensors source_audio_tensor = torch.tensor(source_audio).unsqueeze(0).float().to(device) ref_audio_tensor = torch.tensor(ref_audio).unsqueeze(0).float().to(device) # Resample to 16kHz ref_waves_16k = torchaudio.functional.resample(ref_audio_tensor, sr_current, sampling_rate) converted_waves_16k = torchaudio.functional.resample(source_audio_tensor, sr_current, sampling_rate) # Generate Whisper features print("[INFO] | Generating Whisper features for source audio.") if converted_waves_16k.size(-1) <= sampling_rate * 30: alt_inputs = whisper_feature_extractor([converted_waves_16k.squeeze(0).cpu().numpy()], return_tensors="pt", return_attention_mask=True, sampling_rate=sampling_rate) alt_input_features = whisper_model._mask_input_features(alt_inputs.input_features, attention_mask=alt_inputs.attention_mask).to(device) alt_outputs = whisper_model.encoder(alt_input_features.to(torch.float32), head_mask=None, output_attentions=False, output_hidden_states=False, return_dict=True) S_alt = alt_outputs.last_hidden_state.to(torch.float32) S_alt = S_alt[:, :converted_waves_16k.size(-1) // 320 + 1] print(f"[INFO] | S_alt shape: {S_alt.shape}") else: # Process in chunks print("[INFO] | Processing source audio in chunks.") overlapping_time = 5 # seconds chunk_size = sampling_rate * 30 # 30 seconds overlap_size = sampling_rate * overlapping_time S_alt_list = [] buffer = None traversed_time = 0 total_length = converted_waves_16k.size(-1) while traversed_time < total_length: if buffer is None: chunk = converted_waves_16k[:, traversed_time:traversed_time + chunk_size] else: chunk = torch.cat([buffer, converted_waves_16k[:, traversed_time:traversed_time + chunk_size - overlap_size]], dim=-1) alt_inputs = whisper_feature_extractor([chunk.squeeze(0).cpu().numpy()], return_tensors="pt", return_attention_mask=True, sampling_rate=sampling_rate) alt_input_features = whisper_model._mask_input_features(alt_inputs.input_features, attention_mask=alt_inputs.attention_mask).to(device) alt_outputs = whisper_model.encoder(alt_input_features.to(torch.float32), head_mask=None, output_attentions=False, output_hidden_states=False, return_dict=True) S_chunk = alt_outputs.last_hidden_state.to(torch.float32) S_chunk = S_chunk[:, :chunk.size(-1) // 320 + 1] print(f"[INFO] | Processed chunk with S_chunk shape: {S_chunk.shape}") if traversed_time == 0: S_alt_list.append(S_chunk) else: skip_frames = 50 * overlapping_time S_alt_list.append(S_chunk[:, skip_frames:]) buffer = chunk[:, -overlap_size:] traversed_time += chunk_size - overlap_size S_alt = torch.cat(S_alt_list, dim=1) print(f"[INFO] | Final S_alt shape after chunk processing: {S_alt.shape}") # Original Whisper features print("[INFO] | Generating Whisper features for reference audio.") ori_waves_16k = torchaudio.functional.resample(ref_audio_tensor, sr_current, sampling_rate) ori_inputs = whisper_feature_extractor([ori_waves_16k.squeeze(0).cpu().numpy()], return_tensors="pt", return_attention_mask=True, sampling_rate=sampling_rate) ori_input_features = whisper_model._mask_input_features(ori_inputs.input_features, attention_mask=ori_inputs.attention_mask).to(device) ori_outputs = whisper_model.encoder(ori_input_features.to(torch.float32), head_mask=None, output_attentions=False, output_hidden_states=False, return_dict=True) S_ori = ori_outputs.last_hidden_state.to(torch.float32) S_ori = S_ori[:, :ori_waves_16k.size(-1) // 320 + 1] print(f"[INFO] | S_ori shape: {S_ori.shape}") # Generate mel spectrograms print("[INFO] | Generating mel spectrograms.") mel = mel_fn(source_audio_tensor.float()) mel2 = mel_fn(ref_audio_tensor.float()) print(f"[INFO] | Mel spectrogram shapes: mel={mel.shape}, mel2={mel2.shape}") # Length adjustment target_lengths = torch.LongTensor([int(mel.size(2) * speed)]).to(mel.device) target2_lengths = torch.LongTensor([mel2.size(2)]).to(mel2.device) print(f"[INFO] | Target lengths: {target_lengths.item()}, {target2_lengths.item()}") # Extract style features print("[INFO] | Extracting style features from reference audio.") feat2 = torchaudio.compliance.kaldi.fbank(ref_waves_16k, num_mel_bins=80, dither=0, sample_frequency=sampling_rate) feat2 = feat2 - feat2.mean(dim=0, keepdim=True) style2 = campplus_model(feat2.unsqueeze(0)) print(f"[INFO] | Style2 shape: {style2.shape}") # Length Regulation print("[INFO] | Applying length regulation.") cond, _, _, _, _ = inference_module.length_regulator(S_alt, ylens=target_lengths, n_quantizers=3, f0=None) prompt_condition, _, _, _, _ = inference_module.length_regulator(S_ori, ylens=target2_lengths, n_quantizers=3, f0=None) print(f"[INFO] | Cond shape: {cond.shape}, Prompt condition shape: {prompt_condition.shape}") # Initialize variables for audio generation max_source_window = max_context_window - mel2.size(2) processed_frames = 0 generated_wave_chunks = [] print("[INFO] | Starting inference and audio generation.") 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) # Perform inference vc_target = inference_module.cfm.inference(cat_condition, torch.LongTensor([cat_condition.size(1)]).to(mel2.device), mel2, style2, None, steps, inference_cfg_rate=guidance) vc_target = vc_target[:, :, mel2.size(2):] print(f"[INFO] | vc_target shape: {vc_target.shape}") # Generate waveform using BigVGAN vc_wave = bigvgan_fn(vc_target.float())[0] print(f"[INFO] | vc_wave shape: {vc_wave.shape}") # Handle the generated waveform output_wave = vc_wave[0].cpu().numpy() generated_wave_chunks.append(output_wave) # Ensure processed_frames increments correctly to avoid infinite loop processed_frames += vc_target.size(2) print(f"[INFO] | Processed frames updated to: {processed_frames}") # Concatenate all generated wave chunks final_audio = np.concatenate(generated_wave_chunks).astype(np.float32) # Normalize the audio to ensure it's within [-1.0, 1.0] max_val = np.max(np.abs(final_audio)) if max_val > 1.0: final_audio = final_audio / max_val print("[INFO] | Final audio normalized.") # Smoothen the audio to reduce distorted audio def butter_bandpass_filter_filtfilt(data, lowcut=80, highcut=6000, fs=sr_current, order=4): nyq = 0.5 * fs low = lowcut / nyq high = highcut / nyq b, a = butter(order, [low, high], btype='band') y = filtfilt(b, a, data) return y final_audio = butter_bandpass_filter_filtfilt(final_audio) print("[INFO] | Final audio smoothed with low-pass filter.") noise_profile = nr.get_noise_profile(final_audio, sr_current) final_audio = nr.reduce_noise(y=final_audio, sr=sr_current, y_noise=noise_profile, prop_decrease=1.0) print("[INFO] | Final audio noise reduced using noisereduce.") # Save the audio to a temporary WAV file print("[INFO] | Saving final audio to a temporary WAV file.") with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp_file: sf.write(tmp_file.name, final_audio, sr_current, format='WAV') temp_file_path = tmp_file.name print(f"[INFO] | Final audio saved to {temp_file_path}") return temp_file_path def cloud(): print("[CLOUD] | Space maintained.") @spaces.GPU(duration=15) def gpu(): return # ---------------------------- # Gradio Interface # ---------------------------- with gr.Blocks(css=css) as main: with gr.Column(): gr.Markdown("🪄 Add tone to audio.") with gr.Column(): input = gr.Audio(label="Input Audio", type="filepath") reference_input = gr.Audio(label="Reference Audio", type="filepath") with gr.Column(): steps = gr.Slider(label="Steps", value=1, minimum=1, maximum=100, step=1) guidance = gr.Slider(label="Guidance", value=0.7, minimum=0.0, maximum=1.0, step=0.1) speed = gr.Slider(label="Speed", value=1.0, minimum=0.5, maximum=2.0, step=0.1) with gr.Column(): submit = gr.Button("▶") maintain = gr.Button("☁️") with gr.Column(): output = gr.Audio(label="Output", type="filepath") submit.click(voice_conversion, inputs=[input, reference_input, steps, guidance, speed], outputs=output, queue=False) maintain.click(cloud, inputs=[], outputs=[], queue=False) main.launch(show_api=True)