Lyra v2 TF/JAX precursor model files? (trying to convert models to ONNX format)

[josephrocca]

I'm trying to get Lyra v2 working in the browser via tfjs-tflite, and also ORT Web. On the ONNX side of things, I'm blocked by a conversion error and the tf2onnx team is unsure whether the tflite files are completely "valid", so they've asked that I look into whether there were some bugs in the conversion process, particularly with respect to the "NoOp" sub-graph of the encoder model:

onnx/tensorflow-onnx#2055 (comment)

That thread shows netron.app images of what I'm referring to. There are TFL_ASSIGN_VARIABLE nodes which don't have any outputs, which is apparently not expected/valid.

Since the TF/JAX/(?) precursor to the tflite model isn't available in this repo, I'm not able to investigate further, so I was wondering if those original models could be published, or whether anyone on the Lyra team has any thoughts/expertise here?

If the original model files can't be published, I wonder if it would be possible for the team to convert the original models to ONNX, and then publish those ONNX files in this repo? Maybe there's something about the TF/JAX-->TFLite-->ONNX pathway that is causing the problem, and that wouldn't happen if we just went from TF/JAX straight to ONNX.

The script for the conversion, assuming that the original model is a SavedModel, would be:

# 1. install tf2onnx:
pip install git+https://github.com/onnx/tensorflow-onnx
# 2. convert (assumes SavedModel folder is called 'soundstream_encoder' in current directory):
python -m tf2onnx.convert --opset 16 --saved-model soundstream_encoder --output soundstream_encoder.onnx --verbose --debug
# 3. repeat the above command for the other two models

I've tested with TensorFlow version 2.9.2. Note that versions 2.8 and 2.10 didn't work for me due to this error and another error which I can't remember.

[aluebs]

Thank you for the detailed report. There is no plan to open-source the original model files at the moment. And I don't think anyone on the team has ONNX experience. Let me see if I can figure out what's happening.

[aluebs]

Your instructions to convert the saved model to ONNX seem to work successfully. Unfortunately I don't have the ONNX experience, or bandwidth at the moment, to evaluate the correctness of these ONNX files.

[josephrocca]

@aluebs Thanks for looking into this! Are you able to share the ONNX files so I can have a play with them?

[aluebs]

Releasing them would require some additional internal review, so unfortunately I don't think that it would be possible in the short term. I just thought I'd let you know that your hypothesis of the TF -->TFLite-->ONNX pathway having a problem seems to be correct. Hope that you figure it out.

[josephrocca]

Okay, thanks for your efforts so far. I was hoping that I wouldn't need the ONNX files, but the tfjs-tflite runtime is unfortunately pretty buggy/incomplete:

• [tfjs-tflite] Loading second TFLite model hangs in _emscripten_futex_wait tensorflow/tfjs#6094 (comment)
• [tfjs-tflite] How to use a particular a subgraph of a tflite model? tensorflow/tfjs#6919

so I'll try to convince the tf2onnx people to look further into the tflite-->ONNX conversion issues (hopefully helped by the fact that you've reported that TF -->ONNX worked without issues), but if that stalls then I'll probably have exhausted all the obvious pathways toward getting this working on the web for now, and so in that case I'll probably be awaiting a possible set of ONNX files from your team, assuming it passes internal review at some point 🙏, although I understand that this may not be possible in the near future.

Thanks again! It's a really exciting bit of open source software, and I look forward to getting it working on the web at some point.

---

P.S. in case this stalls: For those hitting this thread in the future and wanting to continue where I left off, here's my progress:

• https://github.com/josephrocca/lyra-v2-soundstream-web
• https://josephrocca.github.io/lyra-v2-soundstream-web

---

Edit: Looks like there are at least 2 problems to overcome to with the ONNX conversion from tflite:

• IndexError: list index out of range in parse_tflite_graph onnx/tensorflow-onnx#2055 (comment)
• [tflite] TFL_VAR_HANDLE and TFL_READ_VARIABLE are not supported onnx/tensorflow-onnx#2059

[aluebs]

Thank you so much for sharing your progress and documenting the blocking issues. I'll definitely keep you posted if there is any progress. I'd love to see Lyra working in the web :)

[josephrocca]

Hey @aluebs, I've been chipping away at getting this working in the browser (see repo), and am wondering whether you could comment on this - i.e. whether the VAR_OPS are required during inference?

Assuming that the above-linked comment is correct: I actually didn't know that TFLite models could be "stateful" like this - I'm not sure if that's possible with ONNX so if not then I'm guessing I'd have to "manually" reach into the model and manage the variable reads/writes between successive inferences? I.e. basically capture all AssignVariable values as "outputs" and feed them back into the places that try to read from that variable during the next inference?

[aluebs]

Yes, these models are stateful indeed. There is possibility to export stateless models that return a state pointer after every call that needs to be copied over to the input for of the next call. That forces the user to handle the state manually, but that is what you need for this use-case. I am no longer at Google, so I don't have access to the exporting pipeline, but maybe there is a way to convert from one model to the other?

[josephrocca]

Thanks for the confirmation! I'll continue to investigate this.

(Apologies for pinging you on this given that you're no longer at Google - I just submitted some feedback to Github that's related to this.)

[aluebs]

Oh, no worries, I recently left and am still super passionate about Lyra. Happy to help with what I can, just mentioned in case you want to direct your questions to the maintainers :)