feat(backends): add Whisper-Medusa transcription

.github/backend-matrix.yml

@@ -166,6 +166,19 @@ include:
     dockerfile: "./backend/Dockerfile.python"
     context: "./"
     ubuntu-version: '2404'
+  - build-type: ''
+    cuda-major-version: ""
+    cuda-minor-version: ""
+    platforms: 'linux/amd64'
+    tag-latest: 'auto'
+    tag-suffix: '-cpu-whisper-medusa'
+    runs-on: 'ubuntu-latest'
+    base-image: "ubuntu:24.04"
+    skip-drivers: 'true'
+    backend: "whisper-medusa"
+    dockerfile: "./backend/Dockerfile.python"
+    context: "./"
+    ubuntu-version: '2404'
   - build-type: ''
     cuda-major-version: ""
     cuda-minor-version: ""
@@ -626,6 +639,19 @@ include:
     dockerfile: "./backend/Dockerfile.python"
     context: "./"
     ubuntu-version: '2404'
+  - build-type: 'cublas'
+    cuda-major-version: "12"
+    cuda-minor-version: "1"
+    platforms: 'linux/amd64'
+    tag-latest: 'auto'
+    tag-suffix: '-gpu-nvidia-cuda-12-whisper-medusa'
+    runs-on: 'ubuntu-latest'
+    base-image: "ubuntu:24.04"
+    skip-drivers: 'false'
+    backend: "whisper-medusa"
+    dockerfile: "./backend/Dockerfile.python"
+    context: "./"
+    ubuntu-version: '2404'
   - build-type: 'cublas'
     cuda-major-version: "12"
     cuda-minor-version: "8"

Makefile

@@ -1,5 +1,5 @@
 # Disable parallel execution for backend builds
-.NOTPARALLEL: backends/diffusers backends/llama-cpp backends/turboquant backends/bonsai backends/outetts backends/piper backends/stablediffusion-ggml backends/trellis2cpp backends/trellis2cpp-darwin backends/whisper backends/crispasr backends/parakeet-cpp backends/moss-transcribe-cpp backends/faster-whisper backends/silero-vad backends/local-store backends/valkey-store backends/cloud-proxy backends/huggingface backends/rfdetr backends/rfdetr-cpp backends/insightface backends/speaker-recognition backends/kitten-tts backends/kokoro backends/chatterbox backends/llama-cpp-darwin backends/neutts build-darwin-python-backend build-darwin-go-backend backends/mlx backends/diffuser-darwin backends/mlx-vlm backends/mlx-audio backends/mlx-distributed backends/stablediffusion-ggml-darwin backends/vllm backends/vllm-omni backends/longcat-video backends/sglang backends/moonshine backends/pocket-tts backends/qwen-tts backends/faster-qwen3-tts backends/qwen-asr backends/nemo backends/voxcpm backends/whisperx backends/ace-step backends/acestep-cpp backends/fish-speech backends/voxtral backends/opus backends/trl backends/llama-cpp-quantization backends/kokoros backends/sam3-cpp backends/qwen3-tts-cpp backends/moss-tts-cpp backends/magpie-tts-cpp backends/vllm-cpp backends/omnivoice-cpp backends/vibevoice-cpp backends/localvqe backends/tinygrad backends/sherpa-onnx backends/ds4 backends/ds4-darwin backends/liquid-audio backends/supertonic backends/depth-anything-cpp backends/privacy-filter backends/privacy-filter-darwin backends/audio-cpp backends/audio-cpp-darwin
+.NOTPARALLEL: backends/diffusers backends/llama-cpp backends/turboquant backends/bonsai backends/outetts backends/piper backends/stablediffusion-ggml backends/trellis2cpp backends/trellis2cpp-darwin backends/whisper backends/crispasr backends/parakeet-cpp backends/moss-transcribe-cpp backends/faster-whisper backends/whisper-medusa backends/silero-vad backends/local-store backends/valkey-store backends/cloud-proxy backends/huggingface backends/rfdetr backends/rfdetr-cpp backends/insightface backends/speaker-recognition backends/kitten-tts backends/kokoro backends/chatterbox backends/llama-cpp-darwin backends/neutts build-darwin-python-backend build-darwin-go-backend backends/mlx backends/diffuser-darwin backends/mlx-vlm backends/mlx-audio backends/mlx-distributed backends/stablediffusion-ggml-darwin backends/vllm backends/vllm-omni backends/longcat-video backends/sglang backends/moonshine backends/pocket-tts backends/qwen-tts backends/faster-qwen3-tts backends/qwen-asr backends/nemo backends/voxcpm backends/whisperx backends/ace-step backends/acestep-cpp backends/fish-speech backends/voxtral backends/opus backends/trl backends/llama-cpp-quantization backends/kokoros backends/sam3-cpp backends/qwen3-tts-cpp backends/moss-tts-cpp backends/magpie-tts-cpp backends/vllm-cpp backends/omnivoice-cpp backends/vibevoice-cpp backends/localvqe backends/tinygrad backends/sherpa-onnx backends/ds4 backends/ds4-darwin backends/liquid-audio backends/supertonic backends/depth-anything-cpp backends/privacy-filter backends/privacy-filter-darwin backends/audio-cpp backends/audio-cpp-darwin
 
 GOCMD=go
 GOTEST=$(GOCMD) test
@@ -611,6 +611,7 @@ prepare-test-extra: protogen-python
 	$(MAKE) -C backend/python/nemo
 	$(MAKE) -C backend/python/voxcpm
 	$(MAKE) -C backend/python/faster-whisper
+	$(MAKE) -C backend/python/whisper-medusa
 	$(MAKE) -C backend/python/whisperx
 	$(MAKE) -C backend/python/ace-step
 	$(MAKE) -C backend/python/trl
@@ -642,6 +643,7 @@ test-extra: prepare-test-extra
 	$(MAKE) -C backend/python/nemo test
 	$(MAKE) -C backend/python/voxcpm test
 	$(MAKE) -C backend/python/faster-whisper test
+	$(MAKE) -C backend/python/whisper-medusa test
 	$(MAKE) -C backend/python/whisperx test
 	$(MAKE) -C backend/python/ace-step test
 	$(MAKE) -C backend/python/trl test
@@ -1317,6 +1319,7 @@ BACKEND_RERANKERS = rerankers|python|.|false|true
 BACKEND_TRANSFORMERS = transformers|python|.|false|true
 BACKEND_OUTETTS = outetts|python|.|false|true
 BACKEND_FASTER_WHISPER = faster-whisper|python|.|false|true
+BACKEND_WHISPER_MEDUSA = whisper-medusa|python|.|false|true
 BACKEND_COQUI = coqui|python|.|false|true
 BACKEND_RFDETR = rfdetr|python|.|false|true
 BACKEND_INSIGHTFACE = insightface|python|.|false|true
@@ -1407,6 +1410,7 @@ $(eval $(call generate-docker-build-target,$(BACKEND_RERANKERS)))
 $(eval $(call generate-docker-build-target,$(BACKEND_TRANSFORMERS)))
 $(eval $(call generate-docker-build-target,$(BACKEND_OUTETTS)))
 $(eval $(call generate-docker-build-target,$(BACKEND_FASTER_WHISPER)))
+$(eval $(call generate-docker-build-target,$(BACKEND_WHISPER_MEDUSA)))
 $(eval $(call generate-docker-build-target,$(BACKEND_COQUI)))
 $(eval $(call generate-docker-build-target,$(BACKEND_RFDETR)))
 $(eval $(call generate-docker-build-target,$(BACKEND_INSIGHTFACE)))
@@ -1456,7 +1460,7 @@ $(eval $(call generate-docker-build-target,$(BACKEND_SUPERTONIC)))
 docker-save-%: backend-images
 	docker save local-ai-backend:$* -o backend-images/$*.tar
 
-docker-build-backends: docker-build-llama-cpp docker-build-ik-llama-cpp docker-build-turboquant docker-build-bonsai docker-build-ds4 docker-build-rerankers docker-build-vllm docker-build-vllm-omni docker-build-longcat-video docker-build-sglang docker-build-transformers docker-build-outetts docker-build-diffusers docker-build-kokoro docker-build-faster-whisper docker-build-crispasr docker-build-coqui docker-build-chatterbox docker-build-vibevoice docker-build-liquid-audio docker-build-moonshine docker-build-pocket-tts docker-build-qwen-tts docker-build-fish-speech docker-build-faster-qwen3-tts docker-build-qwen-asr docker-build-nemo docker-build-voxcpm docker-build-whisperx docker-build-ace-step docker-build-acestep-cpp docker-build-voxtral docker-build-mlx-distributed docker-build-trl docker-build-llama-cpp-quantization docker-build-tinygrad docker-build-kokoros docker-build-sam3-cpp docker-build-rfdetr-cpp docker-build-qwen3-tts-cpp docker-build-moss-tts-cpp docker-build-magpie-tts-cpp docker-build-vllm-cpp docker-build-omnivoice-cpp docker-build-vibevoice-cpp docker-build-localvqe docker-build-insightface docker-build-speaker-recognition docker-build-sherpa-onnx docker-build-cloud-proxy docker-build-supertonic docker-build-depth-anything-cpp docker-build-moss-transcribe-cpp docker-build-privacy-filter docker-build-trellis2cpp docker-build-valkey-store docker-build-audio-cpp
+docker-build-backends: docker-build-llama-cpp docker-build-ik-llama-cpp docker-build-turboquant docker-build-bonsai docker-build-ds4 docker-build-rerankers docker-build-vllm docker-build-vllm-omni docker-build-longcat-video docker-build-sglang docker-build-transformers docker-build-outetts docker-build-diffusers docker-build-kokoro docker-build-faster-whisper docker-build-whisper-medusa docker-build-crispasr docker-build-coqui docker-build-chatterbox docker-build-vibevoice docker-build-liquid-audio docker-build-moonshine docker-build-pocket-tts docker-build-qwen-tts docker-build-fish-speech docker-build-faster-qwen3-tts docker-build-qwen-asr docker-build-nemo docker-build-voxcpm docker-build-whisperx docker-build-ace-step docker-build-acestep-cpp docker-build-voxtral docker-build-mlx-distributed docker-build-trl docker-build-llama-cpp-quantization docker-build-tinygrad docker-build-kokoros docker-build-sam3-cpp docker-build-rfdetr-cpp docker-build-qwen3-tts-cpp docker-build-moss-tts-cpp docker-build-magpie-tts-cpp docker-build-vllm-cpp docker-build-omnivoice-cpp docker-build-vibevoice-cpp docker-build-localvqe docker-build-insightface docker-build-speaker-recognition docker-build-sherpa-onnx docker-build-cloud-proxy docker-build-supertonic docker-build-depth-anything-cpp docker-build-moss-transcribe-cpp docker-build-privacy-filter docker-build-trellis2cpp docker-build-valkey-store docker-build-audio-cpp
 
 ########################################################
 ### Mock Backend for E2E Tests

backend/index.yaml

@@ -502,7 +502,7 @@
   capabilities:
     nvidia: "cuda12-rfdetr"
     intel: "intel-rfdetr"
     #amd: "rocm-rfdetr"
     nvidia-l4t: "nvidia-l4t-arm64-rfdetr"
     metal: "metal-rfdetr"
     default: "cpu-rfdetr"
@@ -1457,6 +1457,25 @@
     nvidia-cuda-12: "cuda12-faster-whisper"
     nvidia-l4t: "nvidia-l4t-arm64-faster-whisper"
     nvidia-l4t-cuda-12: "nvidia-l4t-arm64-faster-whisper"
+- &whisper-medusa
+  description: |
+    Whisper-Medusa accelerates Whisper transcription by predicting multiple tokens per decoding step.
+    The upstream checkpoints use custom Transformers generation code and accept audio clips up to 30 seconds.
+  urls:
+    - https://github.com/aiola-lab/whisper-medusa
+    - https://huggingface.co/collections/aiola/whisper-medusa
+  tags:
+    - speech-to-text
+    - transcription
+    - Whisper
+    - Medusa
+  license: MIT
+  name: "whisper-medusa"
+  alias: "whisper-medusa"
+  capabilities:
+    default: "cpu-whisper-medusa"
+    nvidia: "cuda12-whisper-medusa"
+    nvidia-cuda-12: "cuda12-whisper-medusa"
 - &moonshine
   description: |
     Moonshine is a fast, accurate, and efficient speech-to-text transcription model using ONNX Runtime.
@@ -4171,7 +4190,7 @@
   capabilities:
     nvidia: "cuda12-rfdetr-development"
     intel: "intel-rfdetr-development"
     #amd: "rocm-rfdetr-development"
     nvidia-l4t: "nvidia-l4t-arm64-rfdetr-development"
     metal: "metal-rfdetr-development"
     default: "cpu-rfdetr-development"
@@ -6214,6 +6233,34 @@
   uri: "quay.io/go-skynet/local-ai-backends:master-nvidia-l4t-cuda-13-arm64-faster-qwen3-tts"
   mirrors:
     - localai/localai-backends:master-nvidia-l4t-cuda-13-arm64-faster-qwen3-tts
+## whisper-medusa
+- !!merge <<: *whisper-medusa
+  name: "whisper-medusa-development"
+  capabilities:
+    default: "cpu-whisper-medusa-development"
+    nvidia: "cuda12-whisper-medusa-development"
+    nvidia-cuda-12: "cuda12-whisper-medusa-development"
+- !!merge <<: *whisper-medusa
+  name: "cpu-whisper-medusa"
+  uri: "quay.io/go-skynet/local-ai-backends:latest-cpu-whisper-medusa"
+  mirrors:
+    - localai/localai-backends:latest-cpu-whisper-medusa
+- !!merge <<: *whisper-medusa
+  name: "cpu-whisper-medusa-development"
+  uri: "quay.io/go-skynet/local-ai-backends:master-cpu-whisper-medusa"
+  mirrors:
+    - localai/localai-backends:master-cpu-whisper-medusa
+- !!merge <<: *whisper-medusa
+  name: "cuda12-whisper-medusa"
+  uri: "quay.io/go-skynet/local-ai-backends:latest-gpu-nvidia-cuda-12-whisper-medusa"
+  mirrors:
+    - localai/localai-backends:latest-gpu-nvidia-cuda-12-whisper-medusa
+- !!merge <<: *whisper-medusa
+  name: "cuda12-whisper-medusa-development"
+  uri: "quay.io/go-skynet/local-ai-backends:master-gpu-nvidia-cuda-12-whisper-medusa"
+  mirrors:
+    - localai/localai-backends:master-gpu-nvidia-cuda-12-whisper-medusa
+
 ## qwen-asr
 - !!merge <<: *qwen-asr
   name: "qwen-asr-development"

backend/python/whisper-medusa/Makefile

@@ -0,0 +1,14 @@
+.DEFAULT_GOAL := install
+
+.PHONY: install
+install:
+	bash install.sh
+
+.PHONY: clean
+clean:
+	$(RM) backend_pb2_grpc.py backend_pb2.py
+	rm -rf venv __pycache__
+
+.PHONY: test
+test:
+	python3 -m unittest test_unit.py

backend/python/whisper-medusa/backend.py

@@ -0,0 +1,167 @@
+#!/usr/bin/env python3
+"""LocalAI gRPC backend for aiola Whisper-Medusa speech recognition."""
+
+import argparse
+from concurrent import futures
+import os
+import signal
+import sys
+import time
+
+import backend_pb2
+import backend_pb2_grpc
+import grpc
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "common"))
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "common"))
+from grpc_auth import get_auth_interceptors
+from model_utils import resolve_model_reference
+
+
+SAMPLE_RATE = 16000
+MAX_DURATION_SECONDS = 30
+MAX_WORKERS = int(os.environ.get("PYTHON_GRPC_MAX_WORKERS", "1"))
+_ONE_DAY_IN_SECONDS = 60 * 60 * 24
+
+
+def _parse_options(raw_options):
+    options = {}
+    for option in raw_options:
+        if ":" not in option:
+            continue
+        key, value = option.split(":", 1)
+        try:
+            value = int(value)
+        except ValueError:
+            try:
+                value = float(value)
+            except ValueError:
+                pass
+        options[key] = value
+    return options
+
+
+def _prepare_audio(path, torchaudio):
+    waveform, sample_rate = torchaudio.load(path)
+    if waveform.shape[0] > 1:
+        waveform = waveform.mean(dim=0, keepdim=True)
+    if sample_rate != SAMPLE_RATE:
+        waveform = torchaudio.transforms.Resample(sample_rate, SAMPLE_RATE)(waveform)
+        sample_rate = SAMPLE_RATE
+    return waveform, sample_rate
+
+
+class BackendServicer(backend_pb2_grpc.BackendServicer):
+    def __init__(self):
+        self.model = None
+        self.processor = None
+        self.device = None
+        self.options = {}
+
+    def Health(self, request, context):
+        return backend_pb2.Reply(message=b"OK")
+
+    def LoadModel(self, request, context):
+        try:
+            import torch
+            from transformers import WhisperProcessor
+            from whisper_medusa import WhisperMedusaModel
+
+            if request.CUDA and not torch.cuda.is_available():
+                return backend_pb2.Result(success=False, message="CUDA is not available")
+
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+            if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
+                self.device = torch.device("mps")
+
+            self.options = _parse_options(request.Options)
+            model_path, local_only = resolve_model_reference(
+                request, "aiola/whisper-medusa-linear-libri"
+            )
+            self.model = WhisperMedusaModel.from_pretrained(
+                model_path, local_files_only=local_only
+            ).to(self.device)
+            self.model.eval()
+            self.processor = WhisperProcessor.from_pretrained(
+                model_path, local_files_only=local_only
+            )
+        except Exception as err:
+            print(f"Whisper-Medusa model load failed: {err}", file=sys.stderr)
+            return backend_pb2.Result(success=False, message=str(err))
+
+        return backend_pb2.Result(success=True, message="Model loaded successfully")
+
+    def AudioTranscription(self, request, context):
+        if self.model is None or self.processor is None:
+            return backend_pb2.TranscriptResult(segments=[], text="")
+
+        try:
+            import torch
+            import torchaudio
+
+            waveform, sample_rate = _prepare_audio(request.dst, torchaudio)
+            duration = waveform.shape[-1] / sample_rate
+            if duration > MAX_DURATION_SECONDS:
+                raise ValueError(
+                    f"Whisper-Medusa supports audio clips up to {MAX_DURATION_SECONDS} seconds"
+                )
+
+            language = request.language or str(self.options.get("language", "en"))
+            regulation_start = int(self.options.get("regulation_start", 140))
+            regulation_factor = float(self.options.get("regulation_factor", 1.01))
+            features = self.processor(
+                waveform.squeeze(), return_tensors="pt", sampling_rate=sample_rate
+            ).input_features.to(self.device)
+            with torch.inference_mode():
+                output = self.model.generate(
+                    features,
+                    language=language,
+                    exponential_decay_length_penalty=(
+                        regulation_start,
+                        regulation_factor,
+                    ),
+                )
+            text = self.processor.decode(output[0], skip_special_tokens=True).strip()
+            segment = backend_pb2.TranscriptSegment(
+                id=0,
+                start=0,
+                end=int(duration * 1_000_000_000),
+                text=text,
+            )
+            return backend_pb2.TranscriptResult(segments=[segment], text=text)
+        except Exception as err:
+            print(f"Whisper-Medusa transcription failed: {err}", file=sys.stderr)
+            return backend_pb2.TranscriptResult(segments=[], text="")
+
+
+def serve(address):
+    server = grpc.server(
+        futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
+        options=[
+            ("grpc.max_send_message_length", 50 * 1024 * 1024),
+            ("grpc.max_receive_message_length", 50 * 1024 * 1024),
+        ],
+        interceptors=get_auth_interceptors(),
+    )
+    backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
+    server.add_insecure_port(address)
+    server.start()
+    print(f"Server started. Listening on: {address}", file=sys.stderr)
+
+    def stop_server(_signal, _frame):
+        server.stop(0)
+        sys.exit(0)
+
+    signal.signal(signal.SIGINT, stop_server)
+    signal.signal(signal.SIGTERM, stop_server)
+    try:
+        while True:
+            time.sleep(_ONE_DAY_IN_SECONDS)
+    except KeyboardInterrupt:
+        server.stop(0)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Run the Whisper-Medusa backend")
+    parser.add_argument("--addr", default="localhost:50051")
+    serve(parser.parse_args().addr)

backend/python/whisper-medusa/install.sh

@@ -0,0 +1,13 @@
+#!/bin/bash
+set -e
+
+backend_dir=$(dirname "$0")
+if [ -d "$backend_dir/common" ]; then
+    source "$backend_dir/common/libbackend.sh"
+else
+    source "$backend_dir/../common/libbackend.sh"
+fi
+
+PYTHON_VERSION="3.11"
+PYTHON_PATCH="13"
+installRequirements

backend/python/whisper-medusa/protogen.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+set -e
+python3 -m grpc_tools.protoc -I../../ --python_out=. --grpc_python_out=. ../../backend.proto

backend/python/whisper-medusa/requirements-cpu.txt

@@ -0,0 +1,3 @@
+--extra-index-url https://download.pytorch.org/whl/cpu
+torch==2.2.2
+torchaudio==2.2.2

backend/python/whisper-medusa/requirements-cublas12.txt

@@ -0,0 +1,3 @@
+--extra-index-url https://download.pytorch.org/whl/cu121
+torch==2.2.2
+torchaudio==2.2.2

backend/python/whisper-medusa/requirements.txt

@@ -0,0 +1,5 @@
+grpcio==1.71.0
+protobuf
+grpcio-tools
+transformers==4.49.0
+git+https://github.com/aiola-lab/whisper-medusa.git@19819c37ab15db6e68826e406614a2c86fbb946e

backend/python/whisper-medusa/run.sh

@@ -0,0 +1,9 @@
+#!/bin/bash
+backend_dir=$(dirname "$0")
+if [ -d "$backend_dir/common" ]; then
+    source "$backend_dir/common/libbackend.sh"
+else
+    source "$backend_dir/../common/libbackend.sh"
+fi
+
+startBackend "$@"