[JustWorks] Export and lower the TTS Kokoro model and identify user experience gaps

[tarun292]

🚀 The feature, motivation and pitch

This new model: https://huggingface.co/hexgrad/Kokoro-82M has gone viral recently and seems to be a very high quality TTS model. This is a good model to take as a use case and lower to ExecuTorch and run e2e.

The main aim is to not just enable this model but also during this process as part of the ExecuTorch JustWorks workstream we'll identify user experience issues/gaps and fix those.

Alternatives

No response

Additional context

No response

RFC (Optional)

No response

[jackzhxng]

Is kokoro transformer based?

[NorbertKlockiewicz]

I attempted to export the model, but unfortunately it's using LSTMs, which aren't compatible with dynamic sequence lengths when using TorchDynamo. Unless there's a workaround or some hack to get around this limitation, it seems that exporting the model may not be feasible at the moment — which is a shame, as Kokoro is an excellent model. If you happen to know any potential tricks or solutions, I'd be happy to give it another try.

[jackzhxng]

Yeah, that is a known limitation which the export team is looking into the Dynamo side and the ExecuTorch team is looking into on the op side. I actually just started looking into a static export, were you able to get that exported?

[NorbertKlockiewicz]

I tried exporting with static sizes and after some changes to model, e.g removing pack_padded_sequence and splitting model into few parts, the exports were working, however I didn't check if the outputs are correct as the model with static sizes didn't fit my needs.

[jackzhxng]

Makes sense. We will try to target this in the near future.

[jackzhxng]

Moving to #12344

[Jransom33]

I tried exporting with static sizes and after some changes to model, e.g removing pack_padded_sequence and splitting model into few parts, the exports were working, however I didn't check if the outputs are correct as the model with static sizes didn't fit my needs.

Any chance you would be willing to share that code? I am new to on-device inference and having trouble getting it to work

[geneing]

Exporting LSTM with static size in the case of kokoro is a bad idea afaik. What dynamo does is unroll the fixed length loop. Since LSTM is based on the sequence length, you'll get an enormous model with hundreds of layer replacing each LSTM. Dynamo needs to implement loops properly. It's frustrating that dynamo is now the main export path, but it's not ready to be used.