Add StyleTTS 2

[eustlb]

What does this PR do?

Adds StyleTTS 2 to support the original model but also other checkpoints like Kokoro.

🆕 🔥 This implementation also add batch support (early benchmarks show ~50% inference speed improvement for BS 128) and mask support for padded inputs.

Note

This implementation differs slightly from the original codebase. The aim here is to have the clearest possible correspondence of naming and structure compared to the original papers (StyleTTS 2 that builds on top of StyleTTS).

TODO

• modeling code → ready for review
• tokenization → add other languages + docstring
• processor → add docstring
• tests
• model_doc

Benchmarks

BS 1

• ~8% faster than the original codebase (benchmarking script: original and transformers)

[eustlb]

cc @Cyrilvallez, modeling code is ready for review 🤗

[Cyrilvallez]

Alright! Super sorry about the (very) long delay! Very very nice implementation, not much to complain about here! Great work! 🤗
Mostly just a few recurrent but easy to fix stuff:

• Use only the config in __init__ if possible. I did not check each one, and some may have additional params that do not really make sense in the config as they are hard-coded, such as the boolean args in StyleTextToSpeech2AdainResBlock1d, but let's still try to use config as much as possible (and add additional parameters if still needed for internal-only purposes)
• If possible, all the transpose ops should have explicit dim number, let's try not to use -1 if number of dims is always the same (not sure here)
• The weight_norm function should probably always be one or the other, as layer names depend on it (if needed, let's use a hard pytorch version check for this one)

[Cyrilvallez]

We can only use this one anyway due to the layer naming in the state dicts no? Did you add the check for older torch versions? Also, IMO it would help clarity to wrap the layer declaration directly with the function, as opposed to applying it afterwards

[eustlb]

True! We do not need to check for older torch versions right? since parametrizations.weight_norm is supported anyway for torch>=2.1

[Cyrilvallez]

Humm, do we need this call every forward?

[Cyrilvallez]

Here it's actually different from the default, but with a quick glance at the library i'm not sure it's useful to set it, but not sure

[Cyrilvallez]

Also, I assume modular could not be used at all here?

[monuminu]

any update on this ?

[eustlb]

Hey @monuminu, this PR went stale due to other priorities. I’ll try to get it moving ASAP, though I’m not sure I’ll manage before the summer break. Thanks for checking in, I’ll bump its priority.