# codec cycle raw audio and save to disk import os import glob import torch import argparse import torchaudio import numpy as np from tqdm import tqdm from typing import Optional from suno_utils.tasks.dac_2c_12cb import DAC EXTS = ["wav", "mp3", "flac", "ogg"] def load_codec(codec_path: str): sd = torch.load(codec_path, map_location="cpu") model = DAC(**sd["metadata"]["kwargs"]) model.load_state_dict(sd["state_dict"]) model.eval() return model def codec_cycle( audio: torch.Tensor, sample_rate: int, dac: torch.nn.Module, n_quantizers: Optional[int] = None, noise_level: Optional[float] = 0.0, ): length = audio.shape[-1] with torch.no_grad(): audio = dac.preprocess(audio, sample_rate) z, codes, latents, commitment_loss, codebook_loss = dac.encode( audio, n_quantizers ) if noise_level > 0.0: p = 1.0 # Adjust this probability as needed z_noise = noise_level * torch.randn((z.shape[-2], z.shape[-1])) bernoulli_mask = torch.rand(z.shape[-1]) < p bernoulli_mask = bernoulli_mask.unsqueeze(0).expand(z.shape[-2], -1) z_noise_with_mask = z_noise * bernoulli_mask.float() z += z_noise_with_mask.type_as(z) audio = dac.decode(z) return audio[..., :length] def block_based_codec_cycle( audio: torch.Tensor, sample_rate: int, model: torch.nn.Module, block_size: int = 524288, ): overlap = block_size // 2 window = torch.hann_window(block_size) in_channels, in_frames = audio.size() num_frames = int(np.ceil((audio.shape[-1] - overlap) / (block_size - overlap))) print("num_frames", num_frames) num_samples = int((num_frames - 1) * (block_size - overlap) + block_size) # encode in blocks with torch.no_grad(): decoded_audio = torch.zeros((2, num_samples)) for frame_idx in tqdm(np.arange(num_frames)): start = frame_idx * overlap end = start + block_size frame_audio = audio[:, start:end] frame_audio = frame_audio.cuda() codec_cycle() if num_frames == 1: # no window when only one frame continue elif frame_idx == 0: # only use last half of window on first frame half_window = window.clone() half_window[:overlap] = 1.0 decoded_frame *= half_window elif (frame_idx + 1) == num_frames: continue else: # use full window decoded_frame *= window decoded_audio[:, start:end] += decoded_frame return decoded_audio[:, :in_frames] if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument( "input_dir", help="Path to directory containing audio to encode." ) parser.add_argument( "output_dir", help="Path to directory to store codec cycled audio." ) parser.add_argument( "--codec_paths", nargs="+", help="List of paths to pretrained codec checkpoints" ) parser.add_argument( "--sample_rate", help="Output audio sample rate.", default=48000 ) parser.add_argument("--batch_size", default=16, type=int) parser.add_argument( "--block_size", help="Block size for block-based inference on long files.", default=4194304, type=int, ) parser.add_argument( "--max_frames", help="Maximum number of frames for each example.", default=524288, type=int, ) args = parser.parse_args() # create output directory os.makedirs(args.output_dir, exist_ok=True) # check codec checkpoints if len(args.codec_paths) < 1: raise RuntimeError(f"No codec checkpoint paths supplied.") # find all audio files filepaths = [] for ext in EXTS: search_path = os.path.join(args.input_dir, f"*.{ext}") filepaths += glob.glob(search_path) print(f"Found {len(filepaths)} audio files in {args.input_dir}.") if len(filepaths) < 1: raise RuntimeError(f"No audio files found in {args.input_dir}") # iterate over each codec for codec_path in args.codec_paths: # load checkpoint model = load_codec(codec_path) codec_name = os.path.basename(codec_path).split(".")[0] print(f"Encoding audio with {codec_name}...") os.makedirs(args.output_dir, exist_ok=True) # model model to gpu model.cuda() filenames = [] in_frames = [] # iterate over all audio files for filepath in tqdm(filepaths): filename = os.path.basename(filepath) # load audio x, sr = torchaudio.load(filepath) if sr != args.sample_rate: x = torchaudio.functional.resample(x, sr, args.sample_rate) if x.shape[-1] < args.max_frames: continue # crop to max_frames start = np.random.randint(0, x.shape[-1] - args.max_frames) end = start + args.max_frames x_crop = x[:, start:end] # randomize the input level if np.random.rand() > 0.33: gain_db = -(np.random.rand() * 24) gain_lin = 10 ** (gain_db / 20.0) x_crop *= gain_lin # randomize the noise level if np.random.rand() > 0.33: noise_level = np.random.rand() * 2 else: noise_level = 0.0 x_crop = x_crop.cuda() in_frames.append(x_crop) filenames.append(filename) # run inference if len(in_frames) == args.batch_size: # encode out_frames = codec_cycle( torch.stack(in_frames), args.sample_rate, model, noise_level=noise_level, ) out_frames = torch.chunk(out_frames, len(filenames), 0) for filename, in_frame, out_frame in zip( filenames, in_frames, out_frames ): # peak normalize out_frame /= out_frame.abs().max().clamp(1e-8) in_frame /= in_frame.abs().max().clamp(1e-8) out_frame = out_frame.squeeze(0) in_frame = in_frame.squeeze(0) # save the input and decoded to disk input_filepath = os.path.join( args.output_dir, filename.replace(".wav", ".input.wav") ) decoded_filepath = os.path.join( args.output_dir, filename.replace(".wav", f".{codec_name}.wav") ) torchaudio.save( decoded_filepath, out_frame.cpu(), args.sample_rate, ) torchaudio.save( input_filepath, in_frame.cpu(), args.sample_rate, ) # reset buffers filenames = [] in_frames = [] out_frames = []