Spaces:
Paused
Paused
| import sys | |
| sys.path.append('./BigVGAN') | |
| import torch | |
| import torch.nn as nn | |
| import torchaudio | |
| import argparse | |
| from BigVGAN import bigvgan | |
| from BigVGAN.meldataset import get_mel_spectrogram | |
| from model import OptimizedAudioRestorationModel | |
| device = 'cuda' if torch.cuda.is_available() else 'cpu' | |
| # If running on non-windows system, you can try using cuda kernel for faster processing `use_cuda_kernel=True` | |
| bigvgan_model = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_24khz_100band_256x', use_cuda_kernel=False).to(device) | |
| bigvgan_model.remove_weight_norm() | |
| example_input = torch.randn(1, 16000) # Example input waveform | |
| example_spec = get_mel_spectrogram(example_input, bigvgan_model.h) | |
| def load_model(save_path): | |
| """ | |
| Load the model. | |
| Parameters: | |
| - save_path: The file path where the optimized model is saved. | |
| """ | |
| optimized_model = OptimizedAudioRestorationModel(device=device, bigvgan_model=bigvgan_model) | |
| state_dict = torch.load(save_path, map_location=torch.device(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): | |
| audio, sr = torchaudio.load(input_path) | |
| if sr != model.target_sample_rate: | |
| audio = torchaudio.functional.resample(audio, sr, model.target_sample_rate) | |
| audio = audio.mean(dim=0, keepdim=True) if audio.dim() > 1 else audio # Convert to mono if stereo | |
| with torch.inference_mode(): | |
| with torch.autocast(device): | |
| restored_wav = model(audio, steps=steps, cfg_strength=cfg_strength) | |
| restored_wav = restored_wav.squeeze(0).float().cpu() # Move to CPU after processing | |
| torchaudio.save(output_path, restored_wav, model.target_sample_rate) | |
| if __name__ == "__main__": | |
| # Argument parser setup | |
| parser = argparse.ArgumentParser(description="Audio restoration using OptimizedAudioRestorationModel") | |
| 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") | |
| # Parse arguments | |
| args = parser.parse_args() | |
| # Load the optimized model | |
| optimized_model = load_model(args.checkpoint) | |
| 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) | |