Coqui-TTS训练的VITS模型导出ONNX时输入长度固定问题及动态音频生成方法咨询

I see your issue—you've successfully exported your VITS model to ONNX, but it's stuck generating fixed-length audio, with noise cropping up when using shorter inputs. Let's break down why this happens and how to fix it to get dynamic, input-length-matched audio output.

Root Cause #

The problem comes down to how the ONNX export captures your model's dynamic computation graph. Even though you set dynamic_axes, the fixed-length dummy input (100 tokens) can make ONNX bake in static dimension assumptions—especially in modules like the duration predictor and decoder, which rely directly on input length to generate variable-length audio.

Step-by-Step Solution #

1. Update the ONNX Export Code

Modify the export_onnx method to ensure all dynamic dimensions are properly declared, and the dummy input doesn't enforce a rigid fixed length. Here's the revised code with key improvements:

def export_onnx(self, output_path: str = "coqui_vits.onnx", verbose: bool = True):
    """Export model to ONNX format for inference with dynamic input support"""
    # Save original state to roll back later
    _forward = self.forward
    disc = self.disc if hasattr(self, "disc") else None
    training = self.training

    # Prepare model for export
    if disc:
        self.disc = None
    self.eval()

    def onnx_inference(text, text_lengths, scales, sid=None, langid=None):
        noise_scale, length_scale, noise_scale_dp = scales
        self.noise_scale = noise_scale
        self.length_scale = length_scale
        self.noise_scale_dp = noise_scale_dp
        return self.inference(
            text,
            aux_input={
                "x_lengths": text_lengths,
                "d_vectors": None,
                "speaker_ids": sid,
                "language_ids": langid,
                "durations": None,
            },
        )["model_outputs"]

    self.forward = onnx_inference

    # Use a small, flexible dummy input instead of fixed 100 tokens
    dummy_input_length = 10  # Shorter length helps ONNX track dynamic behavior better
    sequences = torch.randint(low=0, high=2, size=(1, dummy_input_length), dtype=torch.long)
    sequence_lengths = torch.LongTensor([sequences.size(1)])
    scales = torch.FloatTensor([self.inference_noise_scale, self.length_scale, self.inference_noise_scale_dp])
    dummy_input = (sequences, sequence_lengths, scales)
    input_names = ["input", "input_lengths", "scales"]
    
    # Define all dynamic axes explicitly
    dynamic_axes = {
        "input": {0: "batch_size", 1: "phoneme_length"},
        "input_lengths": {0: "batch_size"},
        "scales": {0: "scale_count"},
        "output": {0: "batch_size", 1: "audio_length", 2: "channels"},
    }

    if self.num_speakers > 0:
        speaker_id = torch.LongTensor([0])
        dummy_input += (speaker_id,)
        input_names.append("sid")
        dynamic_axes["sid"] = {0: "batch_size"}  # Add dynamic batch support for speaker IDs

    if hasattr(self, "num_languages") and self.num_languages > 0 and self.embedded_language_dim > 0:
        language_id = torch.LongTensor([0])
        dummy_input += (language_id,)
        input_names.append("langid")
        dynamic_axes["langid"] = {0: "batch_size"}  # Add dynamic batch support for language IDs

    # Export with settings that preserve dynamic behavior
    torch.onnx.export(
        model=self,
        args=dummy_input,
        opset_version=18,
        f=output_path,
        verbose=verbose,
        input_names=input_names,
        output_names=["output"],
        dynamic_axes=dynamic_axes,
        dynamo=False,
        do_constant_folding=False,  # Prevent dynamic operations from being locked to constants
        export_params=True,
        keep_initializers_as_inputs=False,
    )

    # Restore original model state
    self.forward = _forward
    if training:
        self.train()
    if disc:
        self.disc = disc

Key Changes Explained:

• Smaller Dummy Input: Using dummy_input_length=10 avoids forcing ONNX to assume a fixed long input length, which helps it track dynamic paths in the model.
• Expanded Dynamic Axes: Explicitly added dynamic batch dimensions for speaker and language IDs (if applicable) and clarified dimension names for readability.
• Disabled Constant Folding: do_constant_folding=False ensures dynamic operations like duration prediction and audio scaling aren't converted to static constants during export.

2. Verify Dynamic Input Support

After re-exporting, confirm the model accepts variable-length inputs with this quick check:

import onnxruntime as ort

session = ort.InferenceSession("coqui_vits.onnx")
for input in session.get_inputs():
    print(f"Input: {input.name}, Shape: {input.shape}")
# You should see dynamic dimensions like `['batch_size', 'phoneme_length']` instead of fixed numbers

3. Test with Your Serbian Text

Your existing test code is already set up correctly—just make sure:

• The text_to_ids function properly handles your Serbian token set (including the <BLNK> token).
• You're passing input_lengths that exactly match the length of your tokenized text (your code already does this with len(ids)).

Why This Works #

By explicitly declaring all dynamic dimensions and avoiding fixed-length dummy inputs, ONNX will properly track how the model adapts to input length. The duration predictor will generate the correct number of frames for your text, and the decoder will produce audio of the exact required length—no extra padding, no noise.

内容的提问来源于stack exchange，提问作者Darko Milošević