Add a onnx version of model

[NeuroDonu]

Self Checks

• I have thoroughly reviewed the project documentation (installation, training, inference) but couldn't find any relevant information that meets my needs. English 中文 日本語 Portuguese (Brazil)
• I have searched for existing issues search for existing issues, including closed ones.
• I confirm that I am using English to submit this report (我已阅读并同意 Language Policy).
• [FOR CHINESE USERS] 请务必使用英文提交 Issue，否则会被关闭。谢谢！:）
• Please do not modify this template :) and fill in all the required fields.

1. Is this request related to a challenge you're experiencing? Tell us your story.

I have spent the last 10 hours trying to convert your pytorch model to onnx. Because it seems to me that if you convert it to onnx version and use tensorrt for inference, the speed will increase many times! But all my attempts to convert the model were unsuccessful.... I tried to write code with deepseek and qwen, but they couldn't produce anything sensible. So I came here to ask for a variation of onnx model =)

2. What is your suggested solution?

Convert pytorch models vqgan and IF THIS IS POSSIBLE llama. deepseek suggested some code like this to convert vqgan, but my tweaks only broke it.

3. Additional context or comments

No response

4. Can you help us with this feature?

• I am interested in contributing to this feature.

[Grime715]

due to the blocking of Ukraine from 11labs, our team also now needs this improvement, please pay your attention! Sincerely)

[tarun7r]

Hi @NeuroDonu , any luck on converting the models ( vocoder and the semantic generator model ) to the onnx and later convert them to tensorrt for faster inference

[NeuroDonu]

Hi @NeuroDonu , any luck on converting the models ( vocoder and the semantic generator model ) to the onnx and later convert them to tensorrt for faster inference

Alas, no. I've never been able

[tarun7r]

@NeuroDonu apart form the --compile , did you find any way to speed up the inference of the model? if so can you please share it

[NeuroDonu]

@NeuroDonu Помимо --compile, нашли ли вы какой-либо способ ускорить вывод модели? Если да, то не могли бы вы поделиться им?

No, didnt find

[YangGuangzhaoJJJacky]

I converted but facing problem like this when running it

onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL :
Non-zero status code returned while running Where node. Name:'/Where'
Status Message: /Where: condition operand cannot broadcast on dim 1
Condition Shape: {1,44}
X Shape: {}
Y Shape: {1,44,1024}

[YangGuangzhaoJJJacky]

here is the way

def decode_one_token_ar(
    model: DualARTransformer,
    x: torch.Tensor,
    input_pos: torch.Tensor,
    semantic_ids: list,
    previous_tokens: torch.Tensor = None,
    **sampling_kwargs,
) -> torch.Tensor:
    x_out = model.forward_generate(x, input_pos)

    import torch.nn as nn
    device = next(model.parameters()).device
    class ONNXWrapper(nn.Module):
        def __init__(self, model: DualARTransformer):
            super().__init__()
            self.model = model.eval()

        def forward(self, x, input_pos):
        
            out = self.model.forward_generate(x, input_pos)
            return out.logits
    onnx_model = ONNXWrapper(model).to(device).eval()
    dummy_x = x
    dummy_input_pos = input_pos
    print("dummy_x shape:", dummy_x.shape)
    print("dummy_input_pos shape:", dummy_input_pos.shape)

    torch.onnx.export(
    onnx_model,  
    (dummy_x, dummy_input_pos),  
    "onnx/dual_ar_transformer_generate.onnx", 
    input_names=["x", "input_pos"],
    output_names=["logits"],
    dynamic_axes = {
        "x": {0: "batch_size", 2: "seq_len"},
        "input_pos": {0: "seq_len"}, 
        "logits": {0: "batch_size", 1: "seq_len"},
    },
    opset_version=17,  
    )

    print("✅ ONNX！")
    exit()

also I found way to streaming the token2wav part(not perfect)

[NeuroDonu]

Wow! So this would be the way to convert all the parts?

[tarun7r]

@YangGuangzhaoJJJacky could you share the end to end text to speech pipeline using the onnx model extracted from the above code please

[PoTaTo-Mika]

here is the way

def decode_one_token_ar(
    model: DualARTransformer,
    x: torch.Tensor,
    input_pos: torch.Tensor,
    semantic_ids: list,
    previous_tokens: torch.Tensor = None,
    **sampling_kwargs,
) -> torch.Tensor:
    x_out = model.forward_generate(x, input_pos)

    import torch.nn as nn
    device = next(model.parameters()).device
    class ONNXWrapper(nn.Module):
        def __init__(self, model: DualARTransformer):
            super().__init__()
            self.model = model.eval()

        def forward(self, x, input_pos):
        
            out = self.model.forward_generate(x, input_pos)
            return out.logits
    onnx_model = ONNXWrapper(model).to(device).eval()
    dummy_x = x
    dummy_input_pos = input_pos
    print("dummy_x shape:", dummy_x.shape)
    print("dummy_input_pos shape:", dummy_input_pos.shape)

    torch.onnx.export(
    onnx_model,  
    (dummy_x, dummy_input_pos),  
    "onnx/dual_ar_transformer_generate.onnx", 
    input_names=["x", "input_pos"],
    output_names=["logits"],
    dynamic_axes = {
        "x": {0: "batch_size", 2: "seq_len"},
        "input_pos": {0: "seq_len"}, 
        "logits": {0: "batch_size", 1: "seq_len"},
    },
    opset_version=17,  
    )

    print("✅ ONNX！")
    exit()

also I found way to streaming the token2wav part(not perfect)

Nice work, would you mind merge it into the tools/export_onnx.py ?

[YangGuangzhaoJJJacky]

I can merge but still facing problem while running

onnxruntime.capi.onnxruntime_pybind11_state.Fail: [ONNXRuntimeError] : 1 : FAIL :
Non-zero status code returned while running Where node. Name:'/Where'
Status Message: /Where: condition operand cannot broadcast on dim 1
Condition Shape: {1,44}
X Shape: {}
Y Shape: {1,44,1024}

[github-actions]

This issue is stale because it has been open for 30 days with no activity.