Requesting Nitro-E

[jdluzen]

Great to have you back Adam!
I'd like to request Nitro-E support. I see some pieces are already supported, like FlowMatchEulerDiscreteScheduler, and things related to MMDiT. Is that reusable from SD3? I just started looking at the Nitro-E repo and comparing it to TensorStack. So much bigger since the last one.

Models:
meta-llama/Llama-3.2-1B, ONNX: https://huggingface.co/onnx-community/Llama-3.2-1B
mit-han-lab/dc-ae-f32c32-sana-1.0-diffusers, ONNX: available? I think, maybe. After much struggling, I managed to get https://github.com/mit-han-lab/efficientvit/blob/master/assets/onnx_export.py working. But I'm not sure what it exported, it's only 243MB, when the original is 1.25GB. I did not intentionally quantize it.

[saddam213]

Hey Man :)

Looks like there is some onnx versions of Llama-3.2-1B in that link, however we will need a TextGeneration pipeline as this is a decoder model not a basic TextEncoder

I have created a repo with the onnx model and a tokenizer format that is compatible with TensorStack
TensorStack/Llama-3.2-1B

A Llama pipeline should be fairly trivial as TensorStack already supports decoder only models

Converting the diffusion transformer model and vae encoder is an unknown, without that adding support is not likely, I have no clue how to convert models like this

[saddam213]

Llama-3.2-1B support added: 7559adf

Hopefully someone clever can convert the rest to onnx

[saddam213]

I cant run the python example as some of the models are gated

Cannot access gated repo for url https://huggingface.co/meta-llama/Llama-3.2-1B/resolve/main/config.json.
Access to model meta-llama/Llama-3.2-1B is restricted. You must have access to it and be authenticated to access it. Please log in.

Not keen to give meta my data for this

[jdluzen]

I cant run the python example as some of the models are gated

Cannot access gated repo for url https://huggingface.co/meta-llama/Llama-3.2-1B/resolve/main/config.json.
Access to model meta-llama/Llama-3.2-1B is restricted. You must have access to it and be authenticated to access it. Please log in.

Not keen to give meta my data for this

I know the feeling. unsloth/Llama-3.2-1B works out of the box for me.

[jdluzen]

Llama-3.2-1B support added: 7559adf

Hopefully someone clever can convert the rest to onnx

Maybe this will work:
VAE: https://huggingface.co/jdluzen/dc-ae-f32c32-sana-1.0-diffusers-onnx though it might be decoder only, the file size does not match up. I'm guessing I did that wrong. If img2img is to be supported, then we'll need the encoder? Edit: yep it's busted. Can't even load it in an InferenceSession.
Transformer is in progress.

[saddam213]

Had to make a big change to get ready for this pipeline

TensorStack.StableDiffusion now has a dependency on TensorStack.TextGeneration this will allow us to use decoder based models for diffusion pipelines.

I cant run the demo as it needs CUDA??? I only have an AMD card, I have no idea what im doing, LMAO

[jdluzen]

I also only have AMD. Though I am using WSL with a custom Dockerfile and ROCm, Nitro is so fast that you can change your device from cuda to cpu and barely notice.

[saddam213]

ok, think I have encoder and decoder converted, so just transformer to go

[jdluzen]

Wow great! What's your torch.onnx.export look like? Mine's broken somehow. I'm currently stuck on the tensor shapes for the transformer. I'm currently using (torch.randn(768, 32, 1, 1, dtype=dtype), torch.randn(1, 128, 2048, dtype=dtype), torch.tensor([999], dtype=torch.long)) which isn't right.

[saddam213]

im just editing one of the amuse converters
https://github.com/TensorStack-AI/AmuseConverters/tree/main/StableDiffusion3

will upload a Nitro one soon, well what I have so far

[saddam213]

latent_model_input = [1, 32, 16, 16]
encoder_hidden_states = [1, 128, 2048]
encoder_attention_mask = [1, 128]
timestep = [1]

[jdluzen]

Those certainly looks a lot better than mine 😅

[saddam213]

ok, think I have it sorted, have converted the 512px model

Now have to make a pipeline to put it all together and see if it work

I am stuck in bed injured ATM, but hopefully I can get this done over the next few days

[saddam213]

I think I have converted all models TensorStack/Nitro-E (this is the 1024 distilled version)

but wont know if they work until I get a pipeline working, started a PR here: AMD Nitro Pipeline