import sys sys.path.append('./BigVGAN') import time import torch import torchaudio import argparse from tqdm import tqdm import librosa from BigVGAN import bigvgan from BigVGAN.meldataset import get_mel_spectrogram from model import OptimizedAudioRestorationModel # Set the device handle macbooks with M1 chip device = 'cuda' if torch.cuda.is_available() else 'cpu' # Initialize BigVGAN model bigvgan_model = bigvgan.BigVGAN.from_pretrained( 'nvidia/bigvgan_v2_24khz_100band_256x', use_cuda_kernel=False, force_download=False ).to(device) bigvgan_model.remove_weight_norm() def measure_gpu_memory(): if device == 'cuda': torch.cuda.synchronize() return torch.cuda.max_memory_allocated() / (1024 ** 2) # Convert to MB return 0 def apply_overlap_windowing_waveform(waveform, window_size_samples, overlap): step_size = int(window_size_samples * (1 - overlap)) num_chunks = (waveform.shape[-1] - window_size_samples) // step_size + 1 windows = [] for i in range(num_chunks): start_idx = i * step_size end_idx = start_idx + window_size_samples chunk = waveform[..., start_idx:end_idx] windows.append(chunk) return torch.stack(windows) def reconstruct_waveform_from_windows(windows, window_size_samples, overlap): step_size = int(window_size_samples * (1 - overlap)) shape = windows.shape if len(shape) == 2: # windows.shape == (num_windows, window_len) num_windows, window_len = shape channels = 1 windows = windows.unsqueeze(1) # Now windows.shape == (num_windows, 1, window_len) elif len(shape) == 3: num_windows, channels, window_len = shape else: raise ValueError(f"Unexpected windows.shape: {windows.shape}") output_length = (num_windows - 1) * step_size + window_size_samples reconstructed = torch.zeros((channels, output_length)) window_sums = torch.zeros((channels, output_length)) for i in range(num_windows): start_idx = i * step_size end_idx = start_idx + window_len reconstructed[:, start_idx:end_idx] += windows[i] window_sums[:, start_idx:end_idx] += 1 reconstructed = reconstructed / window_sums.clamp(min=1e-6) if channels == 1: reconstructed = reconstructed.squeeze(0) # Remove channel dimension if single channel return reconstructed def load_model(save_path): """ Load the optimized audio restoration model. Parameters: - save_path: Path to the checkpoint file. """ optimized_model = OptimizedAudioRestorationModel(device=device, bigvgan_model=bigvgan_model) state_dict = torch.load(save_path, map_location=device) if 'model_state_dict' in state_dict: state_dict = state_dict['model_state_dict'] optimized_model.voice_restore.load_state_dict(state_dict, strict=True) return optimized_model def restore_audio(model, input_path, output_path, steps=16, cfg_strength=0.5, window_size_sec=5.0, overlap=0.5): # Load the audio file start_time = time.time() initial_gpu_memory = measure_gpu_memory() wav, sr = librosa.load(input_path, sr=24000, mono=True) wav = torch.FloatTensor(wav).unsqueeze(0) # Shape: [1, num_samples] window_size_samples = int(window_size_sec * sr) step_size = int(window_size_samples * (1 - overlap)) # Apply overlapping windowing to the waveform wav_windows = apply_overlap_windowing_waveform(wav, window_size_samples, overlap) restored_wav_windows = [] for wav_window in tqdm(wav_windows): wav_window = wav_window.to(device) # Shape: [1, window_size_samples] # Convert to Mel-spectrogram processed_mel = get_mel_spectrogram(wav_window, bigvgan_model.h).to(device) # Restore audio with torch.no_grad(): with torch.autocast(device): restored_mel = model.voice_restore.sample(processed_mel.transpose(1, 2), steps=steps, cfg_strength=cfg_strength) restored_mel = restored_mel.squeeze(0).transpose(0, 1) # Convert restored mel-spectrogram to waveform with torch.no_grad(): with torch.autocast(device): restored_wav = bigvgan_model(restored_mel.unsqueeze(0)).squeeze(0).float().cpu() # Shape: [num_samples] # Debug: Print shapes # print(f"restored_wav.shape: {restored_wav.shape}") restored_wav_windows.append(restored_wav) del wav_window, processed_mel, restored_mel, restored_wav torch.cuda.empty_cache() restored_wav_windows = torch.stack(restored_wav_windows) # Shape: [num_windows, num_samples] # Debug: Print shapes # print(f"restored_wav_windows.shape: {restored_wav_windows.shape}") # Reconstruct the full waveform from the processed windows restored_wav = reconstruct_waveform_from_windows(restored_wav_windows, window_size_samples, overlap) # Ensure the restored_wav has correct dimensions for saving if restored_wav.dim() == 1: restored_wav = restored_wav.unsqueeze(0) # Shape: [1, num_samples] # Save the restored audio torchaudio.save(output_path, restored_wav, 24000) end_time = time.time() total_time = end_time - start_time peak_gpu_memory = measure_gpu_memory() gpu_memory_used = peak_gpu_memory - initial_gpu_memory print(f"Total inference time: {total_time:.2f} seconds") print(f"Peak GPU memory usage: {peak_gpu_memory:.2f} MB") print(f"GPU memory used: {gpu_memory_used:.2f} MB") if __name__ == "__main__": # Argument parser setup parser = argparse.ArgumentParser(description="Audio restoration using OptimizedAudioRestorationModel for long-form audio.") parser.add_argument('--checkpoint', type=str, required=True, help="Path to the checkpoint file") parser.add_argument('--input', type=str, required=True, help="Path to the input audio file") parser.add_argument('--output', type=str, required=True, help="Path to save the restored audio file") parser.add_argument('--steps', type=int, default=16, help="Number of sampling steps") parser.add_argument('--cfg_strength', type=float, default=0.5, help="CFG strength value") parser.add_argument('--window_size_sec', type=float, default=5.0, help="Window size in seconds for overlapping") parser.add_argument('--overlap', type=float, default=0.5, help="Overlap ratio for windowing") # Parse arguments args = parser.parse_args() # Load the optimized model optimized_model = load_model(args.checkpoint) if device == 'cuda': optimized_model.bfloat16() optimized_model.eval() optimized_model.to(device) # Use the model to restore audio restore_audio( optimized_model, args.input, args.output, steps=args.steps, cfg_strength=args.cfg_strength, window_size_sec=args.window_size_sec, overlap=args.overlap )