| import torch |
| import torchaudio |
| from pathlib import Path |
| import argparse |
| from tqdm import tqdm |
| from acestep.music_dcae.music_dcae_pipeline import MusicDCAE |
|
|
| class AudioVAE: |
| def __init__(self, device: torch.device): |
| self.model = MusicDCAE().to(device) |
| self.model.eval() |
| self.device = device |
| self.latent_mean = torch.tensor( |
| [0.1207, -0.0186, -0.0947, -0.3779, 0.5956, 0.3422, 0.1796, -0.0526], |
| device=device, |
| ).view(1, -1, 1, 1) |
| self.latent_std = torch.tensor( |
| [0.4638, 0.3154, 0.6244, 1.5078, 0.4696, 0.4633, 0.5614, 0.2707], |
| device=device, |
| ).view(1, -1, 1, 1) |
|
|
| def encode(self, audio): |
|
|
| with torch.no_grad(): |
| audio_lengths = torch.tensor([audio.shape[2]] * audio.shape[0]).to(self.device) |
| latents, _ = self.model.encode(audio, audio_lengths, sr=48000) |
| latents = (latents - self.latent_mean) / self.latent_std |
| return latents |
|
|
| def decode(self, latents: torch.Tensor) -> torch.Tensor: |
| with torch.no_grad(): |
| latents = latents * self.latent_std + self.latent_mean |
| _, audio_list = self.model.decode(latents, sr=48000) |
| audio_batch = torch.stack(audio_list).to(self.device) |
| return audio_batch |
|
|
| def load_audio(audio_path, target_sr=48000): |
| """Load and preprocess audio file.""" |
| audio, sr = torchaudio.load(audio_path) |
|
|
| if audio.shape[0] == 1: |
| audio = audio.repeat(2, 1) |
| elif audio.shape[0] > 2: |
| audio = audio[:2] |
|
|
| if sr != target_sr: |
| resampler = torchaudio.transforms.Resample(sr, target_sr) |
| audio = resampler(audio) |
|
|
| return audio |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description='Encode audio files to VAE latents') |
|
|
| parser.add_argument('--audio-dir', type=str, required=True, |
| help='Directory containing audio files') |
| parser.add_argument('--output-dir', type=str, default="latents", |
| help='Directory to save encoded latents') |
|
|
| args = parser.parse_args() |
|
|
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| print(f"Using device: {device}") |
|
|
| output_dir = Path(args.output_dir) |
| output_dir.mkdir(parents=True, exist_ok=True) |
|
|
| audio_dir = Path(args.audio_dir) |
| audio_extensions = ['*.mp3', '*.wav', '*.flac', '*.ogg', '*.m4a'] |
| audio_files = [] |
| for ext in audio_extensions: |
| audio_files.extend(list(audio_dir.glob(ext))) |
| audio_files = sorted(audio_files) |
|
|
| if len(audio_files) == 0: |
| raise ValueError(f"No audio files found in {args.audio_dir}") |
|
|
| print(f"Found {len(audio_files)} audio files") |
|
|
| vae = AudioVAE(device) |
| print("VAE loaded") |
|
|
| |
| print("\nEncoding audio files...") |
| for audio_path in tqdm(audio_files, desc="Encoding"): |
| try: |
| audio = load_audio(audio_path) |
| audio = audio.unsqueeze(0).to(device) |
| latents = vae.encode(audio) |
| latents = latents.squeeze(0) |
|
|
| output_path = output_dir / f"{audio_path.stem}.pt" |
| torch.save(latents.cpu(), output_path) |
|
|
| except Exception as e: |
| print(f"\nError encoding {audio_path.name}: {e}") |
| continue |
|
|
| print(f"\nEncoding complete! Saved {len(list(output_dir.glob('*.pt')))} latent files to {output_dir}") |
|
|
| if __name__ == '__main__': |
| main() |
|
|
