[Audio] SSL Pretraining framework for flow-matching model for audio processing
State: closed | Pull request: yes
Labels: name: Run CICD

name: audio

# What does this PR do ?

This PR adds a ssl-pretraining framework and a flow matching AudioToAudio model

**Collection**: audio

# Changelog 
- Add a flow matching generative model in `audio/models/enhancement.py`
- Create `audio/modules/ssl_pretrain_masking` to partially-masked out input in the ssl pre-training stage.
- Create `audio/parts/submodules/transformerunet.py` to implement a Transformer U-Net structure.
- Create `audio/parts/submodules/flow.py` for training/inferencing with the flow matching generative model
- Create `audio/parts/utils/callbacks.py` to add processed audio in the training logger.

# Usage

Model can be first be pre-trained using

```python
python examples/audio/audio_to_audio_train.py --config-name flow_matching_generative_ssl_pretraining.yaml
```
- Please update the `train_ds.manifest_filepath` and `validation_ds.manifest_filepath` with the correct filepaths

After the ssl-pretraining stage is completed, the pretrained model could be finetuned using 

```python
python examples/audio/audio_to_audio_train.py --config-name flow_matching_generative_finetuning.yaml
```
- Remember to update the `init_from_nemo_model` with the pretrained checkpoint path in the config file.

Alternatively, one can also just train a flow matching generative from scratch:
```python
python examples/audio/audio_to_audio_train.py --config-name flow_matching_generative.yaml
```

Model can be evaluated using
```python
python examples/audio/audio_to_audio_eval.py \
    model_path=${PATH_TO_NEMO_FILE} \
    dataset_manifest=${PATH_TO_DATASET_MANIFEST} \
    output_dir=${PATH_TO_OUTPUT_DIR} \
    input_key=noisy_filepath \
    target_key=clean_filepath \
    metrics=[sdr,estoi]
```

Sampler setting can be easily overridden, e.g., to change the number of sampling steps as

```python
python examples/audio/audio_to_audio_eval.py ... ++sampler.num_steps=${NUM_STEPS}
```


# GitHub Actions CI

The Jenkins CI system has been replaced by GitHub Actions self-hosted runners.

The GitHub Actions CI will run automatically when the "Run CICD" label is added to the PR.
To re-run CI remove and add the label again.
To run CI on an untrusted fork, a NeMo user with write access must first click "Approve and run".

# Before your PR is "Ready for review"
**Pre checks**:
- [x] Make sure you read and followed [Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md)
- [ ] Did you write any new necessary tests?
- [] Did you add or update any necessary documentation?
- [ ] Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)
  - [ ] Reviewer: Does the PR have correct import guards for all optional libraries?
  
**PR Type**:
- [x] New Feature
- [ ] Bugfix
- [ ] Documentation

If you haven't finished some of the above items you can still open "Draft" PR.


## Who can review?

Anyone in the community is free to review the PR once the checks have passed. 
[Contributor guidelines](https://github.com/NVIDIA/NeMo/blob/main/CONTRIBUTING.md) contains specific people who can review PRs to various areas.