[TTS] MagpieTTS Inference Refactoring

[subhankar-ghosh]

Important

The Update branch button must only be pressed in very rare occassions.
An outdated branch is never blocking the merge of a PR.
Please reach out to the automation team before pressing that button.

What does this PR do ?

Copy of [TTS] MagpieTTS Inference Changes #15158
Because after rebasing with main git showed the PR had 231 files changed, which was wrong. So I created a fresh PR. This shows the correct number of files changed.

Collection: [Note which collection this PR will affect]

Changelog

1. Refactored Inference script
   a) Modularized code
   b) Separated inference and evaluation logic
   c) Separated different configs
   d) Lays the groundwork for longform inference

2. infer_batch returns dataclass object now

3. Moved inference code from scripts to examples

4. Added fixed context embeddings instead of getting it through context encoder, in case checkpoint does not have context encoder

5. Removed evalset_config.py and replaced with evalset.json

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this 

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:

• Make sure you read and followed Contributor guidelines
• 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:

• 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 contains specific people who can review PRs to various areas.

Additional Information

• Related to # (issue)

[blisc]

Can you do a find and replace inside of this file of all print statements and switch them to logging?

[rlangman]

Should now say the map contains tuples

[subhankar-ghosh]

done

[rlangman]

The last condition should probably throw an error, as it implies someone tried to load an embedding but did it incorrectly.

[subhankar-ghosh]

has_baked_context_embedding property is used as a check in many places. In cases where the condition should absolutely be true it does throw an error. Hence I think this method just returns if baked_context_embedding is present and properly loaded.

[rlangman]

Why do we want to enforce that baked embedding can only be used with decoder_ce? Its not a big deal since I imagine we will be removing most of the other model_type code paths later. Just as first glance, it looks like the implementation would be easier/simpler if it was not gated by model type.

[rlangman]

A benefit of using an embedding layer is you could validate the shape of the embedding during model initialization. The value of N, num_baked_speakers, would need to be provided as a buffer from the checkpoint, replacing the current definition of it as an @property

[rlangman]

I would have this branch decided by whether speaker_indices is None or not. If speaker indices is provided, then use this path, else use the context encoder. An error should be thrown if someone provides neither or both context_audio and speaker_indices.

[subhankar-ghosh]

The value of N, num_baked_speakers, would need to be provided as a buffer from the checkpoint

Why do we need to have number of baked speakers saved as a buffer? We need to save T and D sure, but I am not sure why we need to save N as a buffer. We can get the num embeddings by embedding_layer.weight.shape[0]
Nonetheless, we are changing the implementation to torch.nn.embedding

[subhankar-ghosh]

I would have this branch decided by whether speaker_indices is None or not. If speaker indices is provided, then use this path, else use the context encoder. An error should be thrown if someone provides neither or both context_audio and speaker_indices.

The idea is if baked_embeddings are there the user should only use that and not context encoder. If speaker_indices are not provided and baked_embeddings are present, the default behavior is to use the first speaker embedding. I thought this is what we had decided - baked_embeddings so we do not let users use any other voice.

[rlangman]

The value of N, num_baked_speakers, would need to be provided as a buffer from the checkpoint

Why do we need to have number of baked speakers saved as a buffer? We need to save T and D sure, but I am not sure why we need to save N as a buffer. We can get the num embeddings by embedding_layer.weight.shape[0]

The request from Jason was to make the parameters a nn.Embedding layer. Creating this layer requires knowing the value of N ahead of time. Like:

num_baked_embeddings = self.register_buffer('num_baked_embeddings ', 0)
self.embedding_layer = torch.nn.Embedding(num_embeddings=num_baked_embeddings, embedding_dim=max_context_len * context_dim, _freeze=True)

This is a clean approach that guarantees the embedding is the correct dimensions when the model is first loaded. Otherwise if the embedding_layer happens to be wrong in the model checkpoint, then you would only know by hitting errors at inference time. Granted this is a pretty minor concern.

[blisc]

I would have this branch decided by whether speaker_indices is None or not. If speaker indices is provided, then use this path, else use the context encoder. An error should be thrown if someone provides neither or both context_audio and speaker_indices.

The idea is if baked_embeddings are there the user should only use that and not context encoder. If speaker_indices are not provided and baked_embeddings are present, the default behavior is to use the first speaker embedding. I thought this is what we had decided - baked_embeddings so we do not let users use any other voice.

I think we assume, and we should add this as a check to ensure that our assumption is correct, that a model can either have backed context embeddings OR a context encoder. Not both. So the path where speaker_indices is None is ok with me

[rlangman]

Unrelated to this release, there should be some inference path that allows a user to pre-compute and reuse context embeddings. So that when hosting the model, it is not necessary to recompute the same context for every request. It would be nice if that use case and the baked context embeddings could be consolidated, but we can revisit that in the later PR which defines the actual inference functions/APIs.

[github-actions]

[🤖]: Hi @subhankar-ghosh 👋,

We wanted to let you know that a CICD pipeline for this PR just finished successfully.

So it might be time to merge this PR or get some approvals.

//cc @chtruong814 @ko3n1g @pablo-garay @thomasdhc