v2.3

.gitignore

@@ -1,6 +1,137 @@
-*.pyc
-__pycache__/
-*/__pycache__/
-alias_free_cuda/build/
+# BigVGAN
+alias_free_activation/cuda/build/
 exp/
-tmp/
+tmp/
+
+# VSCode configs
+.vscode/
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+.idea/

README_model.md

@@ -1,37 +1,96 @@
 ## BigVGAN: A Universal Neural Vocoder with Large-Scale Training
+
 #### Sang-gil Lee, Wei Ping, Boris Ginsburg, Bryan Catanzaro, Sungroh Yoon
 
-<center><img src="https://user-images.githubusercontent.com/15963413/218609148-881e39df-33af-4af9-ab95-1427c4ebf062.png" width="800"></center>
+[[Paper]](https://arxiv.org/abs/2206.04658) - [[Code]](https://github.com/NVIDIA/BigVGAN) - [[Showcase]](https://bigvgan-demo.github.io/) - [[Project Page]](https://research.nvidia.com/labs/adlr/projects/bigvgan/) - [[Weights]](https://huggingface.co/collections/nvidia/bigvgan-66959df3d97fd7d98d97dc9a) - [[Demo]](https://huggingface.co/spaces/nvidia/BigVGAN)
 
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/bigvgan-a-universal-neural-vocoder-with-large/speech-synthesis-on-libritts)](https://paperswithcode.com/sota/speech-synthesis-on-libritts?p=bigvgan-a-universal-neural-vocoder-with-large)
 
-### [Paper](https://arxiv.org/abs/2206.04658) &emsp; [Project page](https://research.nvidia.com/labs/adlr/projects/bigvgan/) &emsp; [Audio demo](https://bigvgan-demo.github.io/)
+<center><img src="https://user-images.githubusercontent.com/15963413/218609148-881e39df-33af-4af9-ab95-1427c4ebf062.png" width="800"></center>
 
 ## News
-[Jul 2024] We release BigVGAN-v2 along with pretrained checkpoints. Below are the highlights:
-* Custom CUDA kernel for inference: we provide a fused upsampling + activation kernel written in CUDA for accelerated inference speed. Our test shows 1.5 - 3x faster speed on a single A100 GPU. 
-* Improved discriminator and loss: BigVGAN-v2 is trained using a [multi-scale sub-band CQT discriminator](https://arxiv.org/abs/2311.14957) and a [multi-scale mel spectrogram loss](https://arxiv.org/abs/2306.06546).
-* Larger training data: BigVGAN-v2 is trained using datasets containing diverse audio types, including speech in multiple languages, environmental sounds, and instruments.
-* We provide pretrained checkpoints of BigVGAN-v2 using diverse audio configurations, supporting up to 44 kHz sampling rate and 512x upsampling ratio.
+- **Jul 2024 (v2.3):**
+  - General refactor and code improvements for improved readability.
+  - Fully fused CUDA kernel of anti-alised activation (upsampling + activation + downsampling) with inference speed benchmark.
+
+- **Jul 2024 (v2.2):** The repository now includes an interactive local demo using gradio.
+
+- **Jul 2024 (v2.1):** BigVGAN is now integrated with 🤗 Hugging Face Hub with easy access to inference using pretrained checkpoints. We also provide an interactive demo on Hugging Face Spaces.
+
+- **Jul 2024 (v2):** We release BigVGAN-v2 along with pretrained checkpoints. Below are the highlights:
+  - Custom CUDA kernel for inference: we provide a fused upsampling + activation kernel written in CUDA for accelerated inference speed. Our test shows 1.5 - 3x faster speed on a single A100 GPU.
+  - Improved discriminator and loss: BigVGAN-v2 is trained using a [multi-scale sub-band CQT discriminator](https://arxiv.org/abs/2311.14957) and a [multi-scale mel spectrogram loss](https://arxiv.org/abs/2306.06546).
+  - Larger training data: BigVGAN-v2 is trained using datasets containing diverse audio types, including speech in multiple languages, environmental sounds, and instruments.
+  - We provide pretrained checkpoints of BigVGAN-v2 using diverse audio configurations, supporting up to 44 kHz sampling rate and 512x upsampling ratio.
 
 ## Installation
+
 The codebase has been tested on Python `3.10` and PyTorch `2.3.1` conda packages with either `pytorch-cuda=12.1` or `pytorch-cuda=11.8`. Below is an example command to create the conda environment:
+
 ```shell
 conda create -n bigvgan python=3.10 pytorch torchvision torchaudio pytorch-cuda=12.1 -c pytorch -c nvidia
 conda activate bigvgan
 ```
 
 Clone the repository and install dependencies:
+
 ```shell
 git clone https://github.com/NVIDIA/BigVGAN
 cd BigVGAN
 pip install -r requirements.txt
 ```
 
+## Inference Quickstart using 🤗 Hugging Face Hub
+
+Below example describes how you can use BigVGAN: load the pretrained BigVGAN generator from Hugging Face Hub, compute mel spectrogram from input waveform, and generate synthesized waveform using the mel spectrogram as the model's input.
+
+```python
+device = 'cuda'
+
+import torch
+import bigvgan
+import librosa
+from meldataset import get_mel_spectrogram
+
+# instantiate the model. You can optionally set use_cuda_kernel=True for faster inference.
+model = bigvgan.BigVGAN.from_pretrained('nvidia/bigvgan_v2_24khz_100band_256x', use_cuda_kernel=False)
+
+# remove weight norm in the model and set to eval mode
+model.remove_weight_norm()
+model = model.eval().to(device)
+
+# load wav file and compute mel spectrogram
+wav_path = '/path/to/your/audio.wav'
+wav, sr = librosa.load(wav_path, sr=model.h.sampling_rate, mono=True) # wav is np.ndarray with shape [T_time] and values in [-1, 1]
+wav = torch.FloatTensor(wav).unsqueeze(0) # wav is FloatTensor with shape [B(1), T_time]
+
+# compute mel spectrogram from the ground truth audio
+mel = get_mel_spectrogram(wav, model.h).to(device) # mel is FloatTensor with shape [B(1), C_mel, T_frame]
+
+# generate waveform from mel
+with torch.inference_mode():
+    wav_gen = model(mel) # wav_gen is FloatTensor with shape [B(1), 1, T_time] and values in [-1, 1]
+wav_gen_float = wav_gen.squeeze(0).cpu() # wav_gen is FloatTensor with shape [1, T_time]
+
+# you can convert the generated waveform to 16 bit linear PCM
+wav_gen_int16 = (wav_gen_float * 32767.0).numpy().astype('int16') # wav_gen is now np.ndarray with shape [1, T_time] and int16 dtype
+```
+
+## Local gradio demo <a href='https://github.com/gradio-app/gradio'><img src='https://img.shields.io/github/stars/gradio-app/gradio'></a>
 
+You can run a local gradio demo using below command:
+
+```python
+pip install -r demo/requirements.txt
+python demo/app.py
+```
+
+## Training
 
 Create symbolic link to the root of the dataset. The codebase uses filelist with the relative path from the dataset. Below are the example commands for LibriTTS dataset:
-``` shell
-cd LibriTTS && \
+
+```shell
+cd filelists/LibriTTS && \
 ln -s /path/to/your/LibriTTS/train-clean-100 train-clean-100 && \
 ln -s /path/to/your/LibriTTS/train-clean-360 train-clean-360 && \
 ln -s /path/to/your/LibriTTS/train-other-500 train-other-500 && \
@@ -39,29 +98,30 @@ ln -s /path/to/your/LibriTTS/dev-clean dev-clean && \
 ln -s /path/to/your/LibriTTS/dev-other dev-other && \
 ln -s /path/to/your/LibriTTS/test-clean test-clean && \
 ln -s /path/to/your/LibriTTS/test-other test-other && \
-cd ..
+cd ../..
 ```
 
-## Training
 Train BigVGAN model. Below is an example command for training BigVGAN-v2 using LibriTTS dataset at 24kHz with a full 100-band mel spectrogram as input:
+
 ```shell
 python train.py \
 --config configs/bigvgan_v2_24khz_100band_256x.json \
---input_wavs_dir LibriTTS \
---input_training_file LibriTTS/train-full.txt \
---input_validation_file LibriTTS/val-full.txt \
---list_input_unseen_wavs_dir LibriTTS LibriTTS \
---list_input_unseen_validation_file LibriTTS/dev-clean.txt LibriTTS/dev-other.txt \
+--input_wavs_dir filelists/LibriTTS \
+--input_training_file filelists/LibriTTS/train-full.txt \
+--input_validation_file filelists/LibriTTS/val-full.txt \
+--list_input_unseen_wavs_dir filelists/LibriTTS filelists/LibriTTS \
+--list_input_unseen_validation_file filelists/LibriTTS/dev-clean.txt filelists/LibriTTS/dev-other.txt \
 --checkpoint_path exp/bigvgan_v2_24khz_100band_256x
 ```
 
-
 ## Synthesis
+
 Synthesize from BigVGAN model. Below is an example command for generating audio from the model.
 It computes mel spectrograms using wav files from `--input_wavs_dir` and saves the generated audio to `--output_dir`.
+
 ```shell
 python inference.py \
---checkpoint_file exp/bigvgan_v2_24khz_100band_256x/g_03000000 \
+--checkpoint_file /path/to/your/bigvgan_v2_24khz_100band_256x/bigvgan_generator.pt \
 --input_wavs_dir /path/to/your/input_wav \
 --output_dir /path/to/your/output_wav
 ```
@@ -70,14 +130,16 @@ python inference.py \
 It loads mel spectrograms from `--input_mels_dir` and saves the generated audio to `--output_dir`.
 
 Make sure that the STFT hyperparameters for mel spectrogram are the same as the model, which are defined in `config.json` of the corresponding model.
+
 ```shell
 python inference_e2e.py \
---checkpoint_file exp/bigvgan_v2_24khz_100band_256x/g_03000000 \
+--checkpoint_file /path/to/your/bigvgan_v2_24khz_100band_256x/bigvgan_generator.pt \
 --input_mels_dir /path/to/your/input_mel \
 --output_dir /path/to/your/output_wav
 ```
 
 ## Using Custom CUDA Kernel for Synthesis
+
 You can apply the fast CUDA inference kernel by using a parameter `use_cuda_kernel` when instantiating BigVGAN:
 
 ```python
@@ -86,15 +148,15 @@ generator = BigVGAN(h, use_cuda_kernel=True)
 
 You can also pass `--use_cuda_kernel` to `inference.py` and `inference_e2e.py` to enable this feature.
 
-When applied for the first time, it builds the kernel using `nvcc` and `ninja`. If the build succeeds, the kernel is saved to `alias_free_cuda/build` and the model automatically loads the kernel. The codebase has been tested using CUDA `12.1`.
+When applied for the first time, it builds the kernel using `nvcc` and `ninja`. If the build succeeds, the kernel is saved to `alias_free_activation/cuda/build` and the model automatically loads the kernel. The codebase has been tested using CUDA `12.1`.
 
 Please make sure that both are installed in your system and `nvcc` installed in your system matches the version your PyTorch build is using.
 
 We recommend running `test_cuda_vs_torch_model.py` first to build and check the correctness of the CUDA kernel. See below example command and its output, where it returns `[Success] test CUDA fused vs. plain torch BigVGAN inference`:
 
 ```python
-python test_cuda_vs_torch_model.py \
---checkpoint_file /path/to/your/bigvgan/g_03000000
+python tests/test_cuda_vs_torch_model.py \
+--checkpoint_file /path/to/your/bigvgan_generator.pt
 ```
 
 ```shell
@@ -102,12 +164,12 @@ loading plain Pytorch BigVGAN
 ...
 loading CUDA kernel BigVGAN with auto-build
 Detected CUDA files, patching ldflags
-Emitting ninja build file /path/to/your/BigVGAN/alias_free_cuda/build/build.ninja...
+Emitting ninja build file /path/to/your/BigVGAN/alias_free_activation/cuda/build/build.ninja..
 Building extension module anti_alias_activation_cuda...
 ...
 Loading extension module anti_alias_activation_cuda...
 ...
-Loading '/path/to/your/bigvgan/g_03000000'
+Loading '/path/to/your/bigvgan_generator.pt'
 ...
 [Success] test CUDA fused vs. plain torch BigVGAN inference
  > mean_difference=0.0007238413265440613
@@ -116,30 +178,34 @@ Loading '/path/to/your/bigvgan/g_03000000'
 
 If you see `[Fail] test CUDA fused vs. plain torch BigVGAN inference`, it means that the CUDA kernel inference is incorrect. Please check if `nvcc` installed in your system is compatible with your PyTorch version.
 
-
 ## Pretrained Models
-We provide the [pretrained models](https://drive.google.com/drive/folders/1L2RDeJMBE7QAI8qV51n0QAf4mkSgUUeE?usp=sharing).
-One can download the checkpoints of the generator weight (e.g., `g_(training_steps)`) and its discriminator/optimizer states (e.g., `do_(training_steps)`) within the listed folders.
-
-|Folder Name|Sampling Rate|Mel band|fmax|Upsampling Ratio|Params.|Dataset|Fine-Tuned|
-|------|---|---|---|---|---|------|---|
-|bigvgan_v2_44khz_128band_512x|44 kHz|128|22050|512|122M|Large-scale Compilation|No|
-|bigvgan_v2_44khz_128band_256x|44 kHz|128|22050|256|112M|Large-scale Compilation|No|
-|bigvgan_v2_24khz_100band_256x|24 kHz|100|12000|256|112M|Large-scale Compilation|No|
-|bigvgan_v2_22khz_80band_256x|22 kHz|80|11025|256|112M|Large-scale Compilation|No|
-|bigvgan_v2_22khz_80band_fmax8k_256x|22 kHz|80|8000|256|112M|Large-scale Compilation|No|
-|bigvgan_24khz_100band|24 kHz|100|12000|256|112M|LibriTTS|No|
-|bigvgan_base_24khz_100band|24 kHz|100|12000|256|14M|LibriTTS|No|
-|bigvgan_22khz_80band|22 kHz|80|8000|256|112M|LibriTTS + VCTK + LJSpeech|No|
-|bigvgan_base_22khz_80band|22 kHz|80|8000|256|14M|LibriTTS + VCTK + LJSpeech|No|
+
+We provide the [pretrained models on Hugging Face Collections](https://huggingface.co/collections/nvidia/bigvgan-66959df3d97fd7d98d97dc9a).
+One can download the checkpoints of the generator weight (named `bigvgan_generator.pt`) and its discriminator/optimizer states (named `bigvgan_discriminator_optimizer.pt`) within the listed model repositories.
+
+| Model Name                                                                                               | Sampling Rate | Mel band | fmax  | Upsampling Ratio | Params | Dataset                    | Steps | Fine-Tuned |
+|:--------------------------------------------------------------------------------------------------------:|:-------------:|:--------:|:-----:|:----------------:|:------:|:--------------------------:|:-----:|:----------:|
+| [bigvgan_v2_44khz_128band_512x](https://huggingface.co/nvidia/bigvgan_v2_44khz_128band_512x)             | 44 kHz        | 128      | 22050 | 512              | 122M   | Large-scale Compilation    | 3M    | No         |
+| [bigvgan_v2_44khz_128band_256x](https://huggingface.co/nvidia/bigvgan_v2_44khz_128band_256x)             | 44 kHz        | 128      | 22050 | 256              | 112M   | Large-scale Compilation    | 3M    | No         |
+| [bigvgan_v2_24khz_100band_256x](https://huggingface.co/nvidia/bigvgan_v2_24khz_100band_256x)             | 24 kHz        | 100      | 12000 | 256              | 112M   | Large-scale Compilation    | 3M    | No         |
+| [bigvgan_v2_22khz_80band_256x](https://huggingface.co/nvidia/bigvgan_v2_22khz_80band_256x)               | 22 kHz        | 80       | 11025 | 256              | 112M   | Large-scale Compilation    | 3M    | No         |
+| [bigvgan_v2_22khz_80band_fmax8k_256x](https://huggingface.co/nvidia/bigvgan_v2_22khz_80band_fmax8k_256x) | 22 kHz        | 80       | 8000  | 256              | 112M   | Large-scale Compilation    | 3M    | No         |
+| [bigvgan_24khz_100band](https://huggingface.co/nvidia/bigvgan_24khz_100band)                             | 24 kHz        | 100      | 12000 | 256              | 112M   | LibriTTS                   | 5M    | No         |
+| [bigvgan_base_24khz_100band](https://huggingface.co/nvidia/bigvgan_base_24khz_100band)                   | 24 kHz        | 100      | 12000 | 256              | 14M    | LibriTTS                   | 5M    | No         |
+| [bigvgan_22khz_80band](https://huggingface.co/nvidia/bigvgan_22khz_80band)                               | 22 kHz        | 80       | 8000  | 256              | 112M   | LibriTTS + VCTK + LJSpeech | 5M    | No         |
+| [bigvgan_base_22khz_80band](https://huggingface.co/nvidia/bigvgan_base_22khz_80band)                     | 22 kHz        | 80       | 8000  | 256              | 14M    | LibriTTS + VCTK + LJSpeech | 5M    | No         |
 
 The paper results are based on the original 24kHz BigVGAN models (`bigvgan_24khz_100band` and `bigvgan_base_24khz_100band`) trained on LibriTTS dataset.
 We also provide 22kHz BigVGAN models with band-limited setup (i.e., fmax=8000) for TTS applications.
-Note that the checkpoints use ``snakebeta`` activation with log scale parameterization, which have the best overall quality.
+Note that the checkpoints use `snakebeta` activation with log scale parameterization, which have the best overall quality.
+
+You can fine-tune the models by:
 
-You can fine-tune the models by downloading the checkpoints (both the generator weight and its discrimiantor/optimizer states) and resuming training using your audio dataset.
+1. downloading the checkpoints (both the generator weight and its discriminator/optimizer states)
+2. resuming training using your audio dataset by specifying `--checkpoint_path` that includes the checkpoints when launching `train.py`
 
 ## Training Details of BigVGAN-v2
+
 Comapred to the original BigVGAN, the pretrained checkpoints of BigVGAN-v2 used `batch_size=32` with a longer `segment_size=65536` and are trained using 8 A100 GPUs.
 
 Note that the BigVGAN-v2 `json` config files in `./configs` use `batch_size=4` as default to fit in a single A100 GPU for training. You can fine-tune the models adjusting `batch_size` depending on your GPUs.
@@ -147,23 +213,50 @@ Note that the BigVGAN-v2 `json` config files in `./configs` use `batch_size=4` a
 When training BigVGAN-v2 from scratch with small batch size, it can potentially encounter the early divergence problem mentioned in the paper. In such case, we recommend lowering the `clip_grad_norm` value (e.g. `100`) for the early training iterations (e.g. 20000 steps) and increase the value to the default `500`.
 
 ## Evaluation Results of BigVGAN-v2
+
 Below are the objective results of the 24kHz model (`bigvgan_v2_24khz_100band_256x`) obtained from the LibriTTS `dev` sets. BigVGAN-v2 shows noticeable improvements of the metrics. The model also exhibits reduced perceptual artifacts, especially for non-speech audio.
 
-|Model|Dataset|Steps|PESQ(↑)|M-STFT(↓)|MCD(↓)|Periodicity(↓)|V/UV F1(↑)|
-|-------|-----|-----|-----|-----|-----|-----|-----|
-|BigVGAN|LibriTTS|1M|4.027|0.7997|0.3745|0.1018|0.9598|
-|BigVGAN|LibriTTS|5M|4.256|0.7409|0.2988|0.0809|0.9698|
-|BigVGAN-v2|Large-scale Compilation|3M|**4.359**|**0.7134**|0.3060|**0.0621**|**0.9777**|
+| Model      | Dataset                 | Steps | PESQ(↑)   | M-STFT(↓)  | MCD(↓) | Periodicity(↓) | V/UV F1(↑) |
+|:----------:|:-----------------------:|:-----:|:---------:|:----------:|:------:|:--------------:|:----------:|
+| BigVGAN    | LibriTTS                | 1M    | 4.027     | 0.7997     | 0.3745 | 0.1018         | 0.9598     |
+| BigVGAN    | LibriTTS                | 5M    | 4.256     | 0.7409     | 0.2988 | 0.0809         | 0.9698     |
+| BigVGAN-v2 | Large-scale Compilation | 3M    | **4.359** | **0.7134** | 0.3060 | **0.0621**     | **0.9777** |
+
+## Speed Benchmark
+
+Below are the speed and VRAM usage benchmark results of BigVGAN from `tests/test_cuda_vs_torch_model.py`, using `bigvgan_v2_24khz_100band_256x` as a reference model.
+
+| GPU                        | num_mel_frame | use_cuda_kernel | Speed (kHz) | Real-time Factor | VRAM (GB) |
+|:--------------------------:|:-------------:|:---------------:|:-----------:|:----------------:|:---------:|
+| NVIDIA A100                | 256           | False           | 1672.1      | 69.7x            | 1.3       |
+|                            |               | True            | 3916.5      | 163.2x           | 1.3       |
+|                            | 2048          | False           | 1899.6      | 79.2x            | 1.7       |
+|                            |               | True            | 5330.1      | 222.1x           | 1.7       |
+|                            | 16384         | False           | 1973.8      | 82.2x            | 5.0       |
+|                            |               | True            | 5761.7      | 240.1x           | 4.4       |
+| NVIDIA GeForce RTX 3080    | 256           | False           | 841.1       | 35.0x            | 1.3       |
+|                            |               | True            | 1598.1      | 66.6x            | 1.3       |
+|                            | 2048          | False           | 929.9       | 38.7x            | 1.7       |
+|                            |               | True            | 1971.3      | 82.1x            | 1.6       |
+|                            | 16384         | False           | 943.4       | 39.3x            | 5.0       |
+|                            |               | True            | 2026.5      | 84.4x            | 3.9       |
+| NVIDIA GeForce RTX 2080 Ti | 256           | False           | 515.6       | 21.5x            | 1.3       |
+|                            |               | True            | 811.3       | 33.8x            | 1.3       |
+|                            | 2048          | False           | 576.5       | 24.0x            | 1.7       |
+|                            |               | True            | 1023.0      | 42.6x            | 1.5       |
+|                            | 16384         | False           | 589.4       | 24.6x            | 5.0       |
+|                            |               | True            | 1068.1      | 44.5x            | 3.2       |
 
 ## Acknowledgements
+
 We thank Vijay Anand Korthikanti and Kevin J. Shih for their generous support in implementing the CUDA kernel for inference.
 
 ## References
-* [HiFi-GAN](https://github.com/jik876/hifi-gan) (for generator and multi-period discriminator)
-* [Snake](https://github.com/EdwardDixon/snake) (for periodic activation)
-* [Alias-free-torch](https://github.com/junjun3518/alias-free-torch) (for anti-aliasing)
-* [Julius](https://github.com/adefossez/julius) (for low-pass filter)
-* [UnivNet](https://github.com/mindslab-ai/univnet) (for multi-resolution discriminator)
-* [descript-audio-codec](https://github.com/descriptinc/descript-audio-codec) and [vocos](https://github.com/gemelo-ai/vocos) (for multi-band multi-scale STFT discriminator and multi-scale mel spectrogram loss)
-* [Amphion](https://github.com/open-mmlab/Amphion) (for multi-scale sub-band CQT discriminator)
 
+- [HiFi-GAN](https://github.com/jik876/hifi-gan) (for generator and multi-period discriminator)
+- [Snake](https://github.com/EdwardDixon/snake) (for periodic activation)
+- [Alias-free-torch](https://github.com/junjun3518/alias-free-torch) (for anti-aliasing)
+- [Julius](https://github.com/adefossez/julius) (for low-pass filter)
+- [UnivNet](https://github.com/mindslab-ai/univnet) (for multi-resolution discriminator)
+- [descript-audio-codec](https://github.com/descriptinc/descript-audio-codec) and [vocos](https://github.com/gemelo-ai/vocos) (for multi-band multi-scale STFT discriminator and multi-scale mel spectrogram loss)
+- [Amphion](https://github.com/open-mmlab/Amphion) (for multi-scale sub-band CQT discriminator)

{alias_free_cuda → alias_free_activation/cuda}/activation1d.py

@@ -3,36 +3,45 @@
 
 import torch
 import torch.nn as nn
-from alias_free_torch.resample import UpSample1d, DownSample1d
+from alias_free_activation.torch.resample import UpSample1d, DownSample1d
+
 # load fused CUDA kernel: this enables importing anti_alias_activation_cuda
-from alias_free_cuda import load
-load.load()
+from alias_free_activation.cuda import load
+
+anti_alias_activation_cuda = load.load()
+
 
 class FusedAntiAliasActivation(torch.autograd.Function):
     """
-    Assumes filter size 12, replication padding on upsampling, and logscale alpha/beta parameters as inputs
+    Assumes filter size 12, replication padding on upsampling/downsampling, and logscale alpha/beta parameters as inputs.
+    The hyperparameters are hard-coded in the kernel to maximize speed.
+    NOTE: The fused kenrel is incorrect for Activation1d with different hyperparameters.
     """
+
     @staticmethod
-    def forward(ctx, inputs, ftr, alpha, beta):
-        import anti_alias_activation_cuda
-        activation_results = anti_alias_activation_cuda.forward(inputs, ftr, alpha, beta)
+    def forward(ctx, inputs, up_ftr, down_ftr, alpha, beta):
+        activation_results = anti_alias_activation_cuda.forward(
+            inputs, up_ftr, down_ftr, alpha, beta
+        )
+
         return activation_results
 
     @staticmethod
     def backward(ctx, output_grads):
-        # TODO: implement bwd pass
         raise NotImplementedError
         return output_grads, None, None
 
+
 class Activation1d(nn.Module):
-    def __init__(self,
-                 activation,
-                 up_ratio: int = 2,
-                 down_ratio: int = 2,
-                 up_kernel_size: int = 12,
-                 down_kernel_size: int = 12,
-                 fused: bool = True
-                 ):
+    def __init__(
+        self,
+        activation,
+        up_ratio: int = 2,
+        down_ratio: int = 2,
+        up_kernel_size: int = 12,
+        down_kernel_size: int = 12,
+        fused: bool = True,
+    ):
         super().__init__()
         self.up_ratio = up_ratio
         self.down_ratio = down_ratio
@@ -40,8 +49,7 @@ class Activation1d(nn.Module):
         self.upsample = UpSample1d(up_ratio, up_kernel_size)
         self.downsample = DownSample1d(down_ratio, down_kernel_size)
 
-        self.fused = fused # whether to use fused CUDA kernel or not
-
+        self.fused = fused  # Whether to use fused CUDA kernel or not
 
     def forward(self, x):
         if not self.fused:
@@ -51,13 +59,19 @@ class Activation1d(nn.Module):
             return x
         else:
             if self.act.__class__.__name__ == "Snake":
-                beta = self.act.alpha.data # snake uses same params for alpha and beta
+                beta = self.act.alpha.data  # Snake uses same params for alpha and beta
             else:
-                beta = self.act.beta.data # snakebeta uses different params for alpha and beta
+                beta = (
+                    self.act.beta.data
+                )  # Snakebeta uses different params for alpha and beta
             alpha = self.act.alpha.data
-            if not self.act.alpha_logscale: # exp baked into cuda kernel, cancel it out with a log
+            if (
+                not self.act.alpha_logscale
+            ):  # Exp baked into cuda kernel, cancel it out with a log
                 alpha = torch.log(alpha)
                 beta = torch.log(beta)
-            x = FusedAntiAliasActivation.apply(x, self.upsample.filter, alpha, beta)
-            x = self.downsample(x)
+
+            x = FusedAntiAliasActivation.apply(
+                x, self.upsample.filter, self.downsample.lowpass.filter, alpha, beta
+            )
             return x

{alias_free_cuda → alias_free_activation/cuda}/anti_alias_activation.cpp

@@ -14,35 +14,10 @@
  * limitations under the License.
  */
 
-#include <cuda_fp16.h>
-#include <torch/extension.h>
-#include <vector>
+ #include <torch/extension.h>
 
-namespace anti_alias_activation {
-
-torch::Tensor fwd_cuda(torch::Tensor const& input,
-                       torch::Tensor const& filter,
-                       torch::Tensor const& alpha,
-		       torch::Tensor const& beta
-		       );
-
-torch::Tensor fwd(torch::Tensor const& input,
-                  torch::Tensor const& filter,
-                  torch::Tensor const& alpha,
-		  torch::Tensor const& beta
-		  ) {
-  AT_ASSERTM(input.dim() == 3, "expected 3D tensor");
-  //AT_ASSERTM((input.scalar_type() == at::ScalarType::Half) ||
-  //	     (input.scalar_type() == at::ScalarType::BFloat16),
-  //    "Only fp16 and bf16 are supported");
-
-  return fwd_cuda(input, filter, alpha, beta);
-}
-
-} // end namespace anti_alias_activation
+extern "C" torch::Tensor fwd_cuda(torch::Tensor const &input, torch::Tensor const &up_filter, torch::Tensor const &down_filter, torch::Tensor const &alpha, torch::Tensor const &beta);
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
-  m.def("forward",
-        &anti_alias_activation::fwd,
-	"Anti Alias Activation -- Forward.");
-}
+    m.def("forward", &fwd_cuda, "Anti-Alias Activation forward (CUDA)");
+}

alias_free_activation/cuda/anti_alias_activation_cuda.cu

@@ -0,0 +1,246 @@
+/* coding=utf-8
+ * Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <ATen/ATen.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+#include "type_shim.h"
+#include <assert.h>
+#include <cfloat>
+#include <limits>
+#include <stdint.h>
+#include <c10/macros/Macros.h>
+
+namespace
+{
+    // Hard-coded hyperparameters
+    // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and
+    constexpr int ELEMENTS_PER_LDG_STG = 1; //(WARP_ITERATIONS < 4) ? 1 : 4;
+    constexpr int BUFFER_SIZE = 32;
+    constexpr int FILTER_SIZE = 12;
+    constexpr int HALF_FILTER_SIZE = 6;
+    constexpr int UPSAMPLE_REPLICATION_PAD = 5; // 5 on each side, matching torch impl
+    constexpr int DOWNSAMPLE_REPLICATION_PAD_LEFT = 5; // matching torch impl
+    constexpr int DOWNSAMPLE_REPLICATION_PAD_RIGHT = 6; // matching torch impl
+
+    template <typename input_t, typename output_t, typename acc_t>
+    __global__ void anti_alias_activation_forward(
+        output_t *dst,
+        const input_t *src,
+        const input_t *up_ftr,
+        const input_t *down_ftr,
+        const input_t *alpha,
+        const input_t *beta,
+        int batch_size,
+        int channels,
+        int seq_len)
+    {
+        // Up and downsample filters
+        input_t up_filter[FILTER_SIZE];
+        input_t down_filter[FILTER_SIZE];
+
+        // Load data from global memory including extra indices reserved for replication paddings
+        input_t elements[2 * FILTER_SIZE + 2 * BUFFER_SIZE + 2 * UPSAMPLE_REPLICATION_PAD] = {0};
+        input_t intermediates[2 * FILTER_SIZE + 2 * BUFFER_SIZE + DOWNSAMPLE_REPLICATION_PAD_LEFT + DOWNSAMPLE_REPLICATION_PAD_RIGHT] = {0};
+
+        // Output stores downsampled output before writing to dst
+        output_t output[BUFFER_SIZE];
+
+        // blockDim/threadIdx = (128, 1, 1)
+        // gridDim/blockIdx = (seq_blocks, channels, batches)
+        int block_offset = (blockIdx.x * 128 * BUFFER_SIZE + seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
+        int local_offset = threadIdx.x * BUFFER_SIZE;
+        int seq_offset = blockIdx.x * 128 * BUFFER_SIZE + local_offset;
+
+        // intermediate have double the seq_len
+        int intermediate_local_offset = threadIdx.x * BUFFER_SIZE * 2;
+        int intermediate_seq_offset = blockIdx.x * 128 * BUFFER_SIZE * 2 + intermediate_local_offset;
+
+        // Get values needed for replication padding before moving pointer
+        const input_t *right_most_pntr = src + (seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
+        input_t seq_left_most_value = right_most_pntr[0];
+        input_t seq_right_most_value = right_most_pntr[seq_len - 1];
+
+        // Move src and dst pointers
+        src += block_offset + local_offset;
+        dst += block_offset + local_offset;
+
+        // Alpha and beta values for snake activatons. Applies exp by default
+        alpha = alpha + blockIdx.y;
+        input_t alpha_val = expf(alpha[0]);
+        beta = beta + blockIdx.y;
+        input_t beta_val = expf(beta[0]);
+
+        #pragma unroll
+        for (int it = 0; it < FILTER_SIZE; it += 1)
+        {
+            up_filter[it] = up_ftr[it];
+            down_filter[it] = down_ftr[it];
+        }
+
+        // Apply replication padding for upsampling, matching torch impl
+        #pragma unroll
+        for (int it = -HALF_FILTER_SIZE; it < BUFFER_SIZE + HALF_FILTER_SIZE; it += 1)
+        {
+            int element_index = seq_offset + it; // index for element
+            if ((element_index < 0) && (element_index >= -UPSAMPLE_REPLICATION_PAD))
+            {
+                elements[2 * (HALF_FILTER_SIZE + it)] = 2 * seq_left_most_value;
+            }
+            if ((element_index >= seq_len) && (element_index < seq_len + UPSAMPLE_REPLICATION_PAD))
+            {
+                elements[2 * (HALF_FILTER_SIZE + it)] = 2 * seq_right_most_value;
+            }
+            if ((element_index >= 0) && (element_index < seq_len))
+            {
+                elements[2 * (HALF_FILTER_SIZE + it)] = 2 * src[it];
+            }
+        }
+
+        // Apply upsampling strided convolution and write to intermediates. It reserves DOWNSAMPLE_REPLICATION_PAD_LEFT for replication padding of the downsampilng conv later
+        #pragma unroll
+        for (int it = 0; it < (2 * BUFFER_SIZE + 2 * FILTER_SIZE); it += 1)
+        {
+            input_t acc = 0.0;
+            int element_index = intermediate_seq_offset + it; // index for intermediate
+            #pragma unroll
+            for (int f_idx = 0; f_idx < FILTER_SIZE; f_idx += 1)
+            {
+                if ((element_index + f_idx) >= 0)
+                {
+                    acc += up_filter[f_idx] * elements[it + f_idx];
+                }
+            }
+            intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] = acc;
+        }
+
+        // Apply activation function. It reserves DOWNSAMPLE_REPLICATION_PAD_LEFT and DOWNSAMPLE_REPLICATION_PAD_RIGHT for replication padding of the downsampilng conv later
+        double no_div_by_zero = 0.000000001;
+        #pragma unroll
+        for (int it = 0; it < 2 * BUFFER_SIZE + 2 * FILTER_SIZE; it += 1)
+        {
+            intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] += (1.0 / (beta_val + no_div_by_zero)) * sinf(intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] * alpha_val) * sinf(intermediates[it + DOWNSAMPLE_REPLICATION_PAD_LEFT] * alpha_val);
+        }
+
+        // Apply replication padding before downsampling conv from intermediates
+        #pragma unroll
+        for (int it = 0; it < DOWNSAMPLE_REPLICATION_PAD_LEFT; it += 1)
+        {
+            intermediates[it] = intermediates[DOWNSAMPLE_REPLICATION_PAD_LEFT];
+        }
+        #pragma unroll
+        for (int it = DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE; it < DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE + DOWNSAMPLE_REPLICATION_PAD_RIGHT; it += 1)
+        {
+            intermediates[it] = intermediates[DOWNSAMPLE_REPLICATION_PAD_LEFT + 2 * BUFFER_SIZE + 2 * FILTER_SIZE - 1];
+        }
+
+        // Apply downsample strided convolution (assuming stride=2) from intermediates
+        #pragma unroll
+        for (int it = 0; it < BUFFER_SIZE; it += 1)
+        {
+            input_t acc = 0.0;
+            #pragma unroll
+            for (int f_idx = 0; f_idx < FILTER_SIZE; f_idx += 1)
+            {
+                // Add constant DOWNSAMPLE_REPLICATION_PAD_RIGHT to match torch implementation
+                acc += down_filter[f_idx] * intermediates[it * 2 + f_idx + DOWNSAMPLE_REPLICATION_PAD_RIGHT];
+            }
+            output[it] = acc;
+        }
+
+        // Write output to dst
+        #pragma unroll
+        for (int it = 0;  it < BUFFER_SIZE;  it += ELEMENTS_PER_LDG_STG)
+        {
+            int element_index = seq_offset + it;
+            if (element_index < seq_len)
+            {
+                dst[it] = output[it];
+            }
+        }
+
+    }
+
+    template <typename input_t, typename output_t, typename acc_t>
+    void dispatch_anti_alias_activation_forward(
+        output_t *dst,
+        const input_t *src,
+        const input_t *up_ftr,
+        const input_t *down_ftr,
+        const input_t *alpha,
+        const input_t *beta,
+        int batch_size,
+        int channels,
+        int seq_len)
+    {
+        if (seq_len == 0)
+        {
+            return;
+        }
+        else
+        {
+            // Use 128 threads per block to maximimize gpu utilization
+            constexpr int threads_per_block = 128;
+            constexpr int seq_len_per_block = 4096;
+            int blocks_per_seq_len = (seq_len + seq_len_per_block - 1) / seq_len_per_block;
+            dim3 blocks(blocks_per_seq_len, channels, batch_size);
+            dim3 threads(threads_per_block, 1, 1);
+
+            anti_alias_activation_forward<input_t, output_t, acc_t>
+                <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, up_ftr, down_ftr, alpha, beta, batch_size, channels, seq_len);
+        }
+    }
+}
+
+extern "C" torch::Tensor fwd_cuda(torch::Tensor const &input, torch::Tensor const &up_filter, torch::Tensor const &down_filter, torch::Tensor const &alpha, torch::Tensor const &beta)
+{
+    // Input is a 3d tensor with dimensions [batches, channels, seq_len]
+    const int batches = input.size(0);
+    const int channels = input.size(1);
+    const int seq_len = input.size(2);
+
+    // Output
+    auto act_options = input.options().requires_grad(false);
+
+    torch::Tensor anti_alias_activation_results =
+        torch::empty({batches, channels, seq_len}, act_options);
+
+    void *input_ptr = static_cast<void *>(input.data_ptr());
+    void *up_filter_ptr = static_cast<void *>(up_filter.data_ptr());
+    void *down_filter_ptr = static_cast<void *>(down_filter.data_ptr());
+    void *alpha_ptr = static_cast<void *>(alpha.data_ptr());
+    void *beta_ptr = static_cast<void *>(beta.data_ptr());
+    void *anti_alias_activation_results_ptr = static_cast<void *>(anti_alias_activation_results.data_ptr());
+
+    DISPATCH_FLOAT_HALF_AND_BFLOAT(
+        input.scalar_type(),
+        "dispatch anti alias activation_forward",
+        dispatch_anti_alias_activation_forward<scalar_t, scalar_t, float>(
+            reinterpret_cast<scalar_t *>(anti_alias_activation_results_ptr),
+            reinterpret_cast<const scalar_t *>(input_ptr),
+            reinterpret_cast<const scalar_t *>(up_filter_ptr),
+            reinterpret_cast<const scalar_t *>(down_filter_ptr),
+            reinterpret_cast<const scalar_t *>(alpha_ptr),
+            reinterpret_cast<const scalar_t *>(beta_ptr),
+            batches,
+            channels,
+            seq_len););
+    return anti_alias_activation_results;
+}

{alias_free_cuda → alias_free_activation/cuda}/compat.h

@@ -18,8 +18,6 @@
  *     https://github.com/NVIDIA/apex
  *     with minor changes. */
 
-
-
 #ifndef TORCH_CHECK
 #define TORCH_CHECK AT_CHECK
 #endif

{alias_free_cuda → alias_free_activation/cuda}/load.py

@@ -7,26 +7,24 @@ import subprocess
 
 from torch.utils import cpp_extension
 
-# Setting this param to a list has a problem of generating different
-# compilation commands (with diferent order of architectures) and
-# leading to recompilation of fused kernels. Set it to empty string
-# to avoid recompilation and assign arch flags explicity in
-# extra_cuda_cflags below
+"""
+Setting this param to a list has a problem of generating different compilation commands (with diferent order of architectures) and leading to recompilation of fused kernels. 
+Set it to empty stringo avoid recompilation and assign arch flags explicity in extra_cuda_cflags below
+"""
 os.environ["TORCH_CUDA_ARCH_LIST"] = ""
 
 
 def load():
     # Check if cuda 11 is installed for compute capability 8.0
     cc_flag = []
-    _, bare_metal_major, _ = _get_cuda_bare_metal_version(
-        cpp_extension.CUDA_HOME)
+    _, bare_metal_major, _ = _get_cuda_bare_metal_version(cpp_extension.CUDA_HOME)
     if int(bare_metal_major) >= 11:
-        cc_flag.append('-gencode')
-        cc_flag.append('arch=compute_80,code=sm_80')
+        cc_flag.append("-gencode")
+        cc_flag.append("arch=compute_80,code=sm_80")
 
     # Build path
     srcpath = pathlib.Path(__file__).parent.absolute()
-    buildpath = srcpath / 'build'
+    buildpath = srcpath / "build"
     _create_build_dir(buildpath)
 
     # Helper function to build the kernels.
@@ -35,26 +33,42 @@ def load():
             name=name,
             sources=sources,
             build_directory=buildpath,
-            extra_cflags=['-O3',],
-            extra_cuda_cflags=['-O3',
-                               '-gencode', 'arch=compute_70,code=sm_70',
-                               '--use_fast_math'] + extra_cuda_flags + cc_flag,
-            verbose=True
+            extra_cflags=[
+                "-O3",
+            ],
+            extra_cuda_cflags=[
+                "-O3",
+                "-gencode",
+                "arch=compute_70,code=sm_70",
+                "--use_fast_math",
+            ]
+            + extra_cuda_flags
+            + cc_flag,
+            verbose=True,
         )
 
-    extra_cuda_flags = ['-U__CUDA_NO_HALF_OPERATORS__',
-                        '-U__CUDA_NO_HALF_CONVERSIONS__',
-                        '--expt-relaxed-constexpr',
-                        '--expt-extended-lambda']
-    
-    sources=[srcpath / 'anti_alias_activation.cpp',
-             srcpath / 'anti_alias_activation_cuda.cu']
+    extra_cuda_flags = [
+        "-U__CUDA_NO_HALF_OPERATORS__",
+        "-U__CUDA_NO_HALF_CONVERSIONS__",
+        "--expt-relaxed-constexpr",
+        "--expt-extended-lambda",
+    ]
+
+    sources = [
+        srcpath / "anti_alias_activation.cpp",
+        srcpath / "anti_alias_activation_cuda.cu",
+    ]
     anti_alias_activation_cuda = _cpp_extention_load_helper(
-        "anti_alias_activation_cuda", sources, extra_cuda_flags)
+        "anti_alias_activation_cuda", sources, extra_cuda_flags
+    )
+
+    return anti_alias_activation_cuda
+
 
 def _get_cuda_bare_metal_version(cuda_dir):
-    raw_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"],
-                                         universal_newlines=True)
+    raw_output = subprocess.check_output(
+        [cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True
+    )
     output = raw_output.split()
     release_idx = output.index("release") + 1
     release = output[release_idx].split(".")
@@ -69,4 +83,4 @@ def _create_build_dir(buildpath):
         os.mkdir(buildpath)
     except OSError:
         if not os.path.isdir(buildpath):
-            print(f"Creation of the build directory {buildpath} failed")
+            print(f"Creation of the build directory {buildpath} failed")

alias_free_activation/cuda/type_shim.h

@@ -0,0 +1,92 @@
+/* coding=utf-8
+ * Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <ATen/ATen.h>
+#include "compat.h"
+
+#define DISPATCH_FLOAT_HALF_AND_BFLOAT(TYPE, NAME, ...)                 \
+	switch (TYPE)                                                       \
+	{                                                                   \
+	case at::ScalarType::Float:                                         \
+	{                                                                   \
+		using scalar_t = float;                                         \
+		__VA_ARGS__;                                                    \
+		break;                                                          \
+	}                                                                   \
+	case at::ScalarType::Half:                                          \
+	{                                                                   \
+		using scalar_t = at::Half;                                      \
+		__VA_ARGS__;                                                    \
+		break;                                                          \
+	}                                                                   \
+	case at::ScalarType::BFloat16:                                      \
+	{                                                                   \
+		using scalar_t = at::BFloat16;                                  \
+		__VA_ARGS__;                                                    \
+		break;                                                          \
+	}                                                                   \
+	default:                                                            \
+		AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'"); \
+	}
+
+#define DISPATCH_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(TYPEIN, TYPEOUT, NAME, ...) \
+	switch (TYPEIN)                                                            \
+	{                                                                          \
+	case at::ScalarType::Float:                                                \
+	{                                                                          \
+		using scalar_t_in = float;                                             \
+		switch (TYPEOUT)                                                       \
+		{                                                                      \
+		case at::ScalarType::Float:                                            \
+		{                                                                      \
+			using scalar_t_out = float;                                        \
+			__VA_ARGS__;                                                       \
+			break;                                                             \
+		}                                                                      \
+		case at::ScalarType::Half:                                             \
+		{                                                                      \
+			using scalar_t_out = at::Half;                                     \
+			__VA_ARGS__;                                                       \
+			break;                                                             \
+		}                                                                      \
+		case at::ScalarType::BFloat16:                                         \
+		{                                                                      \
+			using scalar_t_out = at::BFloat16;                                 \
+			__VA_ARGS__;                                                       \
+			break;                                                             \
+		}                                                                      \
+		default:                                                               \
+			AT_ERROR(#NAME, " not implemented for '", toString(TYPEOUT), "'"); \
+		}                                                                      \
+		break;                                                                 \
+	}                                                                          \
+	case at::ScalarType::Half:                                                 \
+	{                                                                          \
+		using scalar_t_in = at::Half;                                          \
+		using scalar_t_out = at::Half;                                         \
+		__VA_ARGS__;                                                           \
+		break;                                                                 \
+	}                                                                          \
+	case at::ScalarType::BFloat16:                                             \
+	{                                                                          \
+		using scalar_t_in = at::BFloat16;                                      \
+		using scalar_t_out = at::BFloat16;                                     \
+		__VA_ARGS__;                                                           \
+		break;                                                                 \
+	}                                                                          \
+	default:                                                                   \
+		AT_ERROR(#NAME, " not implemented for '", toString(TYPEIN), "'");      \
+	}

{alias_free_torch → alias_free_activation/torch}/__init__.py

@@ -3,4 +3,4 @@
 
 from .filter import *
 from .resample import *
-from .act import *
+from .act import *

{alias_free_torch → alias_free_activation/torch}/act.py

@@ -2,16 +2,18 @@
 #   LICENSE is in incl_licenses directory.
 
 import torch.nn as nn
-from .resample import UpSample1d, DownSample1d
+from alias_free_activation.torch.resample import UpSample1d, DownSample1d
 
 
 class Activation1d(nn.Module):
-    def __init__(self,
-                 activation,
-                 up_ratio: int = 2,
-                 down_ratio: int = 2,
-                 up_kernel_size: int = 12,
-                 down_kernel_size: int = 12):
+    def __init__(
+        self,
+        activation,
+        up_ratio: int = 2,
+        down_ratio: int = 2,
+        up_kernel_size: int = 12,
+        down_kernel_size: int = 12,
+    ):
         super().__init__()
         self.up_ratio = up_ratio
         self.down_ratio = down_ratio
@@ -25,4 +27,4 @@ class Activation1d(nn.Module):
         x = self.act(x)
         x = self.downsample(x)
 
-        return x
+        return x

{alias_free_torch → alias_free_activation/torch}/filter.py

@@ -6,7 +6,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 import math
 
-if 'sinc' in dir(torch):
+if "sinc" in dir(torch):
     sinc = torch.sinc
 else:
     # This code is adopted from adefossez's julius.core.sinc under the MIT License
@@ -17,40 +17,45 @@ else:
         Implementation of sinc, i.e. sin(pi * x) / (pi * x)
         __Warning__: Different to julius.sinc, the input is multiplied by `pi`!
         """
-        return torch.where(x == 0,
-                           torch.tensor(1., device=x.device, dtype=x.dtype),
-                           torch.sin(math.pi * x) / math.pi / x)
+        return torch.where(
+            x == 0,
+            torch.tensor(1.0, device=x.device, dtype=x.dtype),
+            torch.sin(math.pi * x) / math.pi / x,
+        )
 
 
 # This code is adopted from adefossez's julius.lowpass.LowPassFilters under the MIT License
 # https://adefossez.github.io/julius/julius/lowpass.html
 #   LICENSE is in incl_licenses directory.
-def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): # return filter [1,1,kernel_size]
-    even = (kernel_size % 2 == 0)
+def kaiser_sinc_filter1d(
+    cutoff, half_width, kernel_size
+):  # return filter [1,1,kernel_size]
+    even = kernel_size % 2 == 0
     half_size = kernel_size // 2
 
-    #For kaiser window
+    # For kaiser window
     delta_f = 4 * half_width
     A = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95
-    if A > 50.:
+    if A > 50.0:
         beta = 0.1102 * (A - 8.7)
-    elif A >= 21.:
-        beta = 0.5842 * (A - 21)**0.4 + 0.07886 * (A - 21.)
+    elif A >= 21.0:
+        beta = 0.5842 * (A - 21) ** 0.4 + 0.07886 * (A - 21.0)
     else:
-        beta = 0.
+        beta = 0.0
     window = torch.kaiser_window(kernel_size, beta=beta, periodic=False)
 
     # ratio = 0.5/cutoff -> 2 * cutoff = 1 / ratio
     if even:
-        time = (torch.arange(-half_size, half_size) + 0.5)
+        time = torch.arange(-half_size, half_size) + 0.5
     else:
         time = torch.arange(kernel_size) - half_size
     if cutoff == 0:
         filter_ = torch.zeros_like(time)
     else:
         filter_ = 2 * cutoff * window * sinc(2 * cutoff * time)
-        # Normalize filter to have sum = 1, otherwise we will have a small leakage
-        # of the constant component in the input signal.
+        """
+        Normalize filter to have sum = 1, otherwise we will have a small leakage of the constant component in the input signal.
+        """
         filter_ /= filter_.sum()
         filter = filter_.view(1, 1, kernel_size)
 
@@ -58,22 +63,25 @@ def kaiser_sinc_filter1d(cutoff, half_width, kernel_size): # return filter [1,1,
 
 
 class LowPassFilter1d(nn.Module):
-    def __init__(self,
-                 cutoff=0.5,
-                 half_width=0.6,
-                 stride: int = 1,
-                 padding: bool = True,
-                 padding_mode: str = 'replicate',
-                 kernel_size: int = 12):
-        # kernel_size should be even number for stylegan3 setup,
-        # in this implementation, odd number is also possible.
+    def __init__(
+        self,
+        cutoff=0.5,
+        half_width=0.6,
+        stride: int = 1,
+        padding: bool = True,
+        padding_mode: str = "replicate",
+        kernel_size: int = 12,
+    ):
+        """
+        kernel_size should be even number for stylegan3 setup, in this implementation, odd number is also possible.
+        """
         super().__init__()
-        if cutoff < -0.:
+        if cutoff < -0.0:
             raise ValueError("Minimum cutoff must be larger than zero.")
         if cutoff > 0.5:
             raise ValueError("A cutoff above 0.5 does not make sense.")
         self.kernel_size = kernel_size
-        self.even = (kernel_size % 2 == 0)
+        self.even = kernel_size % 2 == 0
         self.pad_left = kernel_size // 2 - int(self.even)
         self.pad_right = kernel_size // 2
         self.stride = stride
@@ -82,14 +90,12 @@ class LowPassFilter1d(nn.Module):
         filter = kaiser_sinc_filter1d(cutoff, half_width, kernel_size)
         self.register_buffer("filter", filter)
 
-    #input [B, C, T]
+    # Input [B, C, T]
     def forward(self, x):
         _, C, _ = x.shape
 
         if self.padding:
-            x = F.pad(x, (self.pad_left, self.pad_right),
-                      mode=self.padding_mode)
-        out = F.conv1d(x, self.filter.expand(C, -1, -1),
-                       stride=self.stride, groups=C)
+            x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode)
+        out = F.conv1d(x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
 
-        return out
+        return out

{alias_free_torch → alias_free_activation/torch}/resample.py

@@ -3,32 +3,37 @@
 
 import torch.nn as nn
 from torch.nn import functional as F
-from .filter import LowPassFilter1d
-from .filter import kaiser_sinc_filter1d
+from alias_free_activation.torch.filter import LowPassFilter1d
+from alias_free_activation.torch.filter import kaiser_sinc_filter1d
 
 
 class UpSample1d(nn.Module):
     def __init__(self, ratio=2, kernel_size=None):
         super().__init__()
         self.ratio = ratio
-        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
         self.stride = ratio
         self.pad = self.kernel_size // ratio - 1
         self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
-        self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
-        filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio,
-                                      half_width=0.6 / ratio,
-                                      kernel_size=self.kernel_size)
+        self.pad_right = (
+            self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
+        )
+        filter = kaiser_sinc_filter1d(
+            cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size
+        )
         self.register_buffer("filter", filter)
 
     # x: [B, C, T]
     def forward(self, x):
         _, C, _ = x.shape
 
-        x = F.pad(x, (self.pad, self.pad), mode='replicate')
+        x = F.pad(x, (self.pad, self.pad), mode="replicate")
         x = self.ratio * F.conv_transpose1d(
-            x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C)
-        x = x[..., self.pad_left:-self.pad_right]
+            x, self.filter.expand(C, -1, -1), stride=self.stride, groups=C
+        )
+        x = x[..., self.pad_left : -self.pad_right]
 
         return x
 
@@ -37,13 +42,17 @@ class DownSample1d(nn.Module):
     def __init__(self, ratio=2, kernel_size=None):
         super().__init__()
         self.ratio = ratio
-        self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
-        self.lowpass = LowPassFilter1d(cutoff=0.5 / ratio,
-                                       half_width=0.6 / ratio,
-                                       stride=ratio,
-                                       kernel_size=self.kernel_size)
+        self.kernel_size = (
+            int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
+        )
+        self.lowpass = LowPassFilter1d(
+            cutoff=0.5 / ratio,
+            half_width=0.6 / ratio,
+            stride=ratio,
+            kernel_size=self.kernel_size,
+        )
 
     def forward(self, x):
         xx = self.lowpass(x)
 
-        return xx
+        return xx

alias_free_cuda/anti_alias_activation_cuda.cu

@@ -1,314 +0,0 @@
-/* coding=utf-8
- * Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include <ATen/ATen.h>
-#include <cuda.h>
-#include <cuda_runtime.h>
-#include <cuda_fp16.h>
-#include <cuda_profiler_api.h>
-#include <ATen/cuda/CUDAContext.h>
-#include <torch/extension.h>
-#include "type_shim.h"
-#include <assert.h>
-#include <cfloat>
-#include <limits>
-#include <stdint.h>
-#include <c10/macros/Macros.h>
-
-namespace {
-
-    /*
-template <typename Datatype, int ELEMENTS_PER_LDG>
-__device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);
-
-template <>
-__device__ __inline__ void copy_vector<c10::BFloat16, 1>(c10::BFloat16 *dst, const c10::BFloat16 *src) { *dst = *src; }
-
-template <>
-__device__ __inline__ void copy_vector<c10::BFloat16, 4>(c10::BFloat16 *dst, const c10::BFloat16 *src) { *((float2*) dst) = *((float2*) src); }
-
-template <>
-__device__ __inline__ void copy_vector<c10::Half, 1>(c10::Half *dst, const c10::Half *src) { *dst = *src; }
-
-template <>
-__device__ __inline__ void copy_vector<c10::Half, 4>(c10::Half *dst, const c10::Half *src) { *((float2*) dst) = *((float2*) src); }
-
-template <>
-__device__ __inline__ void copy_vector<uint8_t, 1>(uint8_t *dst, const uint8_t *src) { *dst = *src; }
-
-template <>
-__device__ __inline__ void copy_vector<uint8_t, 4>(uint8_t *dst, const uint8_t *src) {*((half2*) dst) = *((half2*) src); }
-
-int log2_ceil(int value) {
-    int log2_value = 0;
-    while ((1 << log2_value) < value) ++log2_value;
-    return log2_value;
-}
-
-template<typename T>
-struct Add {
-  __device__ __forceinline__ T operator()(T a, T b) const {
-    return a + b;
-  }
-};
-
-template<typename T>
-struct Max {
-  __device__ __forceinline__ T operator()(T a, T b) const {
-    return a < b ? b : a;
-  }
-};
-
-template <typename T>
-__device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff)
-{
-#if CUDA_VERSION >= 9000
-    return __shfl_xor_sync(mask, value, laneMask, width);
-#else
-    return __shfl_xor(value, laneMask, width);
-#endif
-}
-
-template <typename acc_t, int WARP_BATCH, int WARP_SIZE, template<typename> class ReduceOp>
-__device__ __forceinline__ void warp_reduce(acc_t* sum) {
-    ReduceOp<acc_t> r;
-    #pragma unroll
-    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
-        #pragma unroll
-        for (int i = 0;  i < WARP_BATCH;  ++i) {
-            acc_t b = WARP_SHFL_XOR_NATIVE(sum[i], offset, WARP_SIZE);
-            sum[i] = r(sum[i], b);
-        }
-    }
-}
-*/
-
-template <typename input_t, typename output_t, typename acc_t>
-__global__ void anti_alias_activation_forward(
-    output_t *dst,
-    const input_t *src,
-    const input_t *ftr,
-    const input_t *alpha,
-    const input_t *beta,
-    int batch_size,
-    int channels,
-    int seq_len)
-{
-    // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and
-    constexpr int ELEMENTS_PER_LDG_STG = 1; //(WARP_ITERATIONS < 4) ? 1 : 4;
-    constexpr int BUFFER_SIZE = 32;
-    constexpr int FILTER_SIZE = 12;
-    constexpr int HALF_FILTER_SIZE = 6;
-    constexpr int REPLICATION_PAD = 5; // 5 on each side
-
-    // blockDim/threadIdx = (128, 1, 1)
-    // gridDim/blockIdx = (seq_blocks, channels, batches)
-    int block_offset = (blockIdx.x * 128 * BUFFER_SIZE + seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
-    int local_offset = threadIdx.x * BUFFER_SIZE;
-    int seq_offset = blockIdx.x * 128 * BUFFER_SIZE + local_offset;
-
-
-    //int intermediate_seq_len = seq_len * 2 - 1 + 4 * REPLICATION_PAD;
-    //int intermediate_block_offset = (blockIdx.x * 128 * BUFFER_SIZE * 2 + intermediate_seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
-    //int intermediate_local_offset = threadIdx.x * BUFFER_SIZE * 2;
-
-    int output_seq_len = seq_len * 2 ; //
-    int output_block_offset = (blockIdx.x * 128 * BUFFER_SIZE * 2 + output_seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
-    int output_local_offset = threadIdx.x * BUFFER_SIZE * 2;
-    int output_seq_offset = blockIdx.x * 128 * BUFFER_SIZE *2 + output_local_offset;
-    // get values needed for replication padding before moving pointer
-    const input_t *right_most_pntr = src + (seq_len * (blockIdx.y + gridDim.y * blockIdx.z));
-    input_t seq_left_most_value = right_most_pntr[0];
-    input_t seq_right_most_value = right_most_pntr[seq_len - 1];
-
-    src += block_offset + local_offset;
-    dst += output_block_offset + output_local_offset  ;
-    alpha = alpha + blockIdx.y;
-    input_t alpha_val = expf(alpha[0]);
-    beta = beta + blockIdx.y;
-    input_t beta_val = expf(beta[0]);
-    // load data from global memory
-    input_t elements[2*FILTER_SIZE+2*BUFFER_SIZE] = {0};
-    input_t intermediates[2*FILTER_SIZE+2*BUFFER_SIZE] = {0};
-    //output_t output[2*BUFFER_SIZE];
-    input_t filter[FILTER_SIZE];
-    //input_t temp_data[ELEMENTS_PER_LDG_STG];
-    //uint8_t temp_mask[ELEMENTS_PER_LDG_STG];
-
-    #pragma unroll
-    for (int it = 0; it < FILTER_SIZE; it+=1) {
-        filter[it] = ftr[it];
-    }
-
-
-    #pragma unroll
-    for (int it = -HALF_FILTER_SIZE;  it < BUFFER_SIZE + HALF_FILTER_SIZE ;  it+=1) {
-        int element_index = seq_offset + it;
-	if ((element_index < 0) && (element_index >= -REPLICATION_PAD)) {
-	    elements[2*(HALF_FILTER_SIZE+it)] = 2*seq_left_most_value;
-	}
-	if ((element_index >= seq_len) && (element_index < seq_len + REPLICATION_PAD)) {
-	    elements[2*(HALF_FILTER_SIZE+it)] = 2*seq_right_most_value;
-	}
-        if ((element_index >= 0) && (element_index < seq_len)) {
-	  elements[2*(HALF_FILTER_SIZE+it)] = 2*src[it];
-        }
-    }
-
-
-
-    // apply filter
-    #pragma unroll
-    for (int it = 0;  it < (2 * BUFFER_SIZE + 2*FILTER_SIZE);  it+=1) {
-        input_t acc = 0.0;
-
-	int element_index = output_seq_offset + it; // index for output
-	#pragma unroll
-        for (int f_idx = 0; f_idx < FILTER_SIZE; f_idx+=1){
-	  if ((element_index + f_idx) >= 0){
-            acc += filter[f_idx] * elements[it+f_idx];
-	  }
-	}
-        intermediates[it] = acc;
-    }
-
-    double no_div_by_zero = 0.000000001;
-    #pragma unroll
-    for (int it = 0; it < 12 + 2 * BUFFER_SIZE; it++) {
-        intermediates[it] += (1.0/(beta_val + no_div_by_zero)) *  sinf(intermediates[it] * alpha_val) * sinf(intermediates[it] * alpha_val);
-    }
-
-
-    // now copy to output
-    #pragma unroll
-    for (int it = 0; it < 2*BUFFER_SIZE; it+=1){
-      int element_index = output_seq_offset + it;
-      if (element_index < output_seq_len) {
-	dst[it]  = intermediates[it+6];
-      }
-    }
-
-
-
-    // for (int it = 0;  it < BUFFER_SIZE;  it+=ELEMENTS_PER_LDG_STG) {
-    //     int element_index = seq_offset + it;
-    //     if (element_index < seq_len) {
-    //         dst[it] = output[it];
-    //     }
-    // }
-
-
-    // // Upsample convolution
-    // for (int it = 0;  it < 2 * BUFFER_SIZE + 12;  it+=1) {
-    //     input_t acc = 0.0;
-
-    //     for (int f_idx = 0; f_idx < FILTER_SIZE; f_idx+=1){
-    //         acc += filter[f_idx] * elements[it+f_idx];
-    //     }
-    //     intermediates[it] = acc;
-    // }
-
-    // // correct the corners of intermediates
-    // if (seq_offset == 0) {
-    //     for (int it = 0; it < 6; it+=1)
-    //         intermediates[it] = 0;
-    // }
-
-    // if (seq_offset + 32 >= seq_len) {
-    //     int offset = seq_len % 32 == 0 ? 32 : seq_len % 32;
-
-    //     for (int it = 0; it < 6; it++) {
-    //         intermediates[6+2*offset+it] = 0;
-    //     }
-    // }
-
-
-
-
-    // for (int it = 0;  it < BUFFER_SIZE;  it+=ELEMENTS_PER_LDG_STG) {
-    //     int element_index = seq_offset + it;
-    //     if (element_index < seq_len) {
-    //         dst[it] = output[it];
-    //     }
-    // }
-}
-
-template<typename input_t, typename output_t, typename acc_t>
-void dispatch_anti_alias_activation_forward(
-    output_t *dst,
-    const input_t *src,
-    const input_t *ftr,
-    const input_t *alpha,
-    const input_t *beta,
-    int batch_size,
-    int channels,
-    int seq_len)
-{
-    if (seq_len == 0) {
-        return;
-    } else {
-        // use 128 threads per block to maximimize gpu utilization
-        constexpr int threads_per_block = 128;
-        constexpr int seq_len_per_block = 4096;
-        int blocks_per_seq_len = (seq_len + seq_len_per_block - 1) / seq_len_per_block;
-        dim3 blocks(blocks_per_seq_len, channels, batch_size);
-        dim3 threads(threads_per_block, 1, 1);
-
-        anti_alias_activation_forward<input_t, output_t, acc_t>
-	  <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, ftr, alpha, beta, batch_size, channels, seq_len);
-    }
-}
-}
-
-namespace anti_alias_activation {
-
-  torch::Tensor fwd_cuda(torch::Tensor const& input, torch::Tensor const& filter, torch::Tensor const& alpha, torch::Tensor const& beta)
-{
-  // input is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len]
-  const int batches = input.size(0);
-  const int channels = input.size(1);
-  const int seq_len = input.size(2);
-
-  // Output
-  auto act_options = input.options().requires_grad(false);
-  int output_seq_len = seq_len*2; // we'll be dilating between each element by interspersing with zeros
-
-  torch::Tensor anti_alias_activation_results =
-      torch::empty({batches, channels, output_seq_len}, act_options);
-
-  // Softmax Intermediate Result Ptr
-  void* input_ptr = static_cast<void*>(input.data_ptr());
-  void* filter_ptr = static_cast<void*>(filter.data_ptr());
-  void* alpha_ptr = static_cast<void*>(alpha.data_ptr());
-  void* beta_ptr = static_cast<void*>(beta.data_ptr());
-  void* anti_alias_activation_results_ptr = static_cast<void*>(anti_alias_activation_results.data_ptr());
-
-  DISPATCH_FLOAT_HALF_AND_BFLOAT(
-      input.scalar_type(),
-      "dispatch anti alias activation_forward",
-      dispatch_anti_alias_activation_forward<scalar_t, scalar_t, float>(
-        reinterpret_cast<scalar_t*>(anti_alias_activation_results_ptr),
-	    reinterpret_cast<const scalar_t*>(input_ptr),
-        reinterpret_cast<const scalar_t*>(filter_ptr),
-        reinterpret_cast<const scalar_t*>(alpha_ptr),
-	reinterpret_cast<const scalar_t*>(beta_ptr),
-	    batches,
-        channels,
-        seq_len);
-      );
-  return anti_alias_activation_results;
-}
-}

alias_free_cuda/test_activation.py

@@ -1,55 +0,0 @@
-# Copyright (c) 2024 NVIDIA CORPORATION.
-#   Licensed under the MIT license.
-
-import math
-import torch
-import alias_free_cuda
-from alias_free_cuda import activation1d
-from activations import Snake, SnakeBeta
-
-def test_load_fused_kernels():
-    try:
-        import alias_free_cuda
-        import torch
-        print("[Success] load_fused_kernels")
-    except ImportError as e:
-        print("[Fail] load_fused_kernels")
-        raise e
-
-def test_anti_alias_activation():
-    data = torch.rand((10, 10, 50000), device='cuda')
-
-    # check activations.Snake cuda vs. torch
-    fused_anti_alias_activation = activation1d.Activation1d(activation=Snake(10), fused=True).cuda()
-    fused_activation_output = fused_anti_alias_activation(data)
-
-    torch_anti_alias_activation = activation1d.Activation1d(activation=Snake(10), fused=False).cuda()
-    torch_activation_output = torch_anti_alias_activation(data)
-
-    test_result = (fused_activation_output - torch_activation_output).abs()
-    
-    while test_result.dim() != 1:
-        test_result = test_result.mean(dim=-1)
-
-    diff = test_result.mean(dim=-1)
-
-    if diff <= 1e-3:
-        print(
-            f"\n[Success] test_fused_anti_alias_activation"
-            f"\n > mean_difference={diff}"
-            f"\n > fused_values={fused_activation_output[-1][-1][-100:].tolist()}"
-            f"\n > torch_values={torch_activation_output[-1][-1][-100:].tolist()}"
-        )
-    else:
-        print(
-            f"\n[Fail] test_fused_anti_alias_activation"
-            f"\n > mean_difference={diff}, "
-            f"\n > fused_values={fused_activation_output[-1][-1][-30:].tolist()}, "
-            f"\n > torch_values={torch_activation_output[-1][-1][-30:].tolist()}"
-        )
-        
-if __name__ == "__main__":
-    from alias_free_cuda import load
-    load.load()
-    test_load_fused_kernels()
-    test_anti_alias_activation()

alias_free_cuda/test_activation_snake_beta.py

@@ -1,55 +0,0 @@
-# Copyright (c) 2024 NVIDIA CORPORATION.
-#   Licensed under the MIT license.
-
-import math
-import torch
-import alias_free_cuda
-from alias_free_cuda import activation1d
-from activations import Snake, SnakeBeta
-
-def test_load_fused_kernels():
-    try:
-        import alias_free_cuda
-        import torch
-        print("[Success] load_fused_kernels")
-    except ImportError as e:
-        print("[Fail] load_fused_kernels")
-        raise e
-
-def test_anti_alias_activation():
-    data = torch.rand((10, 10, 50000), device='cuda')
-
-    # check activations.Snake cuda vs. torch
-    fused_anti_alias_activation = activation1d.Activation1d(activation=SnakeBeta(10), fused=True).cuda()
-    fused_activation_output = fused_anti_alias_activation(data)
-
-    torch_anti_alias_activation = activation1d.Activation1d(activation=SnakeBeta(10), fused=False).cuda()
-    torch_activation_output = torch_anti_alias_activation(data)
-    
-    test_result = (fused_activation_output - torch_activation_output).abs()
-
-    while test_result.dim() != 1:
-        test_result = test_result.mean(dim=-1)
-
-    diff = test_result.mean(dim=-1)
-
-    if diff <= 1e-3:
-        print(
-            f"\n[Success] test_fused_anti_alias_activation"
-            f"\n > mean_difference={diff}"
-            f"\n > fused_values={fused_activation_output[-1][-1][-100:].tolist()}"
-            f"\n > torch_values={torch_activation_output[-1][-1][-100:].tolist()}"
-        )
-    else:
-        print(
-            f"\n[Fail] test_fused_anti_alias_activation"
-            f"\n > mean_difference={diff}, "
-            f"\n > fused_values={fused_activation_output[-1][-1][-30:].tolist()}, "
-            f"\n > torch_values={torch_activation_output[-1][-1][-30:].tolist()}"
-        )
-
-if __name__ == "__main__":
-    from alias_free_cuda import load
-    load.load()
-    test_load_fused_kernels()
-    test_anti_alias_activation()

alias_free_cuda/type_shim.h

@@ -1,97 +0,0 @@
-/* coding=utf-8
- * Copyright (c) 2020, NVIDIA CORPORATION.  All rights reserved.
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-
-#include <ATen/ATen.h>
-#include "compat.h"
-
-
-#define DISPATCH_FLOAT_HALF_AND_BFLOAT(TYPE, NAME, ...)			\
-  switch(TYPE)							\
-    {									\
-    case at::ScalarType::Float:			\
-      {									\
-	using scalar_t = float;			    \
-	__VA_ARGS__;						\
-	break;								\
-      }                                 \
-    case at::ScalarType::Half:			\
-      {									\
-	using scalar_t = at::Half;			\
-	__VA_ARGS__;						\
-	break;								\
-      }									\
-    case at::ScalarType::BFloat16:		\
-      {									\
-	using scalar_t = at::BFloat16;		\
-	__VA_ARGS__;						\
-	break;								\
-      }									\
-    default:							\
-      AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'");	\
-      }
-
-
-
-#define DISPATCH_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(TYPEIN, TYPEOUT, NAME, ...) \
-  switch(TYPEIN)							\
-    {									\
-    case at::ScalarType::Float:						\
-      {									\
-	using scalar_t_in = float;					\
-	switch(TYPEOUT)							\
-	  {								\
-	  case at::ScalarType::Float:					\
-	    {								\
-	      using scalar_t_out = float;				\
-	      __VA_ARGS__;						\
-	      break;							\
-	    }								\
-	  case at::ScalarType::Half:					\
-	    {								\
-	      using scalar_t_out = at::Half;				\
-	      __VA_ARGS__;						\
-	      break;							\
-	    }								\
-	  case at::ScalarType::BFloat16:				\
-	    {								\
-	      using scalar_t_out = at::BFloat16;			\
-	      __VA_ARGS__;						\
-	      break;							\
-	    }								\
-	  default:							\
-	    AT_ERROR(#NAME, " not implemented for '", toString(TYPEOUT), "'"); \
-	  }								\
-	break;								\
-      }									\
-    case at::ScalarType::Half:						\
-      {									\
-	using scalar_t_in = at::Half;					\
-	using scalar_t_out = at::Half;					\
-	__VA_ARGS__;							\
-	break;								\
-      }									\
-    case at::ScalarType::BFloat16:					\
-      {									\
-	using scalar_t_in = at::BFloat16;				\
-	using scalar_t_out = at::BFloat16;				\
-	__VA_ARGS__;							\
-	break;								\
-      }									\
-    default:								\
-      AT_ERROR(#NAME, " not implemented for '", toString(TYPEIN), "'");	\
-    }
-

app.py

@@ -273,7 +273,7 @@ with iface:
         <h3>News</h3>
         <p>[Jul 2024] We release BigVGAN-v2 along with pretrained checkpoints. Below are the highlights:</p>
         <ul>
-            <li>Custom CUDA kernel for inference: we provide a fused upsampling + activation kernel written in CUDA for accelerated inference speed. Our test shows 1.5 - 3x faster speed on a single A100 GPU.</li>
+            <li>Custom CUDA kernel for inference: we provide a fused anti-aliased activation kernel written in CUDA for accelerated inference speed. Our test shows 1.5 - 3x faster speed on a single A100 GPU.</li>
             <li>Improved discriminator and loss: BigVGAN-v2 is trained using a <a href="https://arxiv.org/abs/2311.14957" target="_blank">multi-scale sub-band CQT discriminator</a> and a <a href="https://arxiv.org/abs/2306.06546" target="_blank">multi-scale mel spectrogram loss</a>.</li>
             <li>Larger training data: BigVGAN-v2 is trained using datasets containing diverse audio types, including speech in multiple languages, environmental sounds, and instruments.</li>
             <li>We provide pretrained checkpoints of BigVGAN-v2 using diverse audio configurations, supporting up to 44 kHz sampling rate and 512x upsampling ratio. See the table below for the link.</li>

bigvgan.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2024 NVIDIA CORPORATION. 
+# Copyright (c) 2024 NVIDIA CORPORATION.
 #   Licensed under the MIT license.
 
 # Adapted from https://github.com/jik876/hifi-gan under the MIT license.
@@ -7,78 +7,127 @@
 import os
 import json
 from pathlib import Path
-
-from collections import namedtuple
-from typing import Optional, List, Union, Dict
+from typing import Optional, Union, Dict
 
 import torch
-import torch.nn.functional as F
 import torch.nn as nn
 from torch.nn import Conv1d, ConvTranspose1d
 from torch.nn.utils import weight_norm, remove_weight_norm
 
 import activations
 from utils import init_weights, get_padding
-from alias_free_torch.act import Activation1d as TorchActivation1d
+from alias_free_activation.torch.act import Activation1d as TorchActivation1d
 from env import AttrDict
 
 from huggingface_hub import PyTorchModelHubMixin, hf_hub_download
 
+
 def load_hparams_from_json(path) -> AttrDict:
     with open(path) as f:
         data = f.read()
-    h = json.loads(data)
-    return AttrDict(h)
+    return AttrDict(json.loads(data))
+
 
 class AMPBlock1(torch.nn.Module):
-    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3, 5), activation=None):
-        super(AMPBlock1, self).__init__()
+    """
+    AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer.
+    AMPBlock1 has additional self.convs2 that contains additional Conv1d layers with a fixed dilation=1 followed by each layer in self.convs1
+
+    Args:
+        h (AttrDict): Hyperparameters.
+        channels (int): Number of convolution channels.
+        kernel_size (int): Size of the convolution kernel. Default is 3.
+        dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5).
+        activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Default is None.
+    """
+
+    def __init__(
+        self,
+        h: AttrDict,
+        channels: int,
+        kernel_size: int = 3,
+        dilation: tuple = (1, 3, 5),
+        activation: str = None,
+    ):
+        super().__init__()
+        
         self.h = h
 
-        self.convs1 = nn.ModuleList([
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
-                               padding=get_padding(kernel_size, dilation[0]))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
-                               padding=get_padding(kernel_size, dilation[1]))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
-                               padding=get_padding(kernel_size, dilation[2])))
-        ])
+        self.convs1 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilation
+            ]
+        )
         self.convs1.apply(init_weights)
 
-        self.convs2 = nn.ModuleList([
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
-                               padding=get_padding(kernel_size, 1))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
-                               padding=get_padding(kernel_size, 1))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
-                               padding=get_padding(kernel_size, 1)))
-        ])
+        self.convs2 = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=1,
+                        padding=get_padding(kernel_size, 1),
+                    )
+                )
+                for _ in range(len(dilation))
+            ]
+        )
         self.convs2.apply(init_weights)
 
-        self.num_layers = len(self.convs1) + len(self.convs2) # total number of conv layers
+        self.num_layers = len(self.convs1) + len(
+            self.convs2
+        )  # Total number of conv layers
 
-        # select which Activation1d, lazy-load cuda version to ensure backward compatibility
+        # Select which Activation1d, lazy-load cuda version to ensure backward compatibility
         if self.h.get("use_cuda_kernel", False):
-            # faster CUDA kernel implementation of Activation1d
-            from alias_free_cuda.activation1d import Activation1d as CudaActivation1d
+            from alias_free_activation.cuda.activation1d import (
+                Activation1d as CudaActivation1d,
+            )
+
             Activation1d = CudaActivation1d
         else:
             Activation1d = TorchActivation1d
 
-        if activation == 'snake': # periodic nonlinearity with snake function and anti-aliasing
-            self.activations = nn.ModuleList([
-                Activation1d(
-                    activation=activations.Snake(channels, alpha_logscale=h.snake_logscale))
-                for _ in range(self.num_layers)
-            ])
-        elif activation == 'snakebeta': # periodic nonlinearity with snakebeta function and anti-aliasing
-            self.activations = nn.ModuleList([
-                Activation1d(
-                    activation=activations.SnakeBeta(channels, alpha_logscale=h.snake_logscale))
-                 for _ in range(self.num_layers)
-            ])
+        # Activation functions
+        if activation == "snake":
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=activations.Snake(
+                            channels, alpha_logscale=h.snake_logscale
+                        )
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        elif activation == "snakebeta":
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=activations.SnakeBeta(
+                            channels, alpha_logscale=h.snake_logscale
+                        )
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
         else:
-            raise NotImplementedError("activation incorrectly specified. check the config file and look for 'activation'.")
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
 
     def forward(self, x):
         acts1, acts2 = self.activations[::2], self.activations[1::2]
@@ -99,51 +148,93 @@ class AMPBlock1(torch.nn.Module):
 
 
 class AMPBlock2(torch.nn.Module):
-    def __init__(self, h, channels, kernel_size=3, dilation=(1, 3), activation=None):
-        super(AMPBlock2, self).__init__()
+    """
+    AMPBlock applies Snake / SnakeBeta activation functions with trainable parameters that control periodicity, defined for each layer.
+    Unlike AMPBlock1, AMPBlock2 does not contain extra Conv1d layers with fixed dilation=1
+
+    Args:
+        h (AttrDict): Hyperparameters.
+        channels (int): Number of convolution channels.
+        kernel_size (int): Size of the convolution kernel. Default is 3.
+        dilation (tuple): Dilation rates for the convolutions. Each dilation layer has two convolutions. Default is (1, 3, 5).
+        activation (str): Activation function type. Should be either 'snake' or 'snakebeta'. Default is None.
+    """
+
+    def __init__(
+        self,
+        h: AttrDict,
+        channels: int,
+        kernel_size: int = 3,
+        dilation: tuple = (1, 3, 5),
+        activation: str = None,
+    ):
+        super().__init__()
+        
         self.h = h
 
-        self.convs = nn.ModuleList([
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
-                               padding=get_padding(kernel_size, dilation[0]))),
-            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
-                               padding=get_padding(kernel_size, dilation[1])))
-        ])
+        self.convs = nn.ModuleList(
+            [
+                weight_norm(
+                    Conv1d(
+                        channels,
+                        channels,
+                        kernel_size,
+                        stride=1,
+                        dilation=d,
+                        padding=get_padding(kernel_size, d),
+                    )
+                )
+                for d in dilation
+            ]
+        )
         self.convs.apply(init_weights)
 
-        self.num_layers = len(self.convs) # total number of conv layers
+        self.num_layers = len(self.convs)  # Total number of conv layers
 
-        # select which Activation1d, lazy-load cuda version to ensure backward compatibility
+        # Select which Activation1d, lazy-load cuda version to ensure backward compatibility
         if self.h.get("use_cuda_kernel", False):
-            # faster CUDA kernel implementation of Activation1d
-            from alias_free_cuda.activation1d import Activation1d as CudaActivation1d
+            from alias_free_activation.cuda.activation1d import (
+                Activation1d as CudaActivation1d,
+            )
+
             Activation1d = CudaActivation1d
         else:
             Activation1d = TorchActivation1d
 
-        if activation == 'snake': # periodic nonlinearity with snake function and anti-aliasing
-            self.activations = nn.ModuleList([
-                Activation1d(
-                    activation=activations.Snake(channels, alpha_logscale=h.snake_logscale))
-                for _ in range(self.num_layers)
-            ])
-        elif activation == 'snakebeta': # periodic nonlinearity with snakebeta function and anti-aliasing
-            self.activations = nn.ModuleList([
-                Activation1d(
-                    activation=activations.SnakeBeta(channels, alpha_logscale=h.snake_logscale))
-                 for _ in range(self.num_layers)
-            ])
+        # Activation functions
+        if activation == "snake":
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=activations.Snake(
+                            channels, alpha_logscale=h.snake_logscale
+                        )
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
+        elif activation == "snakebeta":
+            self.activations = nn.ModuleList(
+                [
+                    Activation1d(
+                        activation=activations.SnakeBeta(
+                            channels, alpha_logscale=h.snake_logscale
+                        )
+                    )
+                    for _ in range(self.num_layers)
+                ]
+            )
         else:
-            raise NotImplementedError("activation incorrectly specified. check the config file and look for 'activation'.")
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
 
     def forward(self, x):
-        for c, a in zip (self.convs, self.activations):
+        for c, a in zip(self.convs, self.activations):
             xt = a(x)
             xt = c(xt)
             x = xt + x
 
-        return x
-
     def remove_weight_norm(self):
         for l in self.convs:
             remove_weight_norm(l)
@@ -157,83 +248,121 @@ class BigVGAN(
     docs_url="https://github.com/NVIDIA/BigVGAN/blob/main/README.md",
     pipeline_tag="audio-to-audio",
     license="mit",
-    tags=["neural-vocoder", "audio-generation", "arxiv:2206.04658"]
+    tags=["neural-vocoder", "audio-generation", "arxiv:2206.04658"],
 ):
-    # this is our main BigVGAN model. Applies anti-aliased periodic activation for resblocks.
-    # New in v2: if use_cuda_kernel is set to True, it loads optimized CUDA kernels for AMP.
-    # NOTE: use_cuda_kernel=True should be used for inference only (training is not supported).
-    def __init__(
-        self,
-        h,
-        use_cuda_kernel: bool=False
-    ):
-        super(BigVGAN, self).__init__()
+    """
+    BigVGAN is a neural vocoder model that applies anti-aliased periodic activation for residual blocks (resblocks).
+    New in BigVGAN-v2: it can optionally use optimized CUDA kernels for AMP (anti-aliased multi-periodicity) blocks.
+
+    Args:
+        h (AttrDict): Hyperparameters.
+        use_cuda_kernel (bool): If set to True, loads optimized CUDA kernels for AMP. This should be used for inference only, as training is not supported with CUDA kernels.
+
+    Note:
+        - The `use_cuda_kernel` parameter should be used for inference only, as training with CUDA kernels is not supported.
+        - Ensure that the activation function is correctly specified in the hyperparameters (h.activation).
+    """
+
+    def __init__(self, h: AttrDict, use_cuda_kernel: bool = False):
+        super().__init__()
         self.h = h
-        self.h["use_cuda_kernel"] = use_cuda_kernel # add it to global hyperparameters (h)
+        self.h["use_cuda_kernel"] = use_cuda_kernel
+
+        # Select which Activation1d, lazy-load cuda version to ensure backward compatibility
+        if self.h.get("use_cuda_kernel", False):
+            from alias_free_activation.cuda.activation1d import (
+                Activation1d as CudaActivation1d,
+            )
+
+            Activation1d = CudaActivation1d
+        else:
+            Activation1d = TorchActivation1d
 
         self.num_kernels = len(h.resblock_kernel_sizes)
         self.num_upsamples = len(h.upsample_rates)
 
-        # pre conv
-        self.conv_pre = weight_norm(Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3))
+        # Pre-conv
+        self.conv_pre = weight_norm(
+            Conv1d(h.num_mels, h.upsample_initial_channel, 7, 1, padding=3)
+        )
 
-        # define which AMPBlock to use. BigVGAN uses AMPBlock1 as default
-        resblock = AMPBlock1 if h.resblock == '1' else AMPBlock2
+        # Define which AMPBlock to use. BigVGAN uses AMPBlock1 as default
+        if h.resblock == "1":
+            resblock_class = AMPBlock1
+        elif h.resblock == "2":
+            resblock_class = AMPBlock2
+        else:
+            raise ValueError(
+                f"Incorrect resblock class specified in hyperparameters. Got {h.resblock}"
+            )
 
-        # transposed conv-based upsamplers. does not apply anti-aliasing
+        # Transposed conv-based upsamplers. does not apply anti-aliasing
         self.ups = nn.ModuleList()
         for i, (u, k) in enumerate(zip(h.upsample_rates, h.upsample_kernel_sizes)):
-            self.ups.append(nn.ModuleList([
-                weight_norm(ConvTranspose1d(h.upsample_initial_channel // (2 ** i),
-                                            h.upsample_initial_channel // (2 ** (i + 1)),
-                                            k, u, padding=(k - u) // 2))
-            ]))
+            self.ups.append(
+                nn.ModuleList(
+                    [
+                        weight_norm(
+                            ConvTranspose1d(
+                                h.upsample_initial_channel // (2**i),
+                                h.upsample_initial_channel // (2 ** (i + 1)),
+                                k,
+                                u,
+                                padding=(k - u) // 2,
+                            )
+                        )
+                    ]
+                )
+            )
 
-        # residual blocks using anti-aliased multi-periodicity composition modules (AMP)
+        # Residual blocks using anti-aliased multi-periodicity composition modules (AMP)
         self.resblocks = nn.ModuleList()
         for i in range(len(self.ups)):
             ch = h.upsample_initial_channel // (2 ** (i + 1))
-            for j, (k, d) in enumerate(zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)):
-                self.resblocks.append(resblock(h, ch, k, d, activation=h.activation))
-        
-        # select which Activation1d, lazy-load cuda version to ensure backward compatibility
-        if self.h.get("use_cuda_kernel", False):
-            # faster CUDA kernel implementation of Activation1d
-            from alias_free_cuda.activation1d import Activation1d as CudaActivation1d
-            Activation1d = CudaActivation1d
-        else:
-            Activation1d = TorchActivation1d
+            for j, (k, d) in enumerate(
+                zip(h.resblock_kernel_sizes, h.resblock_dilation_sizes)
+            ):
+                self.resblocks.append(
+                    resblock_class(h, ch, k, d, activation=h.activation)
+                )
 
-        # post conv
-        if h.activation == "snake": # periodic nonlinearity with snake function and anti-aliasing
-            activation_post = activations.Snake(ch, alpha_logscale=h.snake_logscale)
-            self.activation_post = Activation1d(activation=activation_post)
-        elif h.activation == "snakebeta": # periodic nonlinearity with snakebeta function and anti-aliasing
-            activation_post = activations.SnakeBeta(ch, alpha_logscale=h.snake_logscale)
-            self.activation_post = Activation1d(activation=activation_post)
-        else:
-            raise NotImplementedError("activation incorrectly specified. check the config file and look for 'activation'.")
-        
-        # whether to use bias for the final conv_post. Defaults to True for backward compatibility
+        # Post-conv
+        activation_post = (
+            activations.Snake(ch, alpha_logscale=h.snake_logscale)
+            if h.activation == "snake"
+            else (
+                activations.SnakeBeta(ch, alpha_logscale=h.snake_logscale)
+                if h.activation == "snakebeta"
+                else None
+            )
+        )
+        if activation_post is None:
+            raise NotImplementedError(
+                "activation incorrectly specified. check the config file and look for 'activation'."
+            )
+
+        self.activation_post = Activation1d(activation=activation_post)
+
+        # Whether to use bias for the final conv_post. Default to True for backward compatibility
         self.use_bias_at_final = h.get("use_bias_at_final", True)
-        self.conv_post = weight_norm(Conv1d(
-            ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final
-        ))
+        self.conv_post = weight_norm(
+            Conv1d(ch, 1, 7, 1, padding=3, bias=self.use_bias_at_final)
+        )
 
-        # weight initialization
+        # Weight initialization
         for i in range(len(self.ups)):
             self.ups[i].apply(init_weights)
         self.conv_post.apply(init_weights)
-        
-        # final tanh activation. Defaults to True for backward compatibility
+
+        # Final tanh activation. Defaults to True for backward compatibility
         self.use_tanh_at_final = h.get("use_tanh_at_final", True)
 
     def forward(self, x):
-        # pre conv
+        # Pre-conv
         x = self.conv_pre(x)
 
         for i in range(self.num_upsamples):
-            # upsampling
+            # Upsampling
             for i_up in range(len(self.ups[i])):
                 x = self.ups[i][i_up](x)
             # AMP blocks
@@ -245,20 +374,20 @@ class BigVGAN(
                     xs += self.resblocks[i * self.num_kernels + j](x)
             x = xs / self.num_kernels
 
-        # post conv
+        # Post-conv
         x = self.activation_post(x)
         x = self.conv_post(x)
-        # final tanh activation
+        # Final tanh activation
         if self.use_tanh_at_final:
             x = torch.tanh(x)
         else:
-            x = torch.clamp(x, min=-1., max=1.) # bound the output to [-1, 1]
+            x = torch.clamp(x, min=-1.0, max=1.0)  # Bound the output to [-1, 1]
 
         return x
 
     def remove_weight_norm(self):
         try:
-            print('Removing weight norm...')
+            print("Removing weight norm...")
             for l in self.ups:
                 for l_i in l:
                     remove_weight_norm(l_i)
@@ -267,23 +396,18 @@ class BigVGAN(
             remove_weight_norm(self.conv_pre)
             remove_weight_norm(self.conv_post)
         except ValueError:
-            print('[INFO] Model already removed weight norm. Skipping!')
+            print("[INFO] Model already removed weight norm. Skipping!")
             pass
 
-    ##################################################################
-    # additional methods for huggingface_hub support
-    ##################################################################
+    # Additional methods for huggingface_hub support
     def _save_pretrained(self, save_directory: Path) -> None:
         """Save weights and config.json from a Pytorch model to a local directory."""
 
-        model_path = save_directory / 'bigvgan_generator.pt'
-        torch.save(
-            {'generator': self.state_dict()},
-            model_path
-        )
+        model_path = save_directory / "bigvgan_generator.pt"
+        torch.save({"generator": self.state_dict()}, model_path)
 
-        config_path = save_directory / 'config.json'
-        with open(config_path, 'w') as config_file:
+        config_path = save_directory / "config.json"
+        with open(config_path, "w") as config_file:
             json.dump(self.h, config_file, indent=4)
 
     @classmethod
@@ -298,23 +422,21 @@ class BigVGAN(
         resume_download: bool,
         local_files_only: bool,
         token: Union[str, bool, None],
-        map_location: str = "cpu", # additional argument
-        strict: bool = False, # additional argument
+        map_location: str = "cpu",  # Additional argument
+        strict: bool = False,  # Additional argument
         use_cuda_kernel: bool = False,
         **model_kwargs,
     ):
         """Load Pytorch pretrained weights and return the loaded model."""
 
-        ##################################################################
-        # download and load hyperparameters (h) used by BigVGAN
-        ##################################################################
+        # Download and load hyperparameters (h) used by BigVGAN
         if os.path.isdir(model_id):
             print("Loading config.json from local directory")
-            config_file = os.path.join(model_id, 'config.json')
+            config_file = os.path.join(model_id, "config.json")
         else:
             config_file = hf_hub_download(
                 repo_id=model_id,
-                filename='config.json',
+                filename="config.json",
                 revision=revision,
                 cache_dir=cache_dir,
                 force_download=force_download,
@@ -325,26 +447,28 @@ class BigVGAN(
             )
         h = load_hparams_from_json(config_file)
 
-        ##################################################################
         # instantiate BigVGAN using h
-        ##################################################################
         if use_cuda_kernel:
-            print(f"[WARNING] You have specified use_cuda_kernel=True during BigVGAN.from_pretrained(). Only inference is supported (training is not implemented)!")
-            print(f"[WARNING] You need nvcc and ninja installed in your system that matches your PyTorch build is using to build the kernel. If not, the model will fail to initialize or generate incorrect waveform!")
-            print(f"[WARNING] For detail, see the official GitHub repository: https://github.com/NVIDIA/BigVGAN?tab=readme-ov-file#using-custom-cuda-kernel-for-synthesis")
+            print(
+                f"[WARNING] You have specified use_cuda_kernel=True during BigVGAN.from_pretrained(). Only inference is supported (training is not implemented)!"
+            )
+            print(
+                f"[WARNING] You need nvcc and ninja installed in your system that matches your PyTorch build is using to build the kernel. If not, the model will fail to initialize or generate incorrect waveform!"
+            )
+            print(
+                f"[WARNING] For detail, see the official GitHub repository: https://github.com/NVIDIA/BigVGAN?tab=readme-ov-file#using-custom-cuda-kernel-for-synthesis"
+            )
         model = cls(h, use_cuda_kernel=use_cuda_kernel)
 
-        ##################################################################
-        # download and load pretrained generator weight
-        ##################################################################
+        # Download and load pretrained generator weight
         if os.path.isdir(model_id):
             print("Loading weights from local directory")
-            model_file = os.path.join(model_id, 'bigvgan_generator.pt')
+            model_file = os.path.join(model_id, "bigvgan_generator.pt")
         else:
             print(f"Loading weights from {model_id}")
             model_file = hf_hub_download(
                 repo_id=model_id,
-                filename='bigvgan_generator.pt',
+                filename="bigvgan_generator.pt",
                 revision=revision,
                 cache_dir=cache_dir,
                 force_download=force_download,
@@ -352,15 +476,17 @@ class BigVGAN(
                 resume_download=resume_download,
                 token=token,
                 local_files_only=local_files_only,
-                )
-        
+            )
+
         checkpoint_dict = torch.load(model_file, map_location=map_location)
 
         try:
-            model.load_state_dict(checkpoint_dict['generator'])
+            model.load_state_dict(checkpoint_dict["generator"])
         except RuntimeError:
-            print(f"[INFO] the pretrained checkpoint does not contain weight norm. Loading the checkpoint after removing weight norm!")
+            print(
+                f"[INFO] the pretrained checkpoint does not contain weight norm. Loading the checkpoint after removing weight norm!"
+            )
             model.remove_weight_norm()
-            model.load_state_dict(checkpoint_dict['generator'])
+            model.load_state_dict(checkpoint_dict["generator"])
 
-        return model
+        return model

meldataset.py

@@ -1,66 +1,354 @@
-# Copyright (c) 2024 NVIDIA CORPORATION. 
+# Copyright (c) 2024 NVIDIA CORPORATION.
 #   Licensed under the MIT license.
 
 # Adapted from https://github.com/jik876/hifi-gan under the MIT license.
 #   LICENSE is in incl_licenses directory.
 
+import math
+import os
+import random
 import torch
 import torch.utils.data
 import numpy as np
+from librosa.util import normalize
 from scipy.io.wavfile import read
 from librosa.filters import mel as librosa_mel_fn
+import pathlib
+from tqdm import tqdm
+
+MAX_WAV_VALUE = 32767.0  # NOTE: 32768.0 -1 to prevent int16 overflow (results in popping sound in corner cases)
+
+
+def load_wav(full_path, sr_target):
+    sampling_rate, data = read(full_path)
+    if sampling_rate != sr_target:
+        raise RuntimeError(
+            f"Sampling rate of the file {full_path} is {sampling_rate} Hz, but the model requires {sr_target} Hz"
+        )
+    return data, sampling_rate
 
-MAX_WAV_VALUE = 32767.0 # NOTE: 32768.0 -1 to prevent int16 overflow (results in popping sound in corner cases)
 
 def dynamic_range_compression(x, C=1, clip_val=1e-5):
     return np.log(np.clip(x, a_min=clip_val, a_max=None) * C)
 
+
 def dynamic_range_decompression(x, C=1):
     return np.exp(x) / C
 
+
 def dynamic_range_compression_torch(x, C=1, clip_val=1e-5):
     return torch.log(torch.clamp(x, min=clip_val) * C)
 
+
 def dynamic_range_decompression_torch(x, C=1):
     return torch.exp(x) / C
 
+
 def spectral_normalize_torch(magnitudes):
-    output = dynamic_range_compression_torch(magnitudes)
-    return output
+    return dynamic_range_compression_torch(magnitudes)
+
 
 def spectral_de_normalize_torch(magnitudes):
-    output = dynamic_range_decompression_torch(magnitudes)
-    return output
+    return dynamic_range_decompression_torch(magnitudes)
+
+
+mel_basis_cache = {}
+hann_window_cache = {}
 
-mel_basis = {}
-hann_window = {}
 
-def mel_spectrogram(y, n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax, center=False):
-    if torch.min(y) < -1.:
-        print('min value is ', torch.min(y))
-    if torch.max(y) > 1.:
-        print('max value is ', torch.max(y))
+def mel_spectrogram(
+    y: torch.Tensor,
+    n_fft: int,
+    num_mels: int,
+    sampling_rate: int,
+    hop_size: int,
+    win_size: int,
+    fmin: int,
+    fmax: int = None,
+    center: bool = False,
+) -> torch.Tensor:
+    """
+    Calculate the mel spectrogram of an input signal.
+    This function uses slaney norm for the librosa mel filterbank (using librosa.filters.mel) and uses Hann window for STFT (using torch.stft).
 
-    global mel_basis, hann_window
-    if fmax not in mel_basis:
-        mel = librosa_mel_fn(sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax)
-        str_key_mel_basis = str(fmax)+'_'+str(y.device)
-        mel_basis[str_key_mel_basis] = torch.from_numpy(mel).float().to(y.device)
-        hann_window[str(y.device)] = torch.hann_window(win_size).to(y.device)
+    Args:
+        y (torch.Tensor): Input signal.
+        n_fft (int): FFT size.
+        num_mels (int): Number of mel bins.
+        sampling_rate (int): Sampling rate of the input signal.
+        hop_size (int): Hop size for STFT.
+        win_size (int): Window size for STFT.
+        fmin (int): Minimum frequency for mel filterbank.
+        fmax (int): Maximum frequency for mel filterbank. If None, defaults to half the sampling rate (fmax = sr / 2.0) inside librosa_mel_fn
+        center (bool): Whether to pad the input to center the frames. Default is False.
 
-    y = torch.nn.functional.pad(y.unsqueeze(1), (int((n_fft-hop_size)/2), int((n_fft-hop_size)/2)), mode='reflect')
-    y = y.squeeze(1)
+    Returns:
+        torch.Tensor: Mel spectrogram.
+    """
+    if torch.min(y) < -1.0:
+        print(f"[WARNING] Min value of input waveform signal is {torch.min(y)}")
+    if torch.max(y) > 1.0:
+        print(f"[WARNING] Max value of input waveform signal is {torch.max(y)}")
 
-    # complex tensor as default, then use view_as_real for future pytorch compatibility
-    spec = torch.stft(y, n_fft, hop_length=hop_size, win_length=win_size, window=hann_window[str(y.device)],
-                      center=center, pad_mode='reflect', normalized=False, onesided=True, return_complex=True)
-    spec = torch.view_as_real(spec)
-    spec = torch.sqrt(spec.pow(2).sum(-1)+(1e-9))
+    device = y.device
+    key = f"{n_fft}_{num_mels}_{sampling_rate}_{hop_size}_{win_size}_{fmin}_{fmax}_{device}"
 
-    spec = torch.matmul(mel_basis[str_key_mel_basis], spec)
-    spec = spectral_normalize_torch(spec)
+    if key not in mel_basis_cache:
+        mel = librosa_mel_fn(
+            sr=sampling_rate, n_fft=n_fft, n_mels=num_mels, fmin=fmin, fmax=fmax
+        )
+        mel_basis_cache[key] = torch.from_numpy(mel).float().to(device)
+        hann_window_cache[key] = torch.hann_window(win_size).to(device)
+
+    mel_basis = mel_basis_cache[key]
+    hann_window = hann_window_cache[key]
+
+    padding = (n_fft - hop_size) // 2
+    y = torch.nn.functional.pad(
+        y.unsqueeze(1), (padding, padding), mode="reflect"
+    ).squeeze(1)
+
+    spec = torch.stft(
+        y,
+        n_fft,
+        hop_length=hop_size,
+        win_length=win_size,
+        window=hann_window,
+        center=center,
+        pad_mode="reflect",
+        normalized=False,
+        onesided=True,
+        return_complex=True,
+    )
+    spec = torch.sqrt(torch.view_as_real(spec).pow(2).sum(-1) + 1e-9)
+
+    mel_spec = torch.matmul(mel_basis, spec)
+    mel_spec = spectral_normalize_torch(mel_spec)
+
+    return mel_spec
 
-    return spec
 
 def get_mel_spectrogram(wav, h):
-    return mel_spectrogram(wav, h.n_fft, h.num_mels, h.sampling_rate, h.hop_size, h.win_size, h.fmin, h.fmax)
+    """
+    Generate mel spectrogram from a waveform using given hyperparameters.
+
+    Args:
+        wav (torch.Tensor): Input waveform.
+        h: Hyperparameters object with attributes n_fft, num_mels, sampling_rate, hop_size, win_size, fmin, fmax.
+
+    Returns:
+        torch.Tensor: Mel spectrogram.
+    """
+    return mel_spectrogram(
+        wav,
+        h.n_fft,
+        h.num_mels,
+        h.sampling_rate,
+        h.hop_size,
+        h.win_size,
+        h.fmin,
+        h.fmax,
+    )
+
+
+def get_dataset_filelist(a):
+    training_files = []
+    validation_files = []
+    list_unseen_validation_files = []
+
+    with open(a.input_training_file, "r", encoding="utf-8") as fi:
+        training_files = [
+            os.path.join(a.input_wavs_dir, x.split("|")[0] + ".wav")
+            for x in fi.read().split("\n")
+            if len(x) > 0
+        ]
+        print(f"first training file: {training_files[0]}")
+
+    with open(a.input_validation_file, "r", encoding="utf-8") as fi:
+        validation_files = [
+            os.path.join(a.input_wavs_dir, x.split("|")[0] + ".wav")
+            for x in fi.read().split("\n")
+            if len(x) > 0
+        ]
+        print(f"first validation file: {validation_files[0]}")
+
+    for i in range(len(a.list_input_unseen_validation_file)):
+        with open(a.list_input_unseen_validation_file[i], "r", encoding="utf-8") as fi:
+            unseen_validation_files = [
+                os.path.join(a.list_input_unseen_wavs_dir[i], x.split("|")[0] + ".wav")
+                for x in fi.read().split("\n")
+                if len(x) > 0
+            ]
+            print(
+                f"first unseen {i}th validation fileset: {unseen_validation_files[0]}"
+            )
+            list_unseen_validation_files.append(unseen_validation_files)
+
+    return training_files, validation_files, list_unseen_validation_files
+
+
+class MelDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        training_files,
+        hparams,
+        segment_size,
+        n_fft,
+        num_mels,
+        hop_size,
+        win_size,
+        sampling_rate,
+        fmin,
+        fmax,
+        split=True,
+        shuffle=True,
+        n_cache_reuse=1,
+        device=None,
+        fmax_loss=None,
+        fine_tuning=False,
+        base_mels_path=None,
+        is_seen=True,
+    ):
+        self.audio_files = training_files
+        random.seed(1234)
+        if shuffle:
+            random.shuffle(self.audio_files)
+        self.hparams = hparams
+        self.is_seen = is_seen
+        if self.is_seen:
+            self.name = pathlib.Path(self.audio_files[0]).parts[0]
+        else:
+            self.name = "-".join(pathlib.Path(self.audio_files[0]).parts[:2]).strip("/")
+
+        self.segment_size = segment_size
+        self.sampling_rate = sampling_rate
+        self.split = split
+        self.n_fft = n_fft
+        self.num_mels = num_mels
+        self.hop_size = hop_size
+        self.win_size = win_size
+        self.fmin = fmin
+        self.fmax = fmax
+        self.fmax_loss = fmax_loss
+        self.cached_wav = None
+        self.n_cache_reuse = n_cache_reuse
+        self._cache_ref_count = 0
+        self.device = device
+        self.fine_tuning = fine_tuning
+        self.base_mels_path = base_mels_path
+
+        print("[INFO] checking dataset integrity...")
+        for i in tqdm(range(len(self.audio_files))):
+            assert os.path.exists(
+                self.audio_files[i]
+            ), f"{self.audio_files[i]} not found"
+
+    def __getitem__(self, index):
+        filename = self.audio_files[index]
+        if self._cache_ref_count == 0:
+            audio, sampling_rate = load_wav(filename, self.sampling_rate)
+            audio = audio / MAX_WAV_VALUE
+            if not self.fine_tuning:
+                audio = normalize(audio) * 0.95
+            self.cached_wav = audio
+            if sampling_rate != self.sampling_rate:
+                raise ValueError(
+                    f"{sampling_rate} SR doesn't match target {self.sampling_rate} SR"
+                )
+            self._cache_ref_count = self.n_cache_reuse
+        else:
+            audio = self.cached_wav
+            self._cache_ref_count -= 1
+
+        audio = torch.FloatTensor(audio)
+        audio = audio.unsqueeze(0)
+
+        if not self.fine_tuning:
+            if self.split:
+                if audio.size(1) >= self.segment_size:
+                    max_audio_start = audio.size(1) - self.segment_size
+                    audio_start = random.randint(0, max_audio_start)
+                    audio = audio[:, audio_start : audio_start + self.segment_size]
+                else:
+                    audio = torch.nn.functional.pad(
+                        audio, (0, self.segment_size - audio.size(1)), "constant"
+                    )
+
+                mel = mel_spectrogram(
+                    audio,
+                    self.n_fft,
+                    self.num_mels,
+                    self.sampling_rate,
+                    self.hop_size,
+                    self.win_size,
+                    self.fmin,
+                    self.fmax,
+                    center=False,
+                )
+            else:  # Validation step
+                # Match audio length to self.hop_size * n for evaluation
+                if (audio.size(1) % self.hop_size) != 0:
+                    audio = audio[:, : -(audio.size(1) % self.hop_size)]
+                mel = mel_spectrogram(
+                    audio,
+                    self.n_fft,
+                    self.num_mels,
+                    self.sampling_rate,
+                    self.hop_size,
+                    self.win_size,
+                    self.fmin,
+                    self.fmax,
+                    center=False,
+                )
+                assert (
+                    audio.shape[1] == mel.shape[2] * self.hop_size
+                ), f"audio shape {audio.shape} mel shape {mel.shape}"
+
+        else:
+            mel = np.load(
+                os.path.join(
+                    self.base_mels_path,
+                    os.path.splitext(os.path.split(filename)[-1])[0] + ".npy",
+                )
+            )
+            mel = torch.from_numpy(mel)
+
+            if len(mel.shape) < 3:
+                mel = mel.unsqueeze(0)
+
+            if self.split:
+                frames_per_seg = math.ceil(self.segment_size / self.hop_size)
+
+                if audio.size(1) >= self.segment_size:
+                    mel_start = random.randint(0, mel.size(2) - frames_per_seg - 1)
+                    mel = mel[:, :, mel_start : mel_start + frames_per_seg]
+                    audio = audio[
+                        :,
+                        mel_start
+                        * self.hop_size : (mel_start + frames_per_seg)
+                        * self.hop_size,
+                    ]
+                else:
+                    mel = torch.nn.functional.pad(
+                        mel, (0, frames_per_seg - mel.size(2)), "constant"
+                    )
+                    audio = torch.nn.functional.pad(
+                        audio, (0, self.segment_size - audio.size(1)), "constant"
+                    )
+
+        mel_loss = mel_spectrogram(
+            audio,
+            self.n_fft,
+            self.num_mels,
+            self.sampling_rate,
+            self.hop_size,
+            self.win_size,
+            self.fmin,
+            self.fmax_loss,
+            center=False,
+        )
+
+        return (mel.squeeze(), audio.squeeze(0), filename, mel_loss.squeeze())
+
+    def __len__(self):
+        return len(self.audio_files)

utils.py

@@ -6,6 +6,7 @@ import os
 import matplotlib
 import torch
 from torch.nn.utils import weight_norm
+
 matplotlib.use("Agg")
 import matplotlib.pylab as plt
 from meldataset import MAX_WAV_VALUE
@@ -14,8 +15,7 @@ from scipy.io.wavfile import write
 
 def plot_spectrogram(spectrogram):
     fig, ax = plt.subplots(figsize=(10, 2))
-    im = ax.imshow(spectrogram, aspect="auto", origin="lower",
-                   interpolation='none')
+    im = ax.imshow(spectrogram, aspect="auto", origin="lower", interpolation="none")
     plt.colorbar(im, ax=ax)
 
     fig.canvas.draw()
@@ -24,10 +24,16 @@ def plot_spectrogram(spectrogram):
     return fig
 
 
-def plot_spectrogram_clipped(spectrogram, clip_max=2.):
+def plot_spectrogram_clipped(spectrogram, clip_max=2.0):
     fig, ax = plt.subplots(figsize=(10, 2))
-    im = ax.imshow(spectrogram, aspect="auto", origin="lower",
-                   interpolation='none', vmin=1e-6, vmax=clip_max)
+    im = ax.imshow(
+        spectrogram,
+        aspect="auto",
+        origin="lower",
+        interpolation="none",
+        vmin=1e-6,
+        vmax=clip_max,
+    )
     plt.colorbar(im, ax=ax)
 
     fig.canvas.draw()
@@ -49,32 +55,45 @@ def apply_weight_norm(m):
 
 
 def get_padding(kernel_size, dilation=1):
-    return int((kernel_size*dilation - dilation)/2)
+    return int((kernel_size * dilation - dilation) / 2)
 
 
 def load_checkpoint(filepath, device):
     assert os.path.isfile(filepath)
-    print("Loading '{}'".format(filepath))
+    print(f"Loading '{filepath}'")
     checkpoint_dict = torch.load(filepath, map_location=device)
     print("Complete.")
     return checkpoint_dict
 
 
 def save_checkpoint(filepath, obj):
-    print("Saving checkpoint to {}".format(filepath))
+    print(f"Saving checkpoint to {filepath}")
     torch.save(obj, filepath)
     print("Complete.")
 
 
-def scan_checkpoint(cp_dir, prefix):
-    pattern = os.path.join(cp_dir, prefix + '????????')
+def scan_checkpoint(cp_dir, prefix, renamed_file=None):
+    # Fallback to original scanning logic first
+    pattern = os.path.join(cp_dir, prefix + "????????")
     cp_list = glob.glob(pattern)
-    if len(cp_list) == 0:
-        return None
-    return sorted(cp_list)[-1]
+
+    if len(cp_list) > 0:
+        last_checkpoint_path = sorted(cp_list)[-1]
+        print(f"[INFO] Resuming from checkpoint: '{last_checkpoint_path}'")
+        return last_checkpoint_path
+
+    # If no pattern-based checkpoints are found, check for renamed file
+    if renamed_file:
+        renamed_path = os.path.join(cp_dir, renamed_file)
+        if os.path.isfile(renamed_path):
+            print(f"[INFO] Resuming from renamed checkpoint: '{renamed_file}'")
+            return renamed_path
+
+    return None
+
 
 def save_audio(audio, path, sr):
     # wav: torch with 1d shape
     audio = audio * MAX_WAV_VALUE
-    audio = audio.cpu().numpy().astype('int16')
-    write(path, sr, audio)
+    audio = audio.cpu().numpy().astype("int16")
+    write(path, sr, audio)