Add EasyMagpie vLLM-Omni serving

.github/actions/test-template/action.yml

@@ -56,6 +56,14 @@ inputs:
     description: "Directory under tests/ containing the test scripts"
     required: false
     default: "functional_tests"
+  download-artifact-name:
+    description: "Optional workflow artifact to place in the test checkout"
+    required: false
+    default: ""
+  download-artifact-path:
+    description: "Destination path relative to the test checkout"
+    required: false
+    default: ".ci-artifacts"
 runs:
   using: "composite"
   steps:
@@ -92,6 +100,14 @@ runs:
       with:
         path: ${{ github.run_id }}/${{steps.uuid.outputs.id }}/NeMo
 
+    - name: Download input artifact
+      if: ${{ inputs.download-artifact-name != '' }}
+      uses: actions/download-artifact@v7
+      with:
+        name: ${{ inputs.download-artifact-name }}
+        path: >-
+          ${{ github.run_id }}/${{ steps.uuid.outputs.id }}/NeMo/${{ inputs.download-artifact-path }}
+
     - name: Start container
       shell: bash
       env:

.github/workflows/cicd-main-speech.yml

@@ -41,6 +41,92 @@ jobs:
       dockerfile: docker/Dockerfile
       runner: ${{ inputs.runner }}
 
+  easymagpie-vllm-serving-changes:
+    runs-on: ubuntu-latest
+    outputs:
+      run: ${{ steps.decision.outputs.run }}
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+        with:
+          fetch-depth: 0
+      - name: Get PR info
+        id: get-pr-info
+        if: startsWith(github.ref, 'refs/heads/pull-request/')
+        uses: nv-gha-runners/get-pr-info@main
+      - name: Get changed files
+        id: changed-files
+        uses: step-security/changed-files@v45.0.1
+        with:
+          sha: ${{ github.sha }}
+          base_sha: ${{ fromJSON(steps.get-pr-info.outputs.pr-info || '{}').base.sha || github.event.pull_request.base.sha }}
+          files: |
+            .github/actions/test-template/action.yml
+            .github/workflows/cicd-main-speech.yml
+            tools/easymagpie_vllm_omni/**
+            tests/collections/tts/easymagpie_vllm_omni/conversion/**
+            tests/collections/tts/easymagpie_vllm_omni/serving/**
+            tests/functional_tests/L0_Unit_Tests_GPU_TTS_EasyMagpie_vLLM.sh
+      - name: Decide whether to run serving tests
+        id: decision
+        env:
+          ANY_CHANGED: ${{ steps.changed-files.outputs.any_changed }}
+          EVENT_NAME: ${{ github.event_name }}
+          TESTS_TO_RUN: ${{ inputs.test_to_run }}
+        run: |
+          run="$ANY_CHANGED"
+
+          if [[ "$EVENT_NAME" == "schedule" ]]; then
+            run=false
+          elif [[ "$EVENT_NAME" == "workflow_dispatch" ]] && \
+              jq -e 'index("all") != null or index("L0_Unit_Tests_GPU_TTS_EasyMagpie_vLLM") != null' \
+                <<<"$TESTS_TO_RUN" >/dev/null; then
+            run=true
+          fi
+
+          echo "run=$run" | tee -a "$GITHUB_OUTPUT"
+
+  easymagpie-vllm-serving-build:
+    needs: [easymagpie-vllm-serving-changes]
+    if: needs.easymagpie-vllm-serving-changes.outputs.run == 'true'
+    uses: ./.github/workflows/_build_container.yml
+    with:
+      image-name: nemo_container_easymagpie_vllm
+      dockerfile: tools/easymagpie_vllm_omni/Dockerfile.ci
+      runner: ${{ inputs.runner }}
+
+  easymagpie-codec-contract:
+    needs: [build, easymagpie-vllm-serving-changes]
+    if: needs.easymagpie-vllm-serving-changes.outputs.run == 'true'
+    runs-on: ${{ inputs.runner }}
+    name: EasyMagpie codec parity contract
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+      - name: Pull Speech image
+        run: docker pull ${{ needs.build.outputs.image }}
+      - name: Generate contract with Speech decoder
+        run: |
+          mkdir -p .ci-artifacts/easymagpie-codec-contract
+          docker run \
+            --rm \
+            --shm-size=8g \
+            --env NEMO_HOME=/tmp/nemo_home \
+            --volume "$(pwd):/workspace" \
+            --workdir /workspace \
+            ${{ needs.build.outputs.image }} \
+            bash -c "python tests/collections/tts/easymagpie_vllm_omni/conversion/generate_codec_contract.py \
+              .ci-artifacts/easymagpie-codec-contract"
+          du -h .ci-artifacts/easymagpie-codec-contract/*
+      - name: Upload codec contract
+        uses: actions/upload-artifact@v6
+        with:
+          name: easymagpie-codec-contract
+          path: .ci-artifacts/easymagpie-codec-contract
+          if-no-files-found: error
+          retention-days: 1
+          compression-level: 0
+
   unit-tests:
     strategy:
       fail-fast: false
@@ -146,6 +232,30 @@ jobs:
           cpu-only: ${{ matrix.cpu-only || false }}
           is_optional: ${{ matrix.is-optional || false }}
 
+  easymagpie-vllm-serving-tests:
+    needs: [easymagpie-vllm-serving-build, easymagpie-codec-contract]
+    runs-on: ${{ inputs.runner }}
+    name: L0_Unit_Tests_GPU_TTS_EasyMagpie_vLLM
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v6
+      - name: Checkout NeMo
+        uses: actions/checkout@v6
+        with:
+          path: ${{ github.run_id }}
+      - name: main
+        uses: ./.github/actions/test-template
+        with:
+          runner: ${{ runner.name }}
+          script: L0_Unit_Tests_GPU_TTS_EasyMagpie_vLLM
+          is_unit_test: true
+          tests_to_run: ${{ inputs.test_to_run }}
+          image: ${{ needs.easymagpie-vllm-serving-build.outputs.image }}
+          test-data-path: ${{ vars.DEFAULT_TEST_DATA_PATH || '/mnt/datadrive/TestData' }}
+          timeout: 15
+          download-artifact-name: easymagpie-codec-contract
+          download-artifact-path: .ci_artifacts/easymagpie-codec-contract
+
   e2e-tests:
     strategy:
       fail-fast: false

tests/collections/tts/easymagpie_vllm_omni/conversion/conftest.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Path setup for Speech-side EasyMagpie checkpoint conversion tests."""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+EASYMAGPIE_ROOT = Path(__file__).resolve().parents[5] / "tools" / "easymagpie_vllm_omni"
+SCRIPTS_DIR = EASYMAGPIE_ROOT / "scripts"
+
+# Conversion uses Speech and pure PyTorch helpers, never the vLLM runtime.
+sys.path.insert(0, str(EASYMAGPIE_ROOT))
+sys.path.insert(0, str(SCRIPTS_DIR))

tests/collections/tts/easymagpie_vllm_omni/conversion/generate_codec_contract.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python3
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Generate a tiny Speech-to-vLLM codec parity contract for the serving CI job."""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+import torch
+from hydra.utils import instantiate
+from safetensors.torch import save_file
+
+ROOT = Path(__file__).resolve().parents[5]
+EASYMAGPIE_ROOT = ROOT / "tools" / "easymagpie_vllm_omni"
+sys.path.insert(0, str(EASYMAGPIE_ROOT))
+sys.path.insert(0, str(EASYMAGPIE_ROOT / "scripts"))
+
+import convert_codec as converter  # noqa: E402
+from easymagpie_vllm_omni.codec.config import EasyMagpieCodecConfig  # noqa: E402
+from easymagpie_vllm_omni.codec.weight_conversion import convert_decoder_state_dict  # noqa: E402
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("output", type=Path)
+    return parser.parse_args()
+
+
+def speech_decoder_config() -> dict:
+    return {
+        "_target_": "nemo.collections.tts.modules.audio_codec_modules.ResNetDecoder",
+        "input_dim": 4,
+        "input_filters": 8,
+        "pre_up_sample_rates": [2],
+        "pre_up_sample_filters": [8],
+        "n_hidden_layers": 2,
+        "hidden_filters": 16,
+        "resblock_up_sample_rates": [2],
+        "resblock_up_sample_filters": [4],
+        "resblock_up_sample_kernel_size": 7,
+        "kernel_size": 3,
+        "activation": "half_snake",
+        "is_causal": True,
+        "pad_mode": "replicate",
+    }
+
+
+def serving_config() -> EasyMagpieCodecConfig:
+    return EasyMagpieCodecConfig(
+        input_dim=4,
+        input_filters=8,
+        hidden_filters=16,
+        num_hidden_layers=2,
+        pre_upsample_rates=[2],
+        pre_upsample_filters=[8],
+        resblock_upsample_rates=[2],
+        resblock_upsample_filters=[4],
+        kernel_size=3,
+        resblock_kernel_size=7,
+        activation="half_snake",
+        num_codebooks=2,
+        codebook_size=4,
+        num_levels_per_group=[2, 2],
+        frame_stacking_factor=2,
+        output_sample_rate=22050,
+    )
+
+
+def generate_contract(output: Path) -> None:
+    torch.manual_seed(20260730)
+    decoder_config = speech_decoder_config()
+    source_decoder = instantiate(decoder_config)
+    source_state = {
+        f"audio_decoder.{name}": tensor.detach().cpu() for name, tensor in source_decoder.state_dict().items()
+    }
+    decoder = converter.restore_speech_decoder(decoder_config, source_state)
+
+    num_frames = 7
+    latent = torch.linspace(-1.0, 1.0, steps=num_frames * decoder_config["input_dim"], dtype=torch.float32)
+    latent = latent.view(num_frames, decoder_config["input_dim"])
+    input_len = torch.tensor([num_frames], dtype=torch.long)
+    with torch.no_grad():
+        audio, audio_len = decoder(inputs=latent.T.unsqueeze(0).contiguous(), input_len=input_len)
+
+    config = serving_config()
+    expected_audio_len = num_frames * config.samples_per_codec_frame
+    if int(audio_len.item()) != expected_audio_len:
+        raise AssertionError(f"Speech decoder returned {int(audio_len.item())} samples, expected {expected_audio_len}")
+
+    output.mkdir(parents=True, exist_ok=True)
+    config.save_pretrained(output)
+    save_file(convert_decoder_state_dict(source_state), output / "model.safetensors")
+    save_file(
+        {
+            "latent": latent.contiguous(),
+            "expected_audio": audio.squeeze(0).contiguous(),
+            "expected_audio_len": audio_len.cpu().contiguous(),
+        },
+        output / "contract.safetensors",
+    )
+
+
+if __name__ == "__main__":
+    generate_contract(parse_args().output)

tests/collections/tts/easymagpie_vllm_omni/conversion/test_convert_codec.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import convert_codec as converter
+import pytest
+
+
+def _valid_decoder_config() -> dict:
+    return {
+        "_target_": "nemo.collections.tts.modules.audio_codec_modules.ResNetDecoder",
+        "is_causal": True,
+        "activation": "half_snake",
+    }
+
+
+@pytest.mark.parametrize(
+    ("field", "value", "message"),
+    [
+        ("_target_", "some.OtherDecoder", "ResNetDecoder"),
+        ("is_causal", False, "causal"),
+        ("activation", "snake", "half_snake"),
+    ],
+)
+def test_validate_decoder_config_rejects_unsupported_codec(field, value, message):
+    config = _valid_decoder_config()
+    config[field] = value
+
+    with pytest.raises(ValueError, match=message):
+        converter.validate_decoder_config(config)