change to fsq

README.md

@@ -18,7 +18,7 @@ This repository provides an implementation of the MiniMax-Speech model, featurin
 ## Architecture
 
 ### Stage 1: Audio to Discrete Tokens
-Converts raw audio into discrete representations using the DAC (Descript Audio Codec) framework.
+Converts raw audio into discrete representations using the FSQ (S3Tokenizer) framework.
 
 ### Stage 2: Discrete Tokens to Continuous Latent Space
 Maps discrete tokens to a continuous latent space using a Variational Autoencoder (VAE).
@@ -29,25 +29,25 @@ Maps discrete tokens to a continuous latent space using a Variational Autoencode
 
 ### 1. Model Training
 
-#### BPE tokens to DAC codec tokens
-- Based on the DAC codec
-- Using Auto Regressive to predict the DAC codec tokens with learnable speaker extractor
+#### BPE tokens to FSQ tokens
+- Based on the FSQ
+- Using Auto Regressive to predict the FSQ tokens with learnable speaker extractor
 
-#### DAC codec tokens to DAC-VAE latent
+#### FSQ tokens to DAC-VAE latent
 - Based on Cosyvoice2 flow matching decoder
 - Learns continuous latent representations from discrete tokens
 
 ### 2. Feature Extraction
 
 Before training the main model:
-1. Extract discrete tokens using the trained DAC codec [Descript Audio Codec](https://github.com/descriptinc/descript-audio-codec)
+1. Extract discrete tokens using the trained FSQ [S3Tokenizer](https://github.com/xingchensong/S3Tokenizer)
 2. Generate continuous latent representations using the trained DAC-VAE - the pretrained I provided here: [DAC-VAE](https://drive.google.com/file/d/1iwZhPlcdDwvPjeON3bFAeYarsV4ZtI2E/view?usp=sharing)
 
 ### 3. Two-Stage Training
 
 Train the models sequentially:
-- **Stage 1**: BPE tokens → Discrete DAC codec 
-- **Stage 2**: Discrete DAC codec → DAC-VAE Continuous latent space
+- **Stage 1**: BPE tokens → Discrete FSQ 
+- **Stage 2**: Discrete FSQ → DAC-VAE Continuous latent space
 
 ## Getting Started
 
@@ -59,7 +59,7 @@ pip install -r requirements.txt
 
 ### Training Pipeline
 
-1. **Extracting DAC Codec** (if not using pretrained)
+1. **Extracting FSQ** (if not using pretrained)
    ```bash
    # Add training command
    ```
@@ -88,13 +88,12 @@ minimax-speech/
 ├── configs/
 │   └── dac_vae.yaml
 ├── models/
-│   ├── dac_codec/
+│   ├── fsq/
 │   └── dac_vae/
 ├── cosyvoice/          # Components from CosyVoice2
 │   ├── flow/
 │   ├── transformer/
 │   └── utils/
-├── train_dac_vae.py
 └── README.md
 ```
 
@@ -130,13 +129,13 @@ If you use this code in your research, please cite:
 
 This project follows the licensing terms of its dependencies:
 - CosyVoice2 components: [Check CosyVoice2 License](https://github.com/FunAudioLLM/CosyVoice/blob/main/LICENSE)
-- DAC components: [Apache 2.0 License](https://github.com/descriptinc/descript-audio-codec/blob/main/LICENSE)
+- FSQ components: [Apache 2.0 License](https://github.com/xingchensong/S3Tokenizer/blob/main/LICENSE)
 - Original contributions: [Specify your license here]
 
 ## Acknowledgments
 
 - **[CosyVoice2](https://github.com/FunAudioLLM/CosyVoice)**: This implementation extensively uses code and architectures from CosyVoice2
-- **[Descript Audio Codec](https://github.com/descriptinc/descript-audio-codec)**: For the DAC implementation
+- **[FSQ](https://github.com/xingchensong/S3Tokenizer)**: For the FSQ implementation
 - **MiniMax team**: For the technical report and methodology
 - **FunAudioLLM team**: For the excellent CosyVoice2 codebase
 

dac-codec/conf/1gpu.yml

@@ -1,6 +0,0 @@
-$include:
-  - conf/base.yml
-
-batch_size: 12
-val_batch_size: 12
-num_workers: 4

dac-codec/conf/ablations/baseline.yml

@@ -1,3 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml

dac-codec/conf/ablations/diff-mb.yml

@@ -1,22 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-Discriminator.sample_rate: 44100
-Discriminator.fft_sizes: [2048, 1024, 512]
-Discriminator.bands:
-  - [0.0, 0.05]
-  - [0.05, 0.1]
-  - [0.1, 0.25]
-  - [0.25, 0.5]
-  - [0.5, 1.0]
-
-
-# re-weight lambdas to make up for
-# lost discriminators vs baseline
-lambdas:
-  mel/loss: 15.0
-  adv/feat_loss: 5.0
-  adv/gen_loss: 1.0
-  vq/commitment_loss: 0.25
-  vq/codebook_loss: 1.0

dac-codec/conf/ablations/equal-mb.yml

@@ -1,22 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-Discriminator.sample_rate: 44100
-Discriminator.fft_sizes: [2048, 1024, 512]
-Discriminator.bands:
-  - [0.0, 0.2]
-  - [0.2, 0.4]
-  - [0.4, 0.6]
-  - [0.6, 0.8]
-  - [0.8, 1.0]
-
-
-# re-weight lambdas to make up for
-# lost discriminators vs baseline
-lambdas:
-  mel/loss: 15.0
-  adv/feat_loss: 5.0
-  adv/gen_loss: 1.0
-  vq/commitment_loss: 0.25
-  vq/codebook_loss: 1.0

dac-codec/conf/ablations/no-adv.yml

@@ -1,9 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-lambdas:
-  mel/loss: 1.0
-  waveform/loss: 1.0
-  vq/commitment_loss: 0.25
-  vq/codebook_loss: 1.0

dac-codec/conf/ablations/no-data-balance.yml

@@ -1,22 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-train/build_dataset.folders:
-  speech:
-    - /data/daps/train
-    - /data/vctk
-    - /data/vocalset
-    - /data/read_speech
-    - /data/french_speech
-    - /data/emotional_speech/
-    - /data/common_voice/
-    - /data/german_speech/
-    - /data/russian_speech/
-    - /data/spanish_speech/
-  music:
-    - /data/musdb/train
-    - /data/jamendo
-  general:
-    - /data/audioset/data/unbalanced_train_segments/
-    - /data/audioset/data/balanced_train_segments/

dac-codec/conf/ablations/no-low-hop.yml

@@ -1,18 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-MelSpectrogramLoss.n_mels: [80]
-MelSpectrogramLoss.window_lengths: [512]
-MelSpectrogramLoss.mel_fmin: [0]
-MelSpectrogramLoss.mel_fmax: [null]
-MelSpectrogramLoss.pow: 1.0
-MelSpectrogramLoss.clamp_eps: 1.0e-5
-MelSpectrogramLoss.mag_weight: 0.0
-
-lambdas:
-  mel/loss: 100.0
-  adv/feat_loss: 2.0
-  adv/gen_loss: 1.0
-  vq/commitment_loss: 0.25
-  vq/codebook_loss: 1.0

dac-codec/conf/ablations/no-mb.yml

@@ -1,17 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-Discriminator.sample_rate: 44100
-Discriminator.fft_sizes: [2048, 1024, 512]
-Discriminator.bands:
-  - [0.0, 1.0]
-
-# re-weight lambdas to make up for
-# lost discriminators vs baseline
-lambdas:
-  mel/loss: 15.0
-  adv/feat_loss: 5.0
-  adv/gen_loss: 1.0
-  vq/commitment_loss: 0.25
-  vq/codebook_loss: 1.0

dac-codec/conf/ablations/no-mpd-msd.yml

@@ -1,21 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-Discriminator.sample_rate: 44100
-Discriminator.rates: []
-Discriminator.periods: []
-Discriminator.fft_sizes: [2048, 1024, 512]
-Discriminator.bands:
-  - [0.0, 0.1]
-  - [0.1, 0.25]
-  - [0.25, 0.5]
-  - [0.5, 0.75]
-  - [0.75, 1.0]
-
-lambdas:
-  mel/loss: 15.0
-  adv/feat_loss: 2.66
-  adv/gen_loss: 1.0
-  vq/commitment_loss: 0.25
-  vq/codebook_loss: 1.0

dac-codec/conf/ablations/no-mpd.yml

@@ -1,21 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-Discriminator.sample_rate: 44100
-Discriminator.rates: [1]
-Discriminator.periods: []
-Discriminator.fft_sizes: [2048, 1024, 512]
-Discriminator.bands:
-  - [0.0, 0.1]
-  - [0.1, 0.25]
-  - [0.25, 0.5]
-  - [0.5, 0.75]
-  - [0.75, 1.0]
-
-lambdas:
-  mel/loss: 15.0
-  adv/feat_loss: 2.5
-  adv/gen_loss: 1.0
-  vq/commitment_loss: 0.25
-  vq/codebook_loss: 1.0

dac-codec/conf/ablations/only-speech.yml

@@ -1,22 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-train/build_dataset.folders:
-  speech_fb:
-    - /data/daps/train
-  speech_hq:
-    - /data/vctk
-    - /data/vocalset
-    - /data/read_speech
-    - /data/french_speech
-  speech_uq:
-    - /data/emotional_speech/
-    - /data/common_voice/
-    - /data/german_speech/
-    - /data/russian_speech/
-    - /data/spanish_speech/
-
-val/build_dataset.folders:
-  speech_hq:
-    - /data/daps/val

dac-codec/conf/base.yml

@@ -1,123 +0,0 @@
-# Model setup
-DAC.sample_rate: 44100
-DAC.encoder_dim: 64
-DAC.encoder_rates: [2, 4, 8, 8]
-DAC.decoder_dim: 1536
-DAC.decoder_rates: [8, 8, 4, 2]
-
-# Quantization
-DAC.n_codebooks: 9
-DAC.codebook_size: 1024
-DAC.codebook_dim: 8
-DAC.quantizer_dropout: 1.0
-
-# Discriminator
-Discriminator.sample_rate: 44100
-Discriminator.rates: []
-Discriminator.periods: [2, 3, 5, 7, 11]
-Discriminator.fft_sizes: [2048, 1024, 512]
-Discriminator.bands:
-  - [0.0, 0.1]
-  - [0.1, 0.25]
-  - [0.25, 0.5]
-  - [0.5, 0.75]
-  - [0.75, 1.0]
-
-# Optimization
-AdamW.betas: [0.8, 0.99]
-AdamW.lr: 0.0001
-ExponentialLR.gamma: 0.999996
-
-amp: false
-val_batch_size: 100
-device: cuda
-num_iters: 250000
-save_iters: [10000, 50000, 100000, 200000]
-valid_freq: 1000
-sample_freq: 10000
-num_workers: 32
-val_idx: [0, 1, 2, 3, 4, 5, 6, 7]
-seed: 0
-lambdas:
-  mel/loss: 15.0
-  adv/feat_loss: 2.0
-  adv/gen_loss: 1.0
-  vq/commitment_loss: 0.25
-  vq/codebook_loss: 1.0
-
-VolumeNorm.db: [const, -16]
-
-# Transforms
-build_transform.preprocess:
-  - Identity
-build_transform.augment_prob: 0.0
-build_transform.augment:
-  - Identity
-build_transform.postprocess:
-  - VolumeNorm
-  - RescaleAudio
-  - ShiftPhase
-
-# Loss setup
-MultiScaleSTFTLoss.window_lengths: [2048, 512]
-MelSpectrogramLoss.n_mels: [5, 10, 20, 40, 80, 160, 320]
-MelSpectrogramLoss.window_lengths: [32, 64, 128, 256, 512, 1024, 2048]
-MelSpectrogramLoss.mel_fmin: [0, 0, 0, 0, 0, 0, 0]
-MelSpectrogramLoss.mel_fmax: [null, null, null, null, null, null, null]
-MelSpectrogramLoss.pow: 1.0
-MelSpectrogramLoss.clamp_eps: 1.0e-5
-MelSpectrogramLoss.mag_weight: 0.0
-
-# Data
-batch_size: 72
-train/AudioDataset.duration: 0.38
-train/AudioDataset.n_examples: 10000000
-
-val/AudioDataset.duration: 5.0
-val/build_transform.augment_prob: 1.0
-val/AudioDataset.n_examples: 250
-
-test/AudioDataset.duration: 10.0
-test/build_transform.augment_prob: 1.0
-test/AudioDataset.n_examples: 1000
-
-AudioLoader.shuffle: true
-AudioDataset.without_replacement: true
-
-train/build_dataset.folders:
-  speech_fb:
-    - /data/daps/train
-  speech_hq:
-    - /data/vctk
-    - /data/vocalset
-    - /data/read_speech
-    - /data/french_speech
-  speech_uq:
-    - /data/emotional_speech/
-    - /data/common_voice/
-    - /data/german_speech/
-    - /data/russian_speech/
-    - /data/spanish_speech/
-  music_hq:
-    - /data/musdb/train
-  music_uq:
-    - /data/jamendo
-  general:
-    - /data/audioset/data/unbalanced_train_segments/
-    - /data/audioset/data/balanced_train_segments/
-
-val/build_dataset.folders:
-  speech_hq:
-    - /data/daps/val
-  music_hq:
-    - /data/musdb/test
-  general:
-    - /data/audioset/data/eval_segments/
-
-test/build_dataset.folders:
-  speech_hq:
-    - /data/daps/test
-  music_hq:
-    - /data/musdb/test
-  general:
-    - /data/audioset/data/eval_segments/

dac-codec/conf/downsampling/1024x.yml

@@ -1,16 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-# Model setup
-DAC.sample_rate: 44100
-DAC.encoder_dim: 64
-DAC.encoder_rates: [2, 8, 8, 8]
-DAC.decoder_dim: 1536
-DAC.decoder_rates: [8, 4, 4, 2, 2, 2]
-
-# Quantization
-DAC.n_codebooks: 19
-DAC.codebook_size: 1024
-DAC.codebook_dim: 8
-DAC.quantizer_dropout: 1.0

dac-codec/conf/downsampling/128x.yml

@@ -1,16 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-# Model setup
-DAC.sample_rate: 44100
-DAC.encoder_dim: 64
-DAC.encoder_rates: [2, 4, 4, 4]
-DAC.decoder_dim: 1536
-DAC.decoder_rates: [4, 4, 2, 2, 2, 1]
-
-# Quantization
-DAC.n_codebooks: 2
-DAC.codebook_size: 1024
-DAC.codebook_dim: 8
-DAC.quantizer_dropout: 1.0

dac-codec/conf/downsampling/1536x.yml

@@ -1,16 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-# Model setup
-DAC.sample_rate: 44100
-DAC.encoder_dim: 96
-DAC.encoder_rates: [2, 8, 8, 12]
-DAC.decoder_dim: 1536
-DAC.decoder_rates: [12, 4, 4, 2, 2, 2]
-
-# Quantization
-DAC.n_codebooks: 28
-DAC.codebook_size: 1024
-DAC.codebook_dim: 8
-DAC.quantizer_dropout: 1.0

dac-codec/conf/downsampling/768x.yml

@@ -1,16 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-# Model setup
-DAC.sample_rate: 44100
-DAC.encoder_dim: 64
-DAC.encoder_rates: [2, 6, 8, 8]
-DAC.decoder_dim: 1536
-DAC.decoder_rates: [6, 4, 4, 2, 2, 2]
-
-# Quantization
-DAC.n_codebooks: 14
-DAC.codebook_size: 1024
-DAC.codebook_dim: 8
-DAC.quantizer_dropout: 1.0

dac-codec/conf/final/16khz.yml

@@ -1,123 +0,0 @@
-# Model setup
-DAC.sample_rate: 16000
-DAC.encoder_dim: 64
-DAC.encoder_rates: [2, 4, 5, 8]
-DAC.decoder_dim: 1536
-DAC.decoder_rates: [8, 5, 4, 2]
-
-# Quantization
-DAC.n_codebooks: 12
-DAC.codebook_size: 1024
-DAC.codebook_dim: 8
-DAC.quantizer_dropout: 0.5
-
-# Discriminator
-Discriminator.sample_rate: 16000
-Discriminator.rates: []
-Discriminator.periods: [2, 3, 5, 7, 11]
-Discriminator.fft_sizes: [2048, 1024, 512]
-Discriminator.bands:
-  - [0.0, 0.1]
-  - [0.1, 0.25]
-  - [0.25, 0.5]
-  - [0.5, 0.75]
-  - [0.75, 1.0]
-
-# Optimization
-AdamW.betas: [0.8, 0.99]
-AdamW.lr: 0.0001
-ExponentialLR.gamma: 0.999996
-
-amp: false
-val_batch_size: 100
-device: cuda
-num_iters: 400000
-save_iters: [10000, 50000, 100000, 200000]
-valid_freq: 1000
-sample_freq: 10000
-num_workers: 32
-val_idx: [0, 1, 2, 3, 4, 5, 6, 7]
-seed: 0
-lambdas:
-  mel/loss: 15.0
-  adv/feat_loss: 2.0
-  adv/gen_loss: 1.0
-  vq/commitment_loss: 0.25
-  vq/codebook_loss: 1.0
-
-VolumeNorm.db: [const, -16]
-
-# Transforms
-build_transform.preprocess:
-  - Identity
-build_transform.augment_prob: 0.0
-build_transform.augment:
-  - Identity
-build_transform.postprocess:
-  - VolumeNorm
-  - RescaleAudio
-  - ShiftPhase
-
-# Loss setup
-MultiScaleSTFTLoss.window_lengths: [2048, 512]
-MelSpectrogramLoss.n_mels: [5, 10, 20, 40, 80, 160, 320]
-MelSpectrogramLoss.window_lengths: [32, 64, 128, 256, 512, 1024, 2048]
-MelSpectrogramLoss.mel_fmin: [0, 0, 0, 0, 0, 0, 0]
-MelSpectrogramLoss.mel_fmax: [null, null, null, null, null, null, null]
-MelSpectrogramLoss.pow: 1.0
-MelSpectrogramLoss.clamp_eps: 1.0e-5
-MelSpectrogramLoss.mag_weight: 0.0
-
-# Data
-batch_size: 72
-train/AudioDataset.duration: 0.38
-train/AudioDataset.n_examples: 10000000
-
-val/AudioDataset.duration: 5.0
-val/build_transform.augment_prob: 1.0
-val/AudioDataset.n_examples: 250
-
-test/AudioDataset.duration: 10.0
-test/build_transform.augment_prob: 1.0
-test/AudioDataset.n_examples: 1000
-
-AudioLoader.shuffle: true
-AudioDataset.without_replacement: true
-
-train/build_dataset.folders:
-  speech_fb:
-    - /data/daps/train
-  speech_hq:
-    - /data/vctk
-    - /data/vocalset
-    - /data/read_speech
-    - /data/french_speech
-  speech_uq:
-    - /data/emotional_speech/
-    - /data/common_voice/
-    - /data/german_speech/
-    - /data/russian_speech/
-    - /data/spanish_speech/
-  music_hq:
-    - /data/musdb/train
-  music_uq:
-    - /data/jamendo
-  general:
-    - /data/audioset/data/unbalanced_train_segments/
-    - /data/audioset/data/balanced_train_segments/
-
-val/build_dataset.folders:
-  speech_hq:
-    - /data/daps/val
-  music_hq:
-    - /data/musdb/test
-  general:
-    - /data/audioset/data/eval_segments/
-
-test/build_dataset.folders:
-  speech_hq:
-    - /data/daps/test
-  music_hq:
-    - /data/musdb/test
-  general:
-    - /data/audioset/data/eval_segments/

dac-codec/conf/final/24khz.yml

@@ -1,123 +0,0 @@
-# Model setup
-DAC.sample_rate: 24000
-DAC.encoder_dim: 64
-DAC.encoder_rates: [2, 4, 5, 8]
-DAC.decoder_dim: 1536
-DAC.decoder_rates: [8, 5, 4, 2]
-
-# Quantization
-DAC.n_codebooks: 32
-DAC.codebook_size: 1024
-DAC.codebook_dim: 8
-DAC.quantizer_dropout: 0.5
-
-# Discriminator
-Discriminator.sample_rate: 24000
-Discriminator.rates: []
-Discriminator.periods: [2, 3, 5, 7, 11]
-Discriminator.fft_sizes: [2048, 1024, 512]
-Discriminator.bands:
-  - [0.0, 0.1]
-  - [0.1, 0.25]
-  - [0.25, 0.5]
-  - [0.5, 0.75]
-  - [0.75, 1.0]
-
-# Optimization
-AdamW.betas: [0.8, 0.99]
-AdamW.lr: 0.0001
-ExponentialLR.gamma: 0.999996
-
-amp: false
-val_batch_size: 100
-device: cuda
-num_iters: 400000
-save_iters: [10000, 50000, 100000, 200000]
-valid_freq: 1000
-sample_freq: 10000
-num_workers: 32
-val_idx: [0, 1, 2, 3, 4, 5, 6, 7]
-seed: 0
-lambdas:
-  mel/loss: 15.0
-  adv/feat_loss: 2.0
-  adv/gen_loss: 1.0
-  vq/commitment_loss: 0.25
-  vq/codebook_loss: 1.0
-
-VolumeNorm.db: [const, -16]
-
-# Transforms
-build_transform.preprocess:
-  - Identity
-build_transform.augment_prob: 0.0
-build_transform.augment:
-  - Identity
-build_transform.postprocess:
-  - VolumeNorm
-  - RescaleAudio
-  - ShiftPhase
-
-# Loss setup
-MultiScaleSTFTLoss.window_lengths: [2048, 512]
-MelSpectrogramLoss.n_mels: [5, 10, 20, 40, 80, 160, 320]
-MelSpectrogramLoss.window_lengths: [32, 64, 128, 256, 512, 1024, 2048]
-MelSpectrogramLoss.mel_fmin: [0, 0, 0, 0, 0, 0, 0]
-MelSpectrogramLoss.mel_fmax: [null, null, null, null, null, null, null]
-MelSpectrogramLoss.pow: 1.0
-MelSpectrogramLoss.clamp_eps: 1.0e-5
-MelSpectrogramLoss.mag_weight: 0.0
-
-# Data
-batch_size: 72
-train/AudioDataset.duration: 0.38
-train/AudioDataset.n_examples: 10000000
-
-val/AudioDataset.duration: 5.0
-val/build_transform.augment_prob: 1.0
-val/AudioDataset.n_examples: 250
-
-test/AudioDataset.duration: 10.0
-test/build_transform.augment_prob: 1.0
-test/AudioDataset.n_examples: 1000
-
-AudioLoader.shuffle: true
-AudioDataset.without_replacement: true
-
-train/build_dataset.folders:
-  speech_fb:
-    - /data/daps/train
-  speech_hq:
-    - /data/vctk
-    - /data/vocalset
-    - /data/read_speech
-    - /data/french_speech
-  speech_uq:
-    - /data/emotional_speech/
-    - /data/common_voice/
-    - /data/german_speech/
-    - /data/russian_speech/
-    - /data/spanish_speech/
-  music_hq:
-    - /data/musdb/train
-  music_uq:
-    - /data/jamendo
-  general:
-    - /data/audioset/data/unbalanced_train_segments/
-    - /data/audioset/data/balanced_train_segments/
-
-val/build_dataset.folders:
-  speech_hq:
-    - /data/daps/val
-  music_hq:
-    - /data/musdb/test
-  general:
-    - /data/audioset/data/eval_segments/
-
-test/build_dataset.folders:
-  speech_hq:
-    - /data/daps/test
-  music_hq:
-    - /data/musdb/test
-  general:
-    - /data/audioset/data/eval_segments/

dac-codec/conf/final/44khz-16kbps.yml

@@ -1,124 +0,0 @@
-# Model setup
-DAC.sample_rate: 44100
-DAC.encoder_dim: 64
-DAC.encoder_rates: [2, 4, 8, 8]
-DAC.latent_dim: 128
-DAC.decoder_dim: 1536
-DAC.decoder_rates: [8, 8, 4, 2]
-
-# Quantization
-DAC.n_codebooks: 18 # Max bitrate of 16kbps
-DAC.codebook_size: 1024
-DAC.codebook_dim: 8
-DAC.quantizer_dropout: 0.5
-
-# Discriminator
-Discriminator.sample_rate: 44100
-Discriminator.rates: []
-Discriminator.periods: [2, 3, 5, 7, 11]
-Discriminator.fft_sizes: [2048, 1024, 512]
-Discriminator.bands:
-  - [0.0, 0.1]
-  - [0.1, 0.25]
-  - [0.25, 0.5]
-  - [0.5, 0.75]
-  - [0.75, 1.0]
-
-# Optimization
-AdamW.betas: [0.8, 0.99]
-AdamW.lr: 0.0001
-ExponentialLR.gamma: 0.999996
-
-amp: false
-val_batch_size: 100
-device: cuda
-num_iters: 400000
-save_iters: [10000, 50000, 100000, 200000]
-valid_freq: 1000
-sample_freq: 10000
-num_workers: 32
-val_idx: [0, 1, 2, 3, 4, 5, 6, 7]
-seed: 0
-lambdas:
-  mel/loss: 15.0
-  adv/feat_loss: 2.0
-  adv/gen_loss: 1.0
-  vq/commitment_loss: 0.25
-  vq/codebook_loss: 1.0
-
-VolumeNorm.db: [const, -16]
-
-# Transforms
-build_transform.preprocess:
-  - Identity
-build_transform.augment_prob: 0.0
-build_transform.augment:
-  - Identity
-build_transform.postprocess:
-  - VolumeNorm
-  - RescaleAudio
-  - ShiftPhase
-
-# Loss setup
-MultiScaleSTFTLoss.window_lengths: [2048, 512]
-MelSpectrogramLoss.n_mels: [5, 10, 20, 40, 80, 160, 320]
-MelSpectrogramLoss.window_lengths: [32, 64, 128, 256, 512, 1024, 2048]
-MelSpectrogramLoss.mel_fmin: [0, 0, 0, 0, 0, 0, 0]
-MelSpectrogramLoss.mel_fmax: [null, null, null, null, null, null, null]
-MelSpectrogramLoss.pow: 1.0
-MelSpectrogramLoss.clamp_eps: 1.0e-5
-MelSpectrogramLoss.mag_weight: 0.0
-
-# Data
-batch_size: 72
-train/AudioDataset.duration: 0.38
-train/AudioDataset.n_examples: 10000000
-
-val/AudioDataset.duration: 5.0
-val/build_transform.augment_prob: 1.0
-val/AudioDataset.n_examples: 250
-
-test/AudioDataset.duration: 10.0
-test/build_transform.augment_prob: 1.0
-test/AudioDataset.n_examples: 1000
-
-AudioLoader.shuffle: true
-AudioDataset.without_replacement: true
-
-train/build_dataset.folders:
-  speech_fb:
-    - /data/daps/train
-  speech_hq:
-    - /data/vctk
-    - /data/vocalset
-    - /data/read_speech
-    - /data/french_speech
-  speech_uq:
-    - /data/emotional_speech/
-    - /data/common_voice/
-    - /data/german_speech/
-    - /data/russian_speech/
-    - /data/spanish_speech/
-  music_hq:
-    - /data/musdb/train
-  music_uq:
-    - /data/jamendo
-  general:
-    - /data/audioset/data/unbalanced_train_segments/
-    - /data/audioset/data/balanced_train_segments/
-
-val/build_dataset.folders:
-  speech_hq:
-    - /data/daps/val
-  music_hq:
-    - /data/musdb/test
-  general:
-    - /data/audioset/data/eval_segments/
-
-test/build_dataset.folders:
-  speech_hq:
-    - /data/daps/test
-  music_hq:
-    - /data/musdb/test
-  general:
-    - /data/audioset/data/eval_segments/

dac-codec/conf/final/44khz.yml

@@ -1,123 +0,0 @@
-# Model setup
-DAC.sample_rate: 44100
-DAC.encoder_dim: 64
-DAC.encoder_rates: [2, 4, 8, 8]
-DAC.decoder_dim: 1536
-DAC.decoder_rates: [8, 8, 4, 2]
-
-# Quantization
-DAC.n_codebooks: 9
-DAC.codebook_size: 1024
-DAC.codebook_dim: 8
-DAC.quantizer_dropout: 0.5
-
-# Discriminator
-Discriminator.sample_rate: 44100
-Discriminator.rates: []
-Discriminator.periods: [2, 3, 5, 7, 11]
-Discriminator.fft_sizes: [2048, 1024, 512]
-Discriminator.bands:
-  - [0.0, 0.1]
-  - [0.1, 0.25]
-  - [0.25, 0.5]
-  - [0.5, 0.75]
-  - [0.75, 1.0]
-
-# Optimization
-AdamW.betas: [0.8, 0.99]
-AdamW.lr: 0.0001
-ExponentialLR.gamma: 0.999996
-
-amp: false
-val_batch_size: 100
-device: cuda
-num_iters: 400000
-save_iters: [10000, 50000, 100000, 200000]
-valid_freq: 1000
-sample_freq: 10000
-num_workers: 32
-val_idx: [0, 1, 2, 3, 4, 5, 6, 7]
-seed: 0
-lambdas:
-  mel/loss: 15.0
-  adv/feat_loss: 2.0
-  adv/gen_loss: 1.0
-  vq/commitment_loss: 0.25
-  vq/codebook_loss: 1.0
-
-VolumeNorm.db: [const, -16]
-
-# Transforms
-build_transform.preprocess:
-  - Identity
-build_transform.augment_prob: 0.0
-build_transform.augment:
-  - Identity
-build_transform.postprocess:
-  - VolumeNorm
-  - RescaleAudio
-  - ShiftPhase
-
-# Loss setup
-MultiScaleSTFTLoss.window_lengths: [2048, 512]
-MelSpectrogramLoss.n_mels: [5, 10, 20, 40, 80, 160, 320]
-MelSpectrogramLoss.window_lengths: [32, 64, 128, 256, 512, 1024, 2048]
-MelSpectrogramLoss.mel_fmin: [0, 0, 0, 0, 0, 0, 0]
-MelSpectrogramLoss.mel_fmax: [null, null, null, null, null, null, null]
-MelSpectrogramLoss.pow: 1.0
-MelSpectrogramLoss.clamp_eps: 1.0e-5
-MelSpectrogramLoss.mag_weight: 0.0
-
-# Data
-batch_size: 72
-train/AudioDataset.duration: 0.38
-train/AudioDataset.n_examples: 10000000
-
-val/AudioDataset.duration: 5.0
-val/build_transform.augment_prob: 1.0
-val/AudioDataset.n_examples: 250
-
-test/AudioDataset.duration: 10.0
-test/build_transform.augment_prob: 1.0
-test/AudioDataset.n_examples: 1000
-
-AudioLoader.shuffle: true
-AudioDataset.without_replacement: true
-
-train/build_dataset.folders:
-  speech_fb:
-    - /data/daps/train
-  speech_hq:
-    - /data/vctk
-    - /data/vocalset
-    - /data/read_speech
-    - /data/french_speech
-  speech_uq:
-    - /data/emotional_speech/
-    - /data/common_voice/
-    - /data/german_speech/
-    - /data/russian_speech/
-    - /data/spanish_speech/
-  music_hq:
-    - /data/musdb/train
-  music_uq:
-    - /data/jamendo
-  general:
-    - /data/audioset/data/unbalanced_train_segments/
-    - /data/audioset/data/balanced_train_segments/
-
-val/build_dataset.folders:
-  speech_hq:
-    - /data/daps/val
-  music_hq:
-    - /data/musdb/test
-  general:
-    - /data/audioset/data/eval_segments/
-
-test/build_dataset.folders:
-  speech_hq:
-    - /data/daps/test
-  music_hq:
-    - /data/musdb/test
-  general:
-    - /data/audioset/data/eval_segments/

dac-codec/conf/quantizer/24kbps.yml

@@ -1,5 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-DAC.n_codebooks: 28

dac-codec/conf/quantizer/256d.yml

@@ -1,5 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-DAC.codebook_dim: 256

dac-codec/conf/quantizer/2d.yml

@@ -1,5 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-DAC.codebook_dim: 2

dac-codec/conf/quantizer/32d.yml

@@ -1,5 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-DAC.codebook_dim: 32

dac-codec/conf/quantizer/4d.yml

@@ -1,5 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-DAC.codebook_dim: 4

dac-codec/conf/quantizer/512d.yml

@@ -1,5 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-DAC.codebook_dim: 512

dac-codec/conf/quantizer/dropout-0.0.yml

@@ -1,5 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-DAC.quantizer_dropout: 0.0

dac-codec/conf/quantizer/dropout-0.25.yml

@@ -1,5 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-DAC.quantizer_dropout: 0.25

dac-codec/conf/quantizer/dropout-0.5.yml

@@ -1,5 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-DAC.quantizer_dropout: 0.5

dac-codec/conf/size/medium.yml

@@ -1,5 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-DAC.decoder_dim: 1024

dac-codec/conf/size/small.yml

@@ -1,5 +0,0 @@
-$include:
-  - conf/base.yml
-  - conf/1gpu.yml
-
-DAC.decoder_dim: 512

dac-codec/dac/__init__.py

@@ -1,16 +0,0 @@
-__version__ = "1.0.0"
-
-# preserved here for legacy reasons
-__model_version__ = "latest"
-
-import audiotools
-
-audiotools.ml.BaseModel.INTERN += ["dac.**"]
-audiotools.ml.BaseModel.EXTERN += ["einops"]
-
-
-from . import nn
-from . import model
-from . import utils
-from .model import DAC
-from .model import DACFile

dac-codec/dac/__main__.py

@@ -1,36 +0,0 @@
-import sys
-
-import argbind
-
-from dac.utils import download
-from dac.utils.decode import decode
-from dac.utils.encode import encode
-
-STAGES = ["encode", "decode", "download"]
-
-
-def run(stage: str):
-    """Run stages.
-
-    Parameters
-    ----------
-    stage : str
-        Stage to run
-    """
-    if stage not in STAGES:
-        raise ValueError(f"Unknown command: {stage}. Allowed commands are {STAGES}")
-    stage_fn = globals()[stage]
-
-    if stage == "download":
-        stage_fn()
-        return
-
-    stage_fn()
-
-
-if __name__ == "__main__":
-    group = sys.argv.pop(1)
-    args = argbind.parse_args(group=group)
-
-    with argbind.scope(args):
-        run(group)

dac-codec/dac/compare/encodec.py

@@ -1,54 +0,0 @@
-import torch
-from audiotools import AudioSignal
-from audiotools.ml import BaseModel
-from encodec import EncodecModel
-
-
-class Encodec(BaseModel):
-    def __init__(self, sample_rate: int = 24000, bandwidth: float = 24.0):
-        super().__init__()
-
-        if sample_rate == 24000:
-            self.model = EncodecModel.encodec_model_24khz()
-        else:
-            self.model = EncodecModel.encodec_model_48khz()
-        self.model.set_target_bandwidth(bandwidth)
-        self.sample_rate = 44100
-
-    def forward(
-        self,
-        audio_data: torch.Tensor,
-        sample_rate: int = 44100,
-        n_quantizers: int = None,
-    ):
-        signal = AudioSignal(audio_data, sample_rate)
-        signal.resample(self.model.sample_rate)
-        recons = self.model(signal.audio_data)
-        recons = AudioSignal(recons, self.model.sample_rate)
-        recons.resample(sample_rate)
-        return {"audio": recons.audio_data}
-
-
-if __name__ == "__main__":
-    import numpy as np
-    from functools import partial
-
-    model = Encodec()
-
-    for n, m in model.named_modules():
-        o = m.extra_repr()
-        p = sum([np.prod(p.size()) for p in m.parameters()])
-        fn = lambda o, p: o + f" {p/1e6:<.3f}M params."
-        setattr(m, "extra_repr", partial(fn, o=o, p=p))
-    print(model)
-    print("Total # of params: ", sum([np.prod(p.size()) for p in model.parameters()]))
-
-    length = 88200 * 2
-    x = torch.randn(1, 1, length).to(model.device)
-    x.requires_grad_(True)
-    x.retain_grad()
-
-    # Make a forward pass
-    out = model(x)["audio"]
-
-    print(x.shape, out.shape)

dac-codec/dac/model/__init__.py

@@ -1,4 +0,0 @@
-from .base import CodecMixin
-from .base import DACFile
-from .dac import DAC
-from .discriminator import Discriminator

dac-codec/dac/model/base.py

@@ -1,294 +0,0 @@
-import math
-from dataclasses import dataclass
-from pathlib import Path
-from typing import Union
-
-import numpy as np
-import torch
-import tqdm
-from audiotools import AudioSignal
-from torch import nn
-
-SUPPORTED_VERSIONS = ["1.0.0"]
-
-
-@dataclass
-class DACFile:
-    codes: torch.Tensor
-
-    # Metadata
-    chunk_length: int
-    original_length: int
-    input_db: float
-    channels: int
-    sample_rate: int
-    padding: bool
-    dac_version: str
-
-    def save(self, path):
-        artifacts = {
-            "codes": self.codes.numpy().astype(np.uint16),
-            "metadata": {
-                "input_db": self.input_db.numpy().astype(np.float32),
-                "original_length": self.original_length,
-                "sample_rate": self.sample_rate,
-                "chunk_length": self.chunk_length,
-                "channels": self.channels,
-                "padding": self.padding,
-                "dac_version": SUPPORTED_VERSIONS[-1],
-            },
-        }
-        path = Path(path).with_suffix(".dac")
-        with open(path, "wb") as f:
-            np.save(f, artifacts)
-        return path
-
-    @classmethod
-    def load(cls, path):
-        artifacts = np.load(path, allow_pickle=True)[()]
-        codes = torch.from_numpy(artifacts["codes"].astype(int))
-        if artifacts["metadata"].get("dac_version", None) not in SUPPORTED_VERSIONS:
-            raise RuntimeError(
-                f"Given file {path} can't be loaded with this version of descript-audio-codec."
-            )
-        return cls(codes=codes, **artifacts["metadata"])
-
-
-class CodecMixin:
-    @property
-    def padding(self):
-        if not hasattr(self, "_padding"):
-            self._padding = True
-        return self._padding
-
-    @padding.setter
-    def padding(self, value):
-        assert isinstance(value, bool)
-
-        layers = [
-            l for l in self.modules() if isinstance(l, (nn.Conv1d, nn.ConvTranspose1d))
-        ]
-
-        for layer in layers:
-            if value:
-                if hasattr(layer, "original_padding"):
-                    layer.padding = layer.original_padding
-            else:
-                layer.original_padding = layer.padding
-                layer.padding = tuple(0 for _ in range(len(layer.padding)))
-
-        self._padding = value
-
-    def get_delay(self):
-        # Any number works here, delay is invariant to input length
-        l_out = self.get_output_length(0)
-        L = l_out
-
-        layers = []
-        for layer in self.modules():
-            if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)):
-                layers.append(layer)
-
-        for layer in reversed(layers):
-            d = layer.dilation[0]
-            k = layer.kernel_size[0]
-            s = layer.stride[0]
-
-            if isinstance(layer, nn.ConvTranspose1d):
-                L = ((L - d * (k - 1) - 1) / s) + 1
-            elif isinstance(layer, nn.Conv1d):
-                L = (L - 1) * s + d * (k - 1) + 1
-
-            L = math.ceil(L)
-
-        l_in = L
-
-        return (l_in - l_out) // 2
-
-    def get_output_length(self, input_length):
-        L = input_length
-        # Calculate output length
-        for layer in self.modules():
-            if isinstance(layer, (nn.Conv1d, nn.ConvTranspose1d)):
-                d = layer.dilation[0]
-                k = layer.kernel_size[0]
-                s = layer.stride[0]
-
-                if isinstance(layer, nn.Conv1d):
-                    L = ((L - d * (k - 1) - 1) / s) + 1
-                elif isinstance(layer, nn.ConvTranspose1d):
-                    L = (L - 1) * s + d * (k - 1) + 1
-
-                L = math.floor(L)
-        return L
-
-    @torch.no_grad()
-    def compress(
-        self,
-        audio_path_or_signal: Union[str, Path, AudioSignal],
-        win_duration: float = 1.0,
-        verbose: bool = False,
-        normalize_db: float = -16,
-        n_quantizers: int = None,
-    ) -> DACFile:
-        """Processes an audio signal from a file or AudioSignal object into
-        discrete codes. This function processes the signal in short windows,
-        using constant GPU memory.
-
-        Parameters
-        ----------
-        audio_path_or_signal : Union[str, Path, AudioSignal]
-            audio signal to reconstruct
-        win_duration : float, optional
-            window duration in seconds, by default 5.0
-        verbose : bool, optional
-            by default False
-        normalize_db : float, optional
-            normalize db, by default -16
-
-        Returns
-        -------
-        DACFile
-            Object containing compressed codes and metadata
-            required for decompression
-        """
-        audio_signal = audio_path_or_signal
-        if isinstance(audio_signal, (str, Path)):
-            audio_signal = AudioSignal.load_from_file_with_ffmpeg(str(audio_signal))
-
-        self.eval()
-        original_padding = self.padding
-        original_device = audio_signal.device
-
-        audio_signal = audio_signal.clone()
-        original_sr = audio_signal.sample_rate
-
-        resample_fn = audio_signal.resample
-        loudness_fn = audio_signal.loudness
-
-        # If audio is > 10 minutes long, use the ffmpeg versions
-        if audio_signal.signal_duration >= 10 * 60 * 60:
-            resample_fn = audio_signal.ffmpeg_resample
-            loudness_fn = audio_signal.ffmpeg_loudness
-
-        original_length = audio_signal.signal_length
-        resample_fn(self.sample_rate)
-        input_db = loudness_fn()
-
-        if normalize_db is not None:
-            audio_signal.normalize(normalize_db)
-        audio_signal.ensure_max_of_audio()
-
-        nb, nac, nt = audio_signal.audio_data.shape
-        audio_signal.audio_data = audio_signal.audio_data.reshape(nb * nac, 1, nt)
-        win_duration = (
-            audio_signal.signal_duration if win_duration is None else win_duration
-        )
-
-        if audio_signal.signal_duration <= win_duration:
-            # Unchunked compression (used if signal length < win duration)
-            self.padding = True
-            n_samples = nt
-            hop = nt
-        else:
-            # Chunked inference
-            self.padding = False
-            # Zero-pad signal on either side by the delay
-            audio_signal.zero_pad(self.delay, self.delay)
-            n_samples = int(win_duration * self.sample_rate)
-            # Round n_samples to nearest hop length multiple
-            n_samples = int(math.ceil(n_samples / self.hop_length) * self.hop_length)
-            hop = self.get_output_length(n_samples)
-
-        codes = []
-        range_fn = range if not verbose else tqdm.trange
-
-        for i in range_fn(0, nt, hop):
-            x = audio_signal[..., i : i + n_samples]
-            x = x.zero_pad(0, max(0, n_samples - x.shape[-1]))
-
-            audio_data = x.audio_data.to(self.device)
-            audio_data = self.preprocess(audio_data, self.sample_rate)
-            _, c, _, _, _ = self.encode(audio_data, n_quantizers)
-            codes.append(c.to(original_device))
-            chunk_length = c.shape[-1]
-
-        codes = torch.cat(codes, dim=-1)
-
-        dac_file = DACFile(
-            codes=codes,
-            chunk_length=chunk_length,
-            original_length=original_length,
-            input_db=input_db,
-            channels=nac,
-            sample_rate=original_sr,
-            padding=self.padding,
-            dac_version=SUPPORTED_VERSIONS[-1],
-        )
-
-        if n_quantizers is not None:
-            codes = codes[:, :n_quantizers, :]
-
-        self.padding = original_padding
-        return dac_file
-
-    @torch.no_grad()
-    def decompress(
-        self,
-        obj: Union[str, Path, DACFile],
-        verbose: bool = False,
-    ) -> AudioSignal:
-        """Reconstruct audio from a given .dac file
-
-        Parameters
-        ----------
-        obj : Union[str, Path, DACFile]
-            .dac file location or corresponding DACFile object.
-        verbose : bool, optional
-            Prints progress if True, by default False
-
-        Returns
-        -------
-        AudioSignal
-            Object with the reconstructed audio
-        """
-        self.eval()
-        if isinstance(obj, (str, Path)):
-            obj = DACFile.load(obj)
-
-        original_padding = self.padding
-        self.padding = obj.padding
-
-        range_fn = range if not verbose else tqdm.trange
-        codes = obj.codes
-        original_device = codes.device
-        chunk_length = obj.chunk_length
-        recons = []
-
-        for i in range_fn(0, codes.shape[-1], chunk_length):
-            c = codes[..., i : i + chunk_length].to(self.device)
-            z = self.quantizer.from_codes(c)[0]
-            r = self.decode(z)
-            recons.append(r.to(original_device))
-
-        recons = torch.cat(recons, dim=-1)
-        recons = AudioSignal(recons, self.sample_rate)
-
-        resample_fn = recons.resample
-        loudness_fn = recons.loudness
-
-        # If audio is > 10 minutes long, use the ffmpeg versions
-        if recons.signal_duration >= 10 * 60 * 60:
-            resample_fn = recons.ffmpeg_resample
-            loudness_fn = recons.ffmpeg_loudness
-
-        recons.normalize(obj.input_db)
-        resample_fn(obj.sample_rate)
-        recons = recons[..., : obj.original_length]
-        loudness_fn()
-        recons.audio_data = recons.audio_data.reshape(
-            -1, obj.channels, obj.original_length
-        )
-
-        self.padding = original_padding
-        return recons

dac-codec/dac/model/dac.py

@@ -1,364 +0,0 @@
-import math
-from typing import List
-from typing import Union
-
-import numpy as np
-import torch
-from audiotools import AudioSignal
-from audiotools.ml import BaseModel
-from torch import nn
-
-from .base import CodecMixin
-from dac.nn.layers import Snake1d
-from dac.nn.layers import WNConv1d
-from dac.nn.layers import WNConvTranspose1d
-from dac.nn.quantize import ResidualVectorQuantize
-
-
-def init_weights(m):
-    if isinstance(m, nn.Conv1d):
-        nn.init.trunc_normal_(m.weight, std=0.02)
-        nn.init.constant_(m.bias, 0)
-
-
-class ResidualUnit(nn.Module):
-    def __init__(self, dim: int = 16, dilation: int = 1):
-        super().__init__()
-        pad = ((7 - 1) * dilation) // 2
-        self.block = nn.Sequential(
-            Snake1d(dim),
-            WNConv1d(dim, dim, kernel_size=7, dilation=dilation, padding=pad),
-            Snake1d(dim),
-            WNConv1d(dim, dim, kernel_size=1),
-        )
-
-    def forward(self, x):
-        y = self.block(x)
-        pad = (x.shape[-1] - y.shape[-1]) // 2
-        if pad > 0:
-            x = x[..., pad:-pad]
-        return x + y
-
-
-class EncoderBlock(nn.Module):
-    def __init__(self, dim: int = 16, stride: int = 1):
-        super().__init__()
-        self.block = nn.Sequential(
-            ResidualUnit(dim // 2, dilation=1),
-            ResidualUnit(dim // 2, dilation=3),
-            ResidualUnit(dim // 2, dilation=9),
-            Snake1d(dim // 2),
-            WNConv1d(
-                dim // 2,
-                dim,
-                kernel_size=2 * stride,
-                stride=stride,
-                padding=math.ceil(stride / 2),
-            ),
-        )
-
-    def forward(self, x):
-        return self.block(x)
-
-
-class Encoder(nn.Module):
-    def __init__(
-        self,
-        d_model: int = 64,
-        strides: list = [2, 4, 8, 8],
-        d_latent: int = 64,
-    ):
-        super().__init__()
-        # Create first convolution
-        self.block = [WNConv1d(1, d_model, kernel_size=7, padding=3)]
-
-        # Create EncoderBlocks that double channels as they downsample by `stride`
-        for stride in strides:
-            d_model *= 2
-            self.block += [EncoderBlock(d_model, stride=stride)]
-
-        # Create last convolution
-        self.block += [
-            Snake1d(d_model),
-            WNConv1d(d_model, d_latent, kernel_size=3, padding=1),
-        ]
-
-        # Wrap black into nn.Sequential
-        self.block = nn.Sequential(*self.block)
-        self.enc_dim = d_model
-
-    def forward(self, x):
-        return self.block(x)
-
-
-class DecoderBlock(nn.Module):
-    def __init__(self, input_dim: int = 16, output_dim: int = 8, stride: int = 1):
-        super().__init__()
-        self.block = nn.Sequential(
-            Snake1d(input_dim),
-            WNConvTranspose1d(
-                input_dim,
-                output_dim,
-                kernel_size=2 * stride,
-                stride=stride,
-                padding=math.ceil(stride / 2),
-            ),
-            ResidualUnit(output_dim, dilation=1),
-            ResidualUnit(output_dim, dilation=3),
-            ResidualUnit(output_dim, dilation=9),
-        )
-
-    def forward(self, x):
-        return self.block(x)
-
-
-class Decoder(nn.Module):
-    def __init__(
-        self,
-        input_channel,
-        channels,
-        rates,
-        d_out: int = 1,
-    ):
-        super().__init__()
-
-        # Add first conv layer
-        layers = [WNConv1d(input_channel, channels, kernel_size=7, padding=3)]
-
-        # Add upsampling + MRF blocks
-        for i, stride in enumerate(rates):
-            input_dim = channels // 2**i
-            output_dim = channels // 2 ** (i + 1)
-            layers += [DecoderBlock(input_dim, output_dim, stride)]
-
-        # Add final conv layer
-        layers += [
-            Snake1d(output_dim),
-            WNConv1d(output_dim, d_out, kernel_size=7, padding=3),
-            nn.Tanh(),
-        ]
-
-        self.model = nn.Sequential(*layers)
-
-    def forward(self, x):
-        return self.model(x)
-
-
-class DAC(BaseModel, CodecMixin):
-    def __init__(
-        self,
-        encoder_dim: int = 64,
-        encoder_rates: List[int] = [2, 4, 8, 8],
-        latent_dim: int = None,
-        decoder_dim: int = 1536,
-        decoder_rates: List[int] = [8, 8, 4, 2],
-        n_codebooks: int = 9,
-        codebook_size: int = 1024,
-        codebook_dim: Union[int, list] = 8,
-        quantizer_dropout: bool = False,
-        sample_rate: int = 44100,
-    ):
-        super().__init__()
-
-        self.encoder_dim = encoder_dim
-        self.encoder_rates = encoder_rates
-        self.decoder_dim = decoder_dim
-        self.decoder_rates = decoder_rates
-        self.sample_rate = sample_rate
-
-        if latent_dim is None:
-            latent_dim = encoder_dim * (2 ** len(encoder_rates))
-
-        self.latent_dim = latent_dim
-
-        self.hop_length = np.prod(encoder_rates)
-        self.encoder = Encoder(encoder_dim, encoder_rates, latent_dim)
-
-        self.n_codebooks = n_codebooks
-        self.codebook_size = codebook_size
-        self.codebook_dim = codebook_dim
-        self.quantizer = ResidualVectorQuantize(
-            input_dim=latent_dim,
-            n_codebooks=n_codebooks,
-            codebook_size=codebook_size,
-            codebook_dim=codebook_dim,
-            quantizer_dropout=quantizer_dropout,
-        )
-
-        self.decoder = Decoder(
-            latent_dim,
-            decoder_dim,
-            decoder_rates,
-        )
-        self.sample_rate = sample_rate
-        self.apply(init_weights)
-
-        self.delay = self.get_delay()
-
-    def preprocess(self, audio_data, sample_rate):
-        if sample_rate is None:
-            sample_rate = self.sample_rate
-        assert sample_rate == self.sample_rate
-
-        length = audio_data.shape[-1]
-        right_pad = math.ceil(length / self.hop_length) * self.hop_length - length
-        audio_data = nn.functional.pad(audio_data, (0, right_pad))
-
-        return audio_data
-
-    def encode(
-        self,
-        audio_data: torch.Tensor,
-        n_quantizers: int = None,
-    ):
-        """Encode given audio data and return quantized latent codes
-
-        Parameters
-        ----------
-        audio_data : Tensor[B x 1 x T]
-            Audio data to encode
-        n_quantizers : int, optional
-            Number of quantizers to use, by default None
-            If None, all quantizers are used.
-
-        Returns
-        -------
-        dict
-            A dictionary with the following keys:
-            "z" : Tensor[B x D x T]
-                Quantized continuous representation of input
-            "codes" : Tensor[B x N x T]
-                Codebook indices for each codebook
-                (quantized discrete representation of input)
-            "latents" : Tensor[B x N*D x T]
-                Projected latents (continuous representation of input before quantization)
-            "vq/commitment_loss" : Tensor[1]
-                Commitment loss to train encoder to predict vectors closer to codebook
-                entries
-            "vq/codebook_loss" : Tensor[1]
-                Codebook loss to update the codebook
-            "length" : int
-                Number of samples in input audio
-        """
-        z = self.encoder(audio_data)
-        z, codes, latents, commitment_loss, codebook_loss = self.quantizer(
-            z, n_quantizers
-        )
-        return z, codes, latents, commitment_loss, codebook_loss
-
-    def decode(self, z: torch.Tensor):
-        """Decode given latent codes and return audio data
-
-        Parameters
-        ----------
-        z : Tensor[B x D x T]
-            Quantized continuous representation of input
-        length : int, optional
-            Number of samples in output audio, by default None
-
-        Returns
-        -------
-        dict
-            A dictionary with the following keys:
-            "audio" : Tensor[B x 1 x length]
-                Decoded audio data.
-        """
-        return self.decoder(z)
-
-    def forward(
-        self,
-        audio_data: torch.Tensor,
-        sample_rate: int = None,
-        n_quantizers: int = None,
-    ):
-        """Model forward pass
-
-        Parameters
-        ----------
-        audio_data : Tensor[B x 1 x T]
-            Audio data to encode
-        sample_rate : int, optional
-            Sample rate of audio data in Hz, by default None
-            If None, defaults to `self.sample_rate`
-        n_quantizers : int, optional
-            Number of quantizers to use, by default None.
-            If None, all quantizers are used.
-
-        Returns
-        -------
-        dict
-            A dictionary with the following keys:
-            "z" : Tensor[B x D x T]
-                Quantized continuous representation of input
-            "codes" : Tensor[B x N x T]
-                Codebook indices for each codebook
-                (quantized discrete representation of input)
-            "latents" : Tensor[B x N*D x T]
-                Projected latents (continuous representation of input before quantization)
-            "vq/commitment_loss" : Tensor[1]
-                Commitment loss to train encoder to predict vectors closer to codebook
-                entries
-            "vq/codebook_loss" : Tensor[1]
-                Codebook loss to update the codebook
-            "length" : int
-                Number of samples in input audio
-            "audio" : Tensor[B x 1 x length]
-                Decoded audio data.
-        """
-        length = audio_data.shape[-1]
-        audio_data = self.preprocess(audio_data, sample_rate)
-        z, codes, latents, commitment_loss, codebook_loss = self.encode(
-            audio_data, n_quantizers
-        )
-
-        x = self.decode(z)
-        return {
-            "audio": x[..., :length],
-            "z": z,
-            "codes": codes,
-            "latents": latents,
-            "vq/commitment_loss": commitment_loss,
-            "vq/codebook_loss": codebook_loss,
-        }
-
-
-if __name__ == "__main__":
-    import numpy as np
-    from functools import partial
-
-    model = DAC().to("cpu")
-
-    for n, m in model.named_modules():
-        o = m.extra_repr()
-        p = sum([np.prod(p.size()) for p in m.parameters()])
-        fn = lambda o, p: o + f" {p/1e6:<.3f}M params."
-        setattr(m, "extra_repr", partial(fn, o=o, p=p))
-    print(model)
-    print("Total # of params: ", sum([np.prod(p.size()) for p in model.parameters()]))
-
-    length = 88200 * 2
-    x = torch.randn(1, 1, length).to(model.device)
-    x.requires_grad_(True)
-    x.retain_grad()
-
-    # Make a forward pass
-    out = model(x)["audio"]
-    print("Input shape:", x.shape)
-    print("Output shape:", out.shape)
-
-    # Create gradient variable
-    grad = torch.zeros_like(out)
-    grad[:, :, grad.shape[-1] // 2] = 1
-
-    # Make a backward pass
-    out.backward(grad)
-
-    # Check non-zero values
-    gradmap = x.grad.squeeze(0)
-    gradmap = (gradmap != 0).sum(0)  # sum across features
-    rf = (gradmap != 0).sum()
-
-    print(f"Receptive field: {rf.item()}")
-
-    x = AudioSignal(torch.randn(1, 1, 44100 * 60), 44100)
-    model.decompress(model.compress(x, verbose=True), verbose=True)

dac-codec/dac/model/discriminator.py

@@ -1,228 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from audiotools import AudioSignal
-from audiotools import ml
-from audiotools import STFTParams
-from einops import rearrange
-from torch.nn.utils import weight_norm
-
-
-def WNConv1d(*args, **kwargs):
-    act = kwargs.pop("act", True)
-    conv = weight_norm(nn.Conv1d(*args, **kwargs))
-    if not act:
-        return conv
-    return nn.Sequential(conv, nn.LeakyReLU(0.1))
-
-
-def WNConv2d(*args, **kwargs):
-    act = kwargs.pop("act", True)
-    conv = weight_norm(nn.Conv2d(*args, **kwargs))
-    if not act:
-        return conv
-    return nn.Sequential(conv, nn.LeakyReLU(0.1))
-
-
-class MPD(nn.Module):
-    def __init__(self, period):
-        super().__init__()
-        self.period = period
-        self.convs = nn.ModuleList(
-            [
-                WNConv2d(1, 32, (5, 1), (3, 1), padding=(2, 0)),
-                WNConv2d(32, 128, (5, 1), (3, 1), padding=(2, 0)),
-                WNConv2d(128, 512, (5, 1), (3, 1), padding=(2, 0)),
-                WNConv2d(512, 1024, (5, 1), (3, 1), padding=(2, 0)),
-                WNConv2d(1024, 1024, (5, 1), 1, padding=(2, 0)),
-            ]
-        )
-        self.conv_post = WNConv2d(
-            1024, 1, kernel_size=(3, 1), padding=(1, 0), act=False
-        )
-
-    def pad_to_period(self, x):
-        t = x.shape[-1]
-        x = F.pad(x, (0, self.period - t % self.period), mode="reflect")
-        return x
-
-    def forward(self, x):
-        fmap = []
-
-        x = self.pad_to_period(x)
-        x = rearrange(x, "b c (l p) -> b c l p", p=self.period)
-
-        for layer in self.convs:
-            x = layer(x)
-            fmap.append(x)
-
-        x = self.conv_post(x)
-        fmap.append(x)
-
-        return fmap
-
-
-class MSD(nn.Module):
-    def __init__(self, rate: int = 1, sample_rate: int = 44100):
-        super().__init__()
-        self.convs = nn.ModuleList(
-            [
-                WNConv1d(1, 16, 15, 1, padding=7),
-                WNConv1d(16, 64, 41, 4, groups=4, padding=20),
-                WNConv1d(64, 256, 41, 4, groups=16, padding=20),
-                WNConv1d(256, 1024, 41, 4, groups=64, padding=20),
-                WNConv1d(1024, 1024, 41, 4, groups=256, padding=20),
-                WNConv1d(1024, 1024, 5, 1, padding=2),
-            ]
-        )
-        self.conv_post = WNConv1d(1024, 1, 3, 1, padding=1, act=False)
-        self.sample_rate = sample_rate
-        self.rate = rate
-
-    def forward(self, x):
-        x = AudioSignal(x, self.sample_rate)
-        x.resample(self.sample_rate // self.rate)
-        x = x.audio_data
-
-        fmap = []
-
-        for l in self.convs:
-            x = l(x)
-            fmap.append(x)
-        x = self.conv_post(x)
-        fmap.append(x)
-
-        return fmap
-
-
-BANDS = [(0.0, 0.1), (0.1, 0.25), (0.25, 0.5), (0.5, 0.75), (0.75, 1.0)]
-
-
-class MRD(nn.Module):
-    def __init__(
-        self,
-        window_length: int,
-        hop_factor: float = 0.25,
-        sample_rate: int = 44100,
-        bands: list = BANDS,
-    ):
-        """Complex multi-band spectrogram discriminator.
-        Parameters
-        ----------
-        window_length : int
-            Window length of STFT.
-        hop_factor : float, optional
-            Hop factor of the STFT, defaults to ``0.25 * window_length``.
-        sample_rate : int, optional
-            Sampling rate of audio in Hz, by default 44100
-        bands : list, optional
-            Bands to run discriminator over.
-        """
-        super().__init__()
-
-        self.window_length = window_length
-        self.hop_factor = hop_factor
-        self.sample_rate = sample_rate
-        self.stft_params = STFTParams(
-            window_length=window_length,
-            hop_length=int(window_length * hop_factor),
-            match_stride=True,
-        )
-
-        n_fft = window_length // 2 + 1
-        bands = [(int(b[0] * n_fft), int(b[1] * n_fft)) for b in bands]
-        self.bands = bands
-
-        ch = 32
-        convs = lambda: nn.ModuleList(
-            [
-                WNConv2d(2, ch, (3, 9), (1, 1), padding=(1, 4)),
-                WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
-                WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
-                WNConv2d(ch, ch, (3, 9), (1, 2), padding=(1, 4)),
-                WNConv2d(ch, ch, (3, 3), (1, 1), padding=(1, 1)),
-            ]
-        )
-        self.band_convs = nn.ModuleList([convs() for _ in range(len(self.bands))])
-        self.conv_post = WNConv2d(ch, 1, (3, 3), (1, 1), padding=(1, 1), act=False)
-
-    def spectrogram(self, x):
-        x = AudioSignal(x, self.sample_rate, stft_params=self.stft_params)
-        x = torch.view_as_real(x.stft())
-        x = rearrange(x, "b 1 f t c -> (b 1) c t f")
-        # Split into bands
-        x_bands = [x[..., b[0] : b[1]] for b in self.bands]
-        return x_bands
-
-    def forward(self, x):
-        x_bands = self.spectrogram(x)
-        fmap = []
-
-        x = []
-        for band, stack in zip(x_bands, self.band_convs):
-            for layer in stack:
-                band = layer(band)
-                fmap.append(band)
-            x.append(band)
-
-        x = torch.cat(x, dim=-1)
-        x = self.conv_post(x)
-        fmap.append(x)
-
-        return fmap
-
-
-class Discriminator(ml.BaseModel):
-    def __init__(
-        self,
-        rates: list = [],
-        periods: list = [2, 3, 5, 7, 11],
-        fft_sizes: list = [2048, 1024, 512],
-        sample_rate: int = 44100,
-        bands: list = BANDS,
-    ):
-        """Discriminator that combines multiple discriminators.
-
-        Parameters
-        ----------
-        rates : list, optional
-            sampling rates (in Hz) to run MSD at, by default []
-            If empty, MSD is not used.
-        periods : list, optional
-            periods (of samples) to run MPD at, by default [2, 3, 5, 7, 11]
-        fft_sizes : list, optional
-            Window sizes of the FFT to run MRD at, by default [2048, 1024, 512]
-        sample_rate : int, optional
-            Sampling rate of audio in Hz, by default 44100
-        bands : list, optional
-            Bands to run MRD at, by default `BANDS`
-        """
-        super().__init__()
-        discs = []
-        discs += [MPD(p) for p in periods]
-        discs += [MSD(r, sample_rate=sample_rate) for r in rates]
-        discs += [MRD(f, sample_rate=sample_rate, bands=bands) for f in fft_sizes]
-        self.discriminators = nn.ModuleList(discs)
-
-    def preprocess(self, y):
-        # Remove DC offset
-        y = y - y.mean(dim=-1, keepdims=True)
-        # Peak normalize the volume of input audio
-        y = 0.8 * y / (y.abs().max(dim=-1, keepdim=True)[0] + 1e-9)
-        return y
-
-    def forward(self, x):
-        x = self.preprocess(x)
-        fmaps = [d(x) for d in self.discriminators]
-        return fmaps
-
-
-if __name__ == "__main__":
-    disc = Discriminator()
-    x = torch.zeros(1, 1, 44100)
-    results = disc(x)
-    for i, result in enumerate(results):
-        print(f"disc{i}")
-        for i, r in enumerate(result):
-            print(r.shape, r.mean(), r.min(), r.max())
-        print()

dac-codec/dac/nn/__init__.py

@@ -1,3 +0,0 @@
-from . import layers
-from . import loss
-from . import quantize