chore: import upstream snapshot with attribution

docs/source/audio/api.rst

@@ -0,0 +1,322 @@
+NeMo Audio API
+==============
+
+Model Classes
+-------------
+Base Classes
+~~~~~~~~~~~~
+.. autoclass:: nemo.collections.audio.models.AudioToAudioModel
+    :show-inheritance:
+    :members:
+    :exclude-members: setup_training_data, setup_validation_data, training_step, on_validation_epoch_end, validation_step, setup_test_data, on_train_epoch_start
+
+
+Processing Models
+~~~~~~~~~~~~~~~~~
+.. autoclass:: nemo.collections.audio.models.EncMaskDecAudioToAudioModel
+    :show-inheritance:
+    :members:
+    :exclude-members: setup_training_data, setup_validation_data, training_step, on_validation_epoch_end, validation_step, setup_test_data, on_train_epoch_start
+
+.. autoclass:: nemo.collections.audio.models.FlowMatchingAudioToAudioModel
+    :show-inheritance:
+    :members:
+    :exclude-members: setup_training_data, setup_validation_data, training_step, on_validation_epoch_end, validation_step, setup_test_data, on_train_epoch_start
+
+.. autoclass:: nemo.collections.audio.models.PredictiveAudioToAudioModel
+    :show-inheritance:
+    :members:
+    :exclude-members: setup_training_data, setup_validation_data, training_step, on_validation_epoch_end, validation_step, setup_test_data, on_train_epoch_start
+
+.. autoclass:: nemo.collections.audio.models.ScoreBasedGenerativeAudioToAudioModel
+    :show-inheritance:
+    :members:
+    :exclude-members: setup_training_data, setup_validation_data, training_step, on_validation_epoch_end, validation_step, setup_test_data, on_train_epoch_start
+
+.. autoclass:: nemo.collections.audio.models.SchroedingerBridgeAudioToAudioModel
+    :show-inheritance:
+    :members:
+    :exclude-members: setup_training_data, setup_validation_data, training_step, on_validation_epoch_end, validation_step, setup_test_data, on_train_epoch_start
+
+
+Modules
+-------
+
+Features
+~~~~~~~~
+.. autoclass:: nemo.collections.audio.modules.features.SpectrogramToMultichannelFeatures
+    :show-inheritance:
+    :members:
+
+
+Masking
+~~~~~~~
+.. autoclass:: nemo.collections.audio.modules.masking.MaskEstimatorRNN
+    :show-inheritance:
+    :members:
+
+.. _audio-api-masking-multi-channel-mask-estimator:
+.. autoclass:: nemo.collections.audio.modules.masking.MaskEstimatorFlexChannels
+    :show-inheritance:
+    :members:
+
+.. _audio-api-masking-guided-source-separation:
+.. autoclass:: nemo.collections.audio.modules.masking.MaskEstimatorGSS
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.modules.masking.MaskReferenceChannel
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.modules.masking.MaskBasedBeamformer
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.modules.masking.MaskBasedDereverbWPE
+    :show-inheritance:
+    :members:
+
+
+Projections
+~~~~~~~~~~~
+
+.. autoclass:: nemo.collections.audio.modules.projections.MixtureConsistencyProjection
+    :show-inheritance:
+    :members:
+
+
+SSL Pretraining
+~~~~~~~~~~~~~~~
+
+.. autoclass:: nemo.collections.audio.modules.ssl_pretrain_masking.SSLPretrainWithMaskedPatch
+    :show-inheritance:
+    :members:
+
+
+Transforms
+~~~~~~~~~~
+
+.. _audio-api-audio-to-spectrogram:
+.. autoclass:: nemo.collections.audio.modules.transforms.AudioToSpectrogram
+    :show-inheritance:
+    :members:
+
+.. _audio-api-spectrogram-to-audio:
+.. autoclass:: nemo.collections.audio.modules.transforms.SpectrogramToAudio
+    :show-inheritance:
+    :members:
+
+
+Parts
+-----
+
+Submodules: Diffusion
+~~~~~~~~~~~~~~~~~~~~~
+.. autoclass:: nemo.collections.audio.parts.submodules.diffusion.StochasticDifferentialEquation
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.diffusion.OrnsteinUhlenbeckVarianceExplodingSDE
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.diffusion.ReverseStochasticDifferentialEquation
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.diffusion.PredictorCorrectorSampler
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.diffusion.Predictor
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.diffusion.ReverseDiffusionPredictor
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.diffusion.Corrector
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.diffusion.AnnealedLangevinDynamics
+    :show-inheritance:
+    :members:
+
+
+Submodules: Flow
+~~~~~~~~~~~~~~~~
+.. autoclass:: nemo.collections.audio.parts.submodules.flow.ConditionalFlow
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.flow.OptimalTransportFlow
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.flow.ConditionalFlowMatchingSampler
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.flow.ConditionalFlowMatchingEulerSampler
+    :show-inheritance:
+    :members:
+
+Submodules: Multichannel
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: nemo.collections.audio.parts.submodules.multichannel.ChannelAugment
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.multichannel.TransformAverageConcatenate
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.multichannel.TransformAttendConcatenate
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.multichannel.ChannelAveragePool
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.multichannel.ChannelAttentionPool
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.multichannel.ParametricMultichannelWienerFilter
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.multichannel.ReferenceChannelEstimatorSNR
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.multichannel.WPEFilter
+    :show-inheritance:
+    :members:
+
+
+Submodules: NCSN++
+~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: nemo.collections.audio.parts.submodules.ncsnpp.SpectrogramNoiseConditionalScoreNetworkPlusPlus
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.ncsnpp.NoiseConditionalScoreNetworkPlusPlus
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.ncsnpp.GaussianFourierProjection
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.ncsnpp.ResnetBlockBigGANPlusPlus
+    :show-inheritance:
+    :members:
+
+
+Submodules: Schrödinger Bridge
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: nemo.collections.audio.parts.submodules.schroedinger_bridge.SBNoiseSchedule
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.schroedinger_bridge.SBNoiseScheduleVE
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.schroedinger_bridge.SBNoiseScheduleVP
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.schroedinger_bridge.SBSampler
+    :show-inheritance:
+    :members:
+
+
+Submodules: TransformerUNet
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: nemo.collections.audio.parts.submodules.transformerunet.LearnedSinusoidalPosEmb
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.transformerunet.ConvPositionEmbed
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.transformerunet.RMSNorm
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.transformerunet.AdaptiveRMSNorm
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.transformerunet.GEGLU
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.transformerunet.TransformerUNet
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.parts.submodules.transformerunet.SpectrogramTransformerUNet
+    :show-inheritance:
+    :members:
+
+
+Losses
+------
+
+.. autoclass:: nemo.collections.audio.losses.MAELoss
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.losses.MSELoss
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.losses.SDRLoss
+    :show-inheritance:
+    :members:
+
+
+Datasets
+--------
+
+NeMo Format
+~~~~~~~~~~~
+
+.. autoclass:: nemo.collections.audio.data.audio_to_audio.BaseAudioDataset
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.data.audio_to_audio.AudioToTargetDataset
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.data.audio_to_audio.AudioToTargetWithReferenceDataset
+    :show-inheritance:
+    :members:
+
+.. autoclass:: nemo.collections.audio.data.audio_to_audio.AudioToTargetWithEmbeddingDataset
+    :show-inheritance:
+    :members:
+
+
+Lhotse Format
+~~~~~~~~~~~~~
+
+.. autoclass:: nemo.collections.audio.data.audio_to_audio_lhotse.LhotseAudioToTargetDataset
+    :show-inheritance:
+    :members:
+
+
+

docs/source/audio/audio_all.bib

@@ -0,0 +1,64 @@
+@inproceedings{jukic2023flexible,
+  title={Flexible multichannel speech enhancement for noise-robust frontend},
+  author={Jukić, Ante and Balam, Jagadeesh and Ginsburg, Boris},
+  booktitle={Proc. WASPAA},
+  year={2023}
+}
+
+@inproceedings{ito2016directional,
+  title={Complex angular central Gaussian mixture model for directional statistics in mask-based microphone array signal processing},
+  author={Ito, Nobutaka and Araki, Shoko and Nakatani, Tomohiro},
+  booktitle={Proc. EUSIPCO},
+  year={2016}
+}
+
+@inproceedings{jukic2024sb,
+  title={Schrödinger Bridge for Generative Speech Enhancement},
+  author={Ante Juki\'{c} and Roman Korostik and Jagadeesh Balam and Boris Ginsburg},
+  year={2024},
+  pages={1175-1179},
+  booktitle={Proc. Interspeech}
+}
+
+@inproceedings{welker2022speech,
+  author={Simon Welker and Julius Richter and Timo Gerkmann},
+  title={Speech Enhancement with Score-Based Generative Models in the Complex {STFT} Domain},
+  year={2022},
+  pages={2928-2932},
+  booktitle={Proc. Interspeech}
+}
+
+@article{richter2023sgmse,
+  author  = {Richter, Julius and Welker, Simon and Lemercier, Jean-Marie and Lay, Bunlong and Gerkmann, Timo},
+  title   = {{Speech Enhancement and Dereverberation with Diffusion-Based Generative Models}},
+  journal = {IEEE/ACM Trans. on Audio, Speech, and Language Process.},
+  volume  = {31},
+  pages   = {2351-2364},
+  year    = {2023}
+}
+
+@article{ku2024generative,
+  title={Generative Speech Foundation Model Pretraining for High-Quality Speech Extraction and Restoration},
+  author={Pin-Jui Ku and Alexander H. Liu and Roman Korostik and Sung-Feng Huang and Szu-Wei Fu and Ante Jukić},
+  journal={arXiv preprint arXiv:2409.16117},
+  year={2024},
+}
+
+@souden{souden2010,
+  author={Souden, Mehrez and Benesty, Jacob and Affes, SofiÈne},
+  journal={IEEE Transactions on Audio, Speech, and Language Processing}, 
+  title={On Optimal Frequency-Domain Multichannel Linear Filtering for Noise Reduction}, 
+  year={2010},
+  volume={18},
+  number={2},
+  pages={260-276}
+}
+
+@inproceedings{
+lipman2023flow,
+title={Flow Matching for Generative Modeling},
+author={Yaron Lipman and Ricky T. Q. Chen and Heli Ben-Hamu and Maximilian Nickel and Matthew Le},
+booktitle={Proc. ICLR},
+year={2023},
+url={https://openreview.net/forum?id=PqvMRDCJT9t}
+}

docs/source/audio/checkpoints.rst

@@ -0,0 +1,59 @@
+Checkpoints
+===========
+
+There are two main ways to load pretrained checkpoints in NeMo:
+
+* Using the :code:`restore_from()` method to load a local checkpoint file (``.nemo``), or
+* Using the :code:`from_pretrained()` method to download and set up a checkpoint from the cloud.
+
+Note that these instructions are for loading fully trained checkpoints for evaluation or fine-tuning. For resuming an unfinished
+training experiment, use the Experiment Manager to do so by setting the ``resume_if_exists`` flag to ``True``.
+
+Local Checkpoints
+-------------------------
+
+* **Save Model Checkpoints**: NeMo automatically saves final model checkpoints with ``.nemo`` suffix. You could also manually save any model checkpoint using :code:`model.save_to(<checkpoint_path>.nemo)`.
+* **Load Model Checkpoints**: if you'd like to load a checkpoint saved at ``<path/to/checkpoint/file.nemo>``, use the :code:`restore_from()` method below, where ``<MODEL_BASE_CLASS>`` is the model class of the original checkpoint.
+
+.. code-block:: python
+
+    import nemo.collections.audio as nemo_audio
+    model = nemo_audio.models.<MODEL_BASE_CLASS>.restore_from(restore_path="<path/to/checkpoint/file.nemo>")
+
+Pretrained Checkpoints
+----------------------
+
+The table below in :ref:`checkpoints_audio_models` list part of available pre-trained audio processing models including speech processing, restoration and extraction.
+
+Load Model Checkpoints
+^^^^^^^^^^^^^^^^^^^^^^
+The models can be accessed via the :code:`from_pretrained()` method inside the audio model class. In general, you can load any of these models with code in the following format,
+
+.. code-block:: python
+
+    import nemo.collections.audio as nemo_audio
+    model = nemo_audio.models.<MODEL_BASE_CLASS>.from_pretrained(model_name="<MODEL_NAME>")
+
+where ``<MODEL_NAME>`` is the value in ``Model Name`` column in the tables in :ref:`checkpoints_audio_models`. These names are predefined in the each model's member function ``self.list_available_models()``. 
+
+
+.. _checkpoints_audio_models:
+
+Audio Models
+------------
+
+Speech Enhancement Models
+^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. csv-table::
+   :file: data/checkpoints_se.csv
+   :align: left
+   :header-rows: 1
+
+SSL Models
+^^^^^^^^^^
+
+.. csv-table::
+   :file: data/checkpoints_ssl.csv
+   :align: left
+   :header-rows: 1

docs/source/audio/configs.rst

@@ -0,0 +1,427 @@
+NeMo Audio Configuration Files
+==============================
+This section describes the NeMo configuration file setup that is specific to models in the audio collection.
+For general information about how to set up and run experiments that is common to all NeMo models (e.g. Experiment Manager and PyTorch Lightning trainer parameters), see the :doc:`../core/core` section.
+
+The model section of the NeMo audio configuration files generally requires information about the dataset(s) being used, parameters for any augmentation being performed, as well as the model architecture specification.
+
+Example configuration files for all of the NeMo audio models can be found in the
+`config directory of the examples <https://github.com/NVIDIA/NeMo/tree/main/examples/audio/conf>`_.
+
+
+.. _audio-configs-nemo-dataset-configuration:
+
+NeMo Dataset Configuration
+--------------------------
+
+Training, validation, and test parameters are specified using the ``model.train_ds``, ``model.validation_ds``, and ``model.test_ds`` sections in the configuration file, respectively.
+Depending on the task, there may be arguments specifying the sample rate or duration of the loaded audio examples. Some fields can be left out and specified via the command-line at runtime.
+Refer to the `Dataset Processing Classes <./api.html#datasets>`__ section of the API for a list of datasets classes and their respective parameters.
+An example train, validation and test datasets can be configured as follows:
+
+.. code-block:: yaml
+
+  model:
+    sample_rate: 16000
+    skip_nan_grad: false
+
+    train_ds:
+      manifest_filepath: ???
+      input_key: audio_filepath # key of the input signal path in the manifest
+      target_key: target_filepath # key of the target signal path in the manifest
+      target_channel_selector: 0 # target signal is the first channel from files in target_key
+      audio_duration: 4.0 # in seconds, audio segment duration for training
+      random_offset: true # if the file is longer than audio_duration, use random offset to select a subsegment
+      min_duration: ${model.train_ds.audio_duration}
+      batch_size: 64 # batch size may be increased based on the available memory
+      shuffle: true
+      num_workers: 8
+      pin_memory: true
+
+    validation_ds:
+      manifest_filepath: ???
+      input_key: audio_filepath # key of the input signal path in the manifest
+      target_key: target_filepath # key of the target signal path in the manifest
+      target_channel_selector: 0 # target signal is the first channel from files in target_key
+      batch_size: 64 # batch size may be increased based on the available memory
+      shuffle: false
+      num_workers: 4
+      pin_memory: true
+
+    test_ds:
+      manifest_filepath: ???
+      input_key: audio_filepath # key of the input signal path in the manifest
+      target_key: target_filepath # key of the target signal path in the manifest
+      target_channel_selector: 0 # target signal is the first channel from files in target_key
+      batch_size: 1 # batch size may be increased based on the available memory
+      shuffle: false
+      num_workers: 4
+      pin_memory: true
+
+More information about online augmentation can found in the `masking example configuration <https://github.com/NVIDIA/NeMo/blob/main/examples/audio/conf/masking.yaml>`_
+
+
+.. _audio-configs-lhotse-dataset-configuration:
+
+Lhotse Dataset Configuration
+----------------------------
+
+Lhotse CutSet
+~~~~~~~~~~~~~
+
+An example train dataset in Lhotse CutSet format can be configured as follows:
+
+.. code-block:: yaml
+
+  train_ds:
+    use_lhotse: true # enable Lhotse data loader
+    cuts_path: ??? # path to Lhotse cuts manifest with input signals and the corresponding target signals (target signals should be in the custom "target_recording" field)
+    truncate_duration: 4.00 # truncate audio to 4 seconds
+    truncate_offset_type: random # if the file is longer than truncate_duration, use random offset to select a subsegment
+    batch_size: 64 # batch size may be increased based on the available memory
+    shuffle: true
+    num_workers: 8
+    pin_memory: true
+
+Lhotse CutSet with Online Augmentation
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+An example train dataset in Lhotse CutSet format using online augmentation with room impulse response (RIR) convolution and additive noise can be configured as follows:
+
+.. code-block:: yaml
+
+  train_ds:
+    use_lhotse: true # enable Lhotse data loader
+    cuts_path: ??? # path to Lhotse cuts manifest with speech signals for augmentation (including custom "target_recording" field with the same signals)
+    truncate_duration: 4.00 # truncate audio to 4 seconds
+    truncate_offset_type: random # if the file is longer than truncate_duration, use random offset to select a subsegment
+    batch_size: 64 # batch size may be increased based on the available memory
+    shuffle: true
+    num_workers: 8
+    pin_memory: true
+    rir_enabled: true # enable room impulse response augmentation
+    rir_path: ??? # path to Lhotse recordings manifest with room impulse response signals
+    noise_path: ??? # path to Lhotse cuts manifest with noise signals
+
+A configuration file with Lhotse online augmentation can found in the `online augmentation example configuration <https://github.com/NVIDIA/NeMo/blob/main/examples/audio/conf/masking_with_online_augmentation.yaml>`_.
+More information about the online augmentation can be found in the `tutorial notebook <https://github.com/NVIDIA/NeMo/blob/main/tutorials/audio/speech_enhancement/Speech_Enhancement_with_Online_Augmentation.ipynb>`_.
+
+
+Lhotse Shar
+~~~~~~~~~~~
+
+An example train dataset in Lhotse shar format can be configured as follows:
+
+.. code-block:: yaml
+
+  train_ds:
+    shar_path: ???
+    use_lhotse: true
+    truncate_duration: 4.00 # truncate audio to 4 seconds
+    truncate_offset_type: random
+    batch_size: 8 # batch size may be increased based on the available memory
+    shuffle: true
+    num_workers: 8
+    pin_memory: true
+
+
+A configuration file with Lhotse shar format can found in the `SSL pretraining example configuration <https://github.com/NVIDIA/NeMo/blob/main/examples/audio/conf/flow_matching_generative_ssl_pretraining.yaml>`_.
+
+
+Dataset Reweighting with Temperature
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When combining multiple datasets using nested ``input_cfg`` groups, you can control the sampling distribution using the ``reweight_temperature`` parameter. This feature allows you to balance dataset sampling without manually recalculating weights **when adding or removing datasets**.
+
+The temperature scaling formula is:
+
+.. math::
+
+    \hat{w}_i = \frac{w_i^{\tau}}{\sum_{j} w_j^{\tau}}
+
+where :math:`w_i` is the original weight of dataset :math:`i`, :math:`\tau` is the temperature, and :math:`\hat{w}_i` is the normalized sampling probability.
+
+**How Temperature Works:**
+
+- ``temperature = 1.0``: Preserves original weight ratios (neutral, no reweighting)
+- ``temperature = 0.0``: Equalizes all datasets (each gets equal probability regardless of original weights)
+- ``0 < temperature < 1.0``: Over-samples smaller datasets relative to larger ones
+- ``temperature > 1.0``: Amplifies differences between dataset weights
+
+**Configuration Options:**
+
+The ``reweight_temperature`` parameter accepts two formats:
+
+1. **Scalar value** (applied to all nesting levels, warning logged):
+
+.. code-block:: yaml
+
+    train_ds:
+      use_lhotse: true
+      reweight_temperature: 0.5  # Applied to all levels, warning logged
+      input_cfg:
+        - type: group
+          input_cfg:
+            - type: lhotse_shar
+              shar_path: /path/to/dataset1
+              weight: 900
+            - type: lhotse_shar
+              shar_path: /path/to/dataset2
+              weight: 100
+        - type: lhotse_shar
+          shar_path: /path/to/dataset3
+          weight: 200
+        - type: nemo_tarred
+          manifest_filepath: /path/to/dataset4/manifest.json
+          tarred_audio_filepath: /path/to/dataset4/audio.tar
+          weight: 300
+
+2. **List matching maximum nesting depth** (one temperature per level):
+
+.. code-block:: yaml
+
+    train_ds:
+      use_lhotse: true
+      reweight_temperature: [1.0, 0.0]  # Level 1: preserve ratios, Level 2: equalize
+      input_cfg:
+        - type: group
+          weight: 0.7
+          input_cfg:
+            - type: lhotse_shar
+              shar_path: /path/to/dataset1
+              weight: 600
+            - type: lhotse_shar
+              shar_path: /path/to/dataset2
+              weight: 400
+        - type: group
+          weight: 0.3
+          input_cfg:
+            - type: lhotse_shar
+              shar_path: /path/to/dataset3
+              weight: 100
+
+.. note::
+
+    If ``reweight_temperature`` is provided as a list, its length **must** exactly match the maximum nesting depth of ``input_cfg``.
+    A mismatch (too few or too many values) raises a ``ValueError``.
+    Use a scalar value instead if you want the same temperature applied uniformly to all levels.
+
+**Maximum Nesting Depth Calculation:**
+
+The maximum nesting depth is calculated as the maximum depth of ``input_cfg`` keys in the configuration. Sibling groups at the same level share the same temperature value.
+
+.. code-block:: yaml
+
+    # This has maximum nesting depth = 2
+    input_cfg:                          # Level 1
+      - type: group
+        input_cfg:                      # Level 2
+          - type: lhotse_shar
+      - type: group                     # Same level as above (sibling)
+        input_cfg:                      # Level 2 (same as above)
+          - type: lhotse_shar
+
+When ``input_cfg`` is overridden via CLI to a YAML file path (e.g.
+``model.train_ds.input_cfg=train_all.yaml``), the depth calculation loads the
+referenced file and traverses its contents to count nested ``input_cfg`` keys.
+This also works with multi-level file references:
+
+.. code-block:: yaml
+
+    # train_all.yaml (referenced via input_cfg=train_all.yaml)
+    - type: group
+      weight: 100
+      input_cfg: ${oc.env:MANIFEST_ROOT}/train_en.yaml   # resolved at runtime
+    - type: group
+      weight: 200
+      input_cfg: ${oc.env:MANIFEST_ROOT}/train_de.yaml
+
+.. note::
+
+    Paths containing OmegaConf interpolations (e.g. ``${oc.env:MANIFEST_ROOT}``)
+    cannot be resolved during depth counting -- they are resolved later at runtime
+    by ``OmegaConf.create()``. Such paths are treated as a single additional
+    nesting level.
+
+**Example: Balancing Multiple Task Groups**
+
+.. code-block:: yaml
+
+    train_ds:
+      use_lhotse: true
+      reweight_temperature: [1.0, 0.0]  # Level 1: Preserve task ratios, Level 2: Equalize within tasks
+      input_cfg:
+        - type: group
+          weight: 0.7
+          tags:
+            task: asr
+          input_cfg:
+            - type: nemo_tarred
+              manifest_filepath: /path/to/asr1/manifest.json
+              tarred_audio_filepath: /path/to/asr1/audio.tar
+              weight: 600  # Large dataset
+            - type: nemo_tarred
+              manifest_filepath: /path/to/asr2/manifest.json
+              tarred_audio_filepath: /path/to/asr2/audio.tar
+              weight: 100  # Small dataset (will be upsampled with temp=0.0)
+        - type: group
+          weight: 0.3
+          tags:
+            task: ast
+          input_cfg:
+            - type: nemo_tarred
+              manifest_filepath: /path/to/ast1/manifest.json
+              tarred_audio_filepath: /path/to/ast1/audio.tar
+              weight: 50
+            - type: nemo_tarred
+              manifest_filepath: /path/to/ast2/manifest.json
+              tarred_audio_filepath: /path/to/ast2/audio.tar
+              weight: 200
+
+In this example:
+
+- Level 1 temperature is ``1.0``: The 70/30 split between ASR and AST groups is preserved
+- Level 2 temperature is ``0.0``: Within each group, all datasets are sampled equally regardless of their original weights
+
+
+Model Architecture Configuration
+--------------------------------
+Each configuration file should describe the model architecture being used for the experiment.
+An example of a simple predictive model configuration is shown below:
+
+.. code-block:: yaml
+
+  model:
+    type: predictive
+    sample_rate: 16000
+    skip_nan_grad: false
+    num_outputs: 1
+    normalize_input: true # normalize the input signal to 0dBFS
+
+    train_ds:
+      manifest_filepath: ???
+      input_key: noisy_filepath
+      target_key: clean_filepath
+      audio_duration: 2.00 # trim audio to 2 seconds
+      random_offset: true
+      normalization_signal: input_signal
+      batch_size: 8 # batch size may be increased based on the available memory
+      shuffle: true
+      num_workers: 8
+      pin_memory: true
+
+    validation_ds:
+      manifest_filepath: ???
+      input_key: noisy_filepath
+      target_key: clean_filepath
+      batch_size: 8
+      shuffle: false
+      num_workers: 4
+      pin_memory: true
+
+    encoder:
+      _target_: nemo.collections.audio.modules.transforms.AudioToSpectrogram
+      fft_length: 510 # Number of subbands in the STFT = fft_length // 2 + 1 = 256
+      hop_length: 128
+      magnitude_power: 0.5
+      scale: 0.33
+
+    decoder:
+      _target_: nemo.collections.audio.modules.transforms.SpectrogramToAudio
+      fft_length: ${model.encoder.fft_length}
+      hop_length: ${model.encoder.hop_length}
+      magnitude_power: ${model.encoder.magnitude_power}
+      scale: ${model.encoder.scale}
+
+    estimator:
+      _target_: nemo.collections.audio.parts.submodules.ncsnpp.SpectrogramNoiseConditionalScoreNetworkPlusPlus
+      in_channels: 1 # single-channel noisy input
+      out_channels: 1 # single-channel estimate
+      num_res_blocks: 3 # increased number of res blocks
+      pad_time_to: 64 # pad to 64 frames for the time dimension
+      pad_dimension_to: 0 # no padding in the frequency dimension
+
+    loss:
+      _target_: nemo.collections.audio.losses.MSELoss # computed in the time domain
+
+    metrics:
+      val:
+        sisdr: # output SI-SDR
+          _target_: torchmetrics.audio.ScaleInvariantSignalDistortionRatio
+
+    optim:
+      name: adam
+      lr: 1e-4
+      # optimizer arguments
+      betas: [0.9, 0.999]
+      weight_decay: 0.0
+
+
+Complete configuration file can found in the `example configuration <https://github.com/NVIDIA/NeMo/blob/main/examples/audio/conf/predictive.yaml>`_.
+
+
+Finetuning Configuration
+--------------------------
+
+All scripts support easy finetuning by partially/fully loading the pretrained weights from a checkpoint into the currently instantiated model.
+Note that the currently instantiated model should have parameters that match the pre-trained checkpoint so the weights may load properly.
+
+Pre-trained weights can be provided by:
+
+* Providing a path to a NeMo model (via ``init_from_nemo_model``)
+* Providing a name of a pretrained NeMo model (which will be downloaded via the cloud) (via ``init_from_pretrained_model``)
+
+
+Training from scratch
+~~~~~~~~~~~~~~~~~~~~~
+
+A model can be trained from scratch using the following command:
+
+.. code-block:: shell
+
+    python examples/audio/audio_to_audio_train.py \
+        --config-path=<path to dir of configs>
+        --config-name=<name of config without .yaml>) \
+        model.train_ds.manifest_filepath="<path to manifest file>" \
+        model.validation_ds.manifest_filepath="<path to manifest file>" \
+        trainer.devices=1 \
+        trainer.accelerator='gpu' \
+        trainer.max_epochs=50
+
+Fine-tuning via a NeMo model
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+A model can be finetuned from an existing NeMo model using the following command:
+
+.. code-block:: shell
+    :emphasize-lines: 9
+
+    python examples/audio/audio_to_audio_train.py \
+        --config-path=<path to dir of configs>
+        --config-name=<name of config without .yaml>) \
+        model.train_ds.manifest_filepath="<path to manifest file>" \
+        model.validation_ds.manifest_filepath="<path to manifest file>" \
+        trainer.devices=1 \
+        trainer.accelerator='gpu' \
+        trainer.max_epochs=50 \
+        +init_from_nemo_model="<path to .nemo model file>"
+
+
+Fine-tuning via a NeMo pretrained model name
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+A model can be finetuned from an pre-trained NeMo model using the following command:
+
+.. code-block:: shell
+    :emphasize-lines: 9
+
+    python examples/audio/audio_to_audio_train.py \
+        --config-path=<path to dir of configs>
+        --config-name=<name of config without .yaml>) \
+        model.train_ds.manifest_filepath="<path to manifest file>" \
+        model.validation_ds.manifest_filepath="<path to manifest file>" \
+        trainer.devices=1 \
+        trainer.accelerator='gpu' \
+        trainer.max_epochs=50 \
+        +init_from_pretrained_model="<name of pretrained checkpoint>"
+

docs/source/audio/data/checkpoints_se.csv

@@ -0,0 +1,3 @@
+Model Name,Dataset,Sampling Rate,Model Class,Model Card
+nvidia/se_den_sb_16k_small,WSJ0+CHiME,16000Hz,nemo.collections.audio.models.SchroedingerBridgeAudioToAudioModel,`se_den_sb_16k_small <https://huggingface.co/nvidia/se_den_sb_16k_small>`_
+nvidia/se_der_sb_16k_small,WSJ0+Reverb,16000Hz,nemo.collections.audio.models.SchroedingerBridgeAudioToAudioModel,`se_der_sb_16k_small <https://huggingface.co/nvidia/se_der_sb_16k_small>`_

docs/source/audio/data/checkpoints_ssl.csv

@@ -0,0 +1,2 @@
+Model Name,Dataset,Sampling Rate,Model Class,Model Card
+nvidia/sr_ssl_flowmatching_16k_430m,Libri-Light,16000Hz,nemo.collections.audio.models.FlowMatchingAudioToAudioModel,`sr_ssl_flowmatching_16k_430m <https://huggingface.co/nvidia/sr_ssl_flowmatching_16k_430m>`_

docs/source/audio/datasets.rst

@@ -0,0 +1,91 @@
+Datasets
+========
+
+The `audio` collection expect the training, validation and tests datasets in either NeMo format or Lhotse format.
+
+NeMo Format
+-----------
+
+Each dataset consists of a set of utterances in individual audio files plus a manifest that describes the dataset, with information about one utterance per line (``.json``).
+There should be one manifest file per dataset that will be passed in, therefore, if the user wants separate training and validation datasets, they should also have separate manifests. Otherwise, they will be loading validation data with their training data and vice versa.
+
+
+In most applications, such as speech restoration or enhancement, the model aims to transform the input audio signal into the target audio signal. In this case, each line of the manifest should have the following structure:
+
+.. code-block:: json
+
+  {"input_filepath": "/path/to/input_audio.wav", "target_filepath": "/path/to/target_audio.wav", "duration": 3.147}
+
+The :code:`input_filepath` field should provide either an absolute path to the audio file corresponding to the utterance, or a relative path with respect to the directory containing the manifest.
+The :code:`target_filepath` field should provide either an absolute path to the audio file corresponding to the utterance, or a relative path with respect to the directory containing the manifest.
+Note that keys for input and target audio can be custom, and can be configured in the model configuration file, as described in :ref:`Configs <audio-configs-nemo-dataset-configuration>`.
+
+Each entry in the manifest (describing one audio file) should be bordered by ``"{"`` and ``"}"`` and must be placed on one line. The ``"key": value`` pairs should be separated by a commas as shown above. NeMo enforces no blank lines in the manifest so that the total number of lines indicates the total number of audio files in the dataset.
+
+Once there is a manifest that describes each audio file in the dataset, assign the ``JSON`` manifest file path in the experiment config file, for example, ``training_ds.manifest_filepath=<path/to/manifest.json>``.
+
+For more information about the individual tarred datasets and the parameters available, including shuffling options, see the corresponding class APIs in the `Datasets <./api.html#datasets>`_ section.
+
+
+Lhotse Format
+-------------
+
+NeMo supports using `Lhotse`_, a speech data handling library, as a dataloading option.
+Lhotse can be easily enabled using `use_lhotse=True` in the dataset configuration file, as described in :ref:`Lhotse dataset configuration <audio-configs-lhotse-dataset-configuration>`.
+
+Lhotse dataloading supports the following types of inputs:
+
+* Lhotse CutSet manifests
+    Regular Lhotse CutSet manifests (typically gzipped JSONL).
+    See `Lhotse Cuts documentation`_ to learn more about Lhotse data formats.
+* Lhotse Shar data
+    Lhotse Shar is a data format that also uses tar files for sequential data loading,
+    but is designed to be modular (i.e., easily extensible with new data sources and with new feature fields).
+    More details can be found in the tutorial notebook: |tutorial_shar|.
+
+
+.. _Lhotse: https://github.com/lhotse-speech/lhotse
+.. _Lhotse Cuts documentation: https://lhotse.readthedocs.io/en/latest/cuts.html
+.. |tutorial_shar| image:: https://colab.research.google.com/assets/colab-badge.svg
+    :target: https://colab.research.google.com/github/lhotse-speech/lhotse/blob/master/examples/04-lhotse-shar.ipynb
+
+
+
+Converting NeMo manifest to Lhotse
+----------------------------------
+
+A dataset with a manifest in NeMo format can be converted to Lhotse format using the provided `conversion script <https://github.com/NVIDIA/NeMo/blob/main/scripts/audio_to_audio/convert_nemo_to_lhotse.py>`_.
+
+.. code:: shell
+
+    CONVERT_SCRIPT=scripts/audio_to_audio/convert_nemo_to_lhotse.py
+    INPUT_MANIFEST=/path/to/data/nemo_manifest.json
+    OUTPUT_MANIFEST=/path/to/data/lhotse_manifest.jsonl
+
+    python ${CONVERT_SCRIPT} ${INPUT_MANIFEST} ${OUTPUT_MANIFEST} -i input_filepath -t target_filepath
+
+
+Creating Lhotse shar dataset
+----------------------------
+
+First, convert the NeMo manifest to Lhotse format with absolute paths.
+
+.. code:: shell
+
+    CONVERT_SCRIPT=scripts/audio_to_audio/convert_nemo_to_lhotse.py
+    INPUT_MANIFEST=/path/to/data/nemo_manifest.json
+    OUTPUT_MANIFEST=/path/to/data/lhotse_manifest_absolute_paths.jsonl
+
+    # convert
+    python ${CONVERT_SCRIPT} ${INPUT_MANIFEST} ${OUTPUT_MANIFEST} -i input_filepath -t target_filepath --force_absolute_paths
+
+
+Then, create the Lhotse shar dataset.
+
+.. code:: shell
+
+    LHOTSE_MANIFEST=/path/to/data/lhotse_manifest_absolute_paths.jsonl
+    OUTPUT_DIR=/path/to/data/shar
+
+    # create shars, each with 2084 examples, flac audio format
+    lhotse shar export --num-jobs 16 --verbose --shard-size 2084 --audio flac ${LHOTSE_MANIFEST} ${OUTPUT_DIR}

docs/source/audio/intro.rst

@@ -0,0 +1,17 @@
+Speech and Audio Processing
+===========================
+
+Speech and audio processing refers to a system that processes audio signals, such as speech, music, and environmental sounds. This collection includes models for speech enhancement, restoration and extraction.
+
+We will illustrate details in the following sections.
+
+.. toctree::
+    :maxdepth: 2
+
+    models
+    datasets
+    checkpoints
+    configs
+    api
+
+.. include:: resources.rst

docs/source/audio/models.rst

@@ -0,0 +1,46 @@
+Models
+=======
+This section provides a brief overview of models that NeMo's audio collection currently supports.
+
+* **Model Recipes** can be accessed through `examples/audio <https://github.com/NVIDIA/NeMo/tree/main/examples/audio>`_.
+* **Configuration Files** can be found in the directory of `examples/audio/conf <https://github.com/NVIDIA/NeMo/tree/main/examples/audio/conf>`_. For detailed information about configuration files and how they
+  should be structured, please refer to the section :doc:`./configs`.
+* **Pretrained Model Checkpoints** are available for any users for immediately synthesizing speech or fine-tuning models on
+  your custom datasets. Please follow the section :doc:`./checkpoints` for instructions on how to use those pretrained models.
+
+
+Encoder-Mask-Decoder Model
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+Encoder-Mask-Decoder model is a general model consisting of an encoder, a mask estimator, a mask processor and a decoder. The encoder processes the input audio signal and produces a latent representation. The mask estimator estimates the mask from the latent representation. The mask processor processes the mask and the latent representation to produce a processed latent representation. The decoder processes the processed latent representation to produce the output audio signal. The model can be used for various tasks such as speech enhancement or speech separation.
+The encoder and decoder can be learned or fixed, such as the :ref:`short-time Fourier transform (STFT) <audio-api-audio-to-spectrogram>` and :ref:`inverse STFT <audio-api-spectrogram-to-audio>` modules, respectively.
+The mask estimator can be a neural model, such as :ref:`multi-channel mask estimator <audio-api-masking-multi-channel-mask-estimator>` :cite:`audio-models-jukic2023flexible` or a non-neural model, such as :ref:`guided source separation (GSS) <audio-api-masking-guided-source-separation>` :cite:`audio-models-ito2016directional`.
+The mask processor can be either simple masking, or a parametric multichannel Wiener filter :cite:`audio-models-jukic2023flexible`.
+
+
+Predictive Model
+~~~~~~~~~~~~~~~~
+Predictive model is similar to the encoder-mask-decoder model, but the mask estimator and mask processor are replaced by a neural estimator. The predictive model estimates the latent representation of the target output signal from the input audio signal :cite:`audio-models-richter2023sgmse,audio-models-jukic2024sb`. The model can be used for various tasks such as speech enhancement or speech separation.
+
+
+Score-Based Generative Model
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Score-based generative model is a diffusion-based generative model that estimates the score function of the data distribution :cite:`audio-models-welker2022speech,audio-models-richter2023sgmse`. The model consists of an encoder and decoder, a neural score estimator, a stochastic differential equation (SDE) model and a sampler.
+
+
+Schrödinger Bridge Model
+~~~~~~~~~~~~~~~~~~~~~~~~
+Schrödinger bridge model is a generative model using a data-to-data process to transform the input (degraded) audio signal into the target (clean) audio signal :cite:`audio-models-jukic2024sb`. The model consists of an encoder and decoder, a neural estimator, noise schedule and a sampler.
+
+
+Flow Matching Model
+~~~~~~~~~~~~~~~~~~~
+Flow matching model is a generative model using a noise-to-data process to transform the input (degraded) audio signal into the target (clean) audio signal :cite:`audio-models-ku2024generative`. The model consists of an encoder and decoder, a neural estimator, a flow model and a sampler.
+
+
+References
+----------
+
+.. bibliography:: audio_all.bib
+    :style: plain
+    :labelprefix: AUDIO-
+    :keyprefix: audio-models-

docs/source/audio/resources.rst

@@ -0,0 +1,10 @@
+Resources and Documentation
+===========================
+
+Tutorial notebooks can be found under `the audio tutorials folder <https://github.com/NVIDIA/NeMo/tree/main/tutorials/audio>`_. If you are just starting with NeMo, consider trying out the tutorials of `NeMo Primer <https://github.com/NVIDIA/NeMo/blob/main/tutorials/00_NeMo_Primer.ipynb>`_ and `NeMo Model <https://github.com/NVIDIA/NeMo/blob/main/tutorials/01_NeMo_Models.ipynb>`_. These tutorials can be run on Google Colab by specifying the link to the notebooks' GitHub pages on Colab.
+
+If you are looking for information about a particular model, or would like to find out more about the model architectures available in the directory of `nemo.collections.audio <https://github.com/NVIDIA/NeMo/tree/main/nemo/collections/audio>`_, refer to the :doc:`Models <./models>` section.
+
+Information about how to load model checkpoints (either local files or pretrained ones from NGC), as well as a list of the checkpoints available on NGC are located on the :doc:`Checkpoints <./checkpoints>` section.
+
+Documentation regarding the configuration files specific to the NeMo audio models can be found on the :doc:`Configuration Files <./configs>` section.