diff --git a/.github/backend-matrix.yml b/.github/backend-matrix.yml index f95311be85dd..d134111dec62 100644 --- a/.github/backend-matrix.yml +++ b/.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" diff --git a/Makefile b/Makefile index 6d64540b8140..4ce90dc3e2e4 100644 --- a/Makefile +++ b/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 diff --git a/backend/index.yaml b/backend/index.yaml index a8f426f1bf84..0ec241e2bf7c 100644 --- a/backend/index.yaml +++ b/backend/index.yaml @@ -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. @@ -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" diff --git a/backend/python/whisper-medusa/Makefile b/backend/python/whisper-medusa/Makefile new file mode 100644 index 000000000000..34f3ad4f01d3 --- /dev/null +++ b/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 diff --git a/backend/python/whisper-medusa/backend.py b/backend/python/whisper-medusa/backend.py new file mode 100755 index 000000000000..e7116fa4cbcc --- /dev/null +++ b/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) diff --git a/backend/python/whisper-medusa/install.sh b/backend/python/whisper-medusa/install.sh new file mode 100755 index 000000000000..e455c67febfe --- /dev/null +++ b/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 diff --git a/backend/python/whisper-medusa/protogen.sh b/backend/python/whisper-medusa/protogen.sh new file mode 100755 index 000000000000..130c8b83ea89 --- /dev/null +++ b/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 diff --git a/backend/python/whisper-medusa/requirements-cpu.txt b/backend/python/whisper-medusa/requirements-cpu.txt new file mode 100644 index 000000000000..bc99310da135 --- /dev/null +++ b/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 diff --git a/backend/python/whisper-medusa/requirements-cublas12.txt b/backend/python/whisper-medusa/requirements-cublas12.txt new file mode 100644 index 000000000000..a7892a753c4e --- /dev/null +++ b/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 diff --git a/backend/python/whisper-medusa/requirements.txt b/backend/python/whisper-medusa/requirements.txt new file mode 100644 index 000000000000..64bb72432503 --- /dev/null +++ b/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 diff --git a/backend/python/whisper-medusa/run.sh b/backend/python/whisper-medusa/run.sh new file mode 100755 index 000000000000..94540e7498a2 --- /dev/null +++ b/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 "$@" diff --git a/backend/python/whisper-medusa/test.sh b/backend/python/whisper-medusa/test.sh new file mode 100755 index 000000000000..d942f2d4150a --- /dev/null +++ b/backend/python/whisper-medusa/test.sh @@ -0,0 +1,26 @@ +#!/bin/bash +set -e +python3 -m unittest test_unit.py + +tmp_dir=$(mktemp -d) +trap 'rm -rf "$tmp_dir"' EXIT +mkdir -p "$tmp_dir/bin" +cat >"$tmp_dir/bin/curl" <<'EOF' +#!/bin/bash +printf '%s\n' "${@: -1}" >"$PORTABLE_PY_URL_CAPTURE" +exit 22 +EOF +chmod +x "$tmp_dir/bin/curl" + +PORTABLE_PY_URL_CAPTURE="$tmp_dir/url" \ + ENV_DIR="$tmp_dir/backend" \ + PORTABLE_PYTHON=true \ + PATH="$tmp_dir/bin:$PATH" \ + bash install.sh >/dev/null 2>&1 || true + +expected_url="https://github.com/astral-sh/python-build-standalone/releases/download/20250818/cpython-3.11.13+20250818-x86_64-unknown-linux-gnu-install_only.tar.gz" +actual_url=$(cat "$tmp_dir/url") +if [ "$actual_url" != "$expected_url" ]; then + echo "unexpected portable Python URL: $actual_url" >&2 + exit 1 +fi diff --git a/backend/python/whisper-medusa/test_unit.py b/backend/python/whisper-medusa/test_unit.py new file mode 100644 index 000000000000..407fea5c2a35 --- /dev/null +++ b/backend/python/whisper-medusa/test_unit.py @@ -0,0 +1,89 @@ +import importlib.util +import pathlib +import sys +import types +import unittest + + +class _Message: + def __init__(self, **kwargs): + self.__dict__.update(kwargs) + + +backend_pb2 = types.ModuleType("backend_pb2") +for name in ("Reply", "Result", "TranscriptResult", "TranscriptSegment"): + setattr(backend_pb2, name, _Message) +sys.modules["backend_pb2"] = backend_pb2 + +backend_pb2_grpc = types.ModuleType("backend_pb2_grpc") +backend_pb2_grpc.BackendServicer = object +backend_pb2_grpc.add_BackendServicer_to_server = lambda *args: None +sys.modules["backend_pb2_grpc"] = backend_pb2_grpc + +grpc = types.ModuleType("grpc") +grpc.server = lambda *args, **kwargs: None +sys.modules["grpc"] = grpc + +grpc_auth = types.ModuleType("grpc_auth") +grpc_auth.get_auth_interceptors = lambda: [] +sys.modules["grpc_auth"] = grpc_auth + +model_utils = types.ModuleType("model_utils") +model_utils.resolve_model_reference = lambda request, default: (request.Model or default, False) +sys.modules["model_utils"] = model_utils + +spec = importlib.util.spec_from_file_location( + "whisper_medusa_backend", pathlib.Path(__file__).with_name("backend.py") +) +backend = importlib.util.module_from_spec(spec) +spec.loader.exec_module(backend) + + +class FakeTensor: + def __init__(self, channels=1, samples=16000): + self.channels = channels + self.samples = samples + self.shape = (channels, samples) + self.mean_calls = [] + + def mean(self, dim, keepdim): + self.mean_calls.append((dim, keepdim)) + return FakeTensor(1, self.samples) + + def squeeze(self): + return self + + def to(self, device): + return self + + +class FakeTorchaudio: + class transforms: + class Resample: + def __init__(self, source, target): + self.source = source + self.target = target + + def __call__(self, waveform): + return FakeTensor(waveform.channels, waveform.samples * self.target // self.source) + + @staticmethod + def load(path): + return FakeTensor(2, 8000), 8000 + + +class BackendHelpersTest(unittest.TestCase): + def test_parse_options_converts_supported_scalar_types(self): + self.assertEqual( + backend._parse_options(["regulation_start:120", "regulation_factor:1.05", "ignored"]), + {"regulation_start": 120, "regulation_factor": 1.05}, + ) + + def test_prepare_audio_mixes_to_mono_and_resamples_to_16khz(self): + waveform, sample_rate = backend._prepare_audio("clip.wav", FakeTorchaudio) + self.assertEqual(sample_rate, 16000) + self.assertEqual(waveform.shape, (1, 16000)) + + +if __name__ == "__main__": + unittest.main() diff --git a/docs/content/features/backends.md b/docs/content/features/backends.md index 0e7c960a6d1c..f278ba58e81a 100644 --- a/docs/content/features/backends.md +++ b/docs/content/features/backends.md @@ -130,7 +130,7 @@ For getting started, see the available backends in LocalAI here: https://github. LocalAI supports various types of backends: - **LLM Backends**: For running language models (e.g., llama.cpp, vLLM, vllm.cpp, SGLang, transformers, MLX) -- **Speech-to-Text Backends**: For transcription, forced alignment and speaker diarization (e.g., whisper.cpp, parakeet.cpp, moss-transcribe.cpp, faster-whisper, NeMo, [audio.cpp]({{%relref "features/audio-cpp" %}})) +- **Speech-to-Text Backends**: For transcription, forced alignment and speaker diarization (e.g., whisper.cpp, parakeet.cpp, moss-transcribe.cpp, faster-whisper, [Whisper-Medusa]({{%relref "features/whisper-medusa" %}}), NeMo, [audio.cpp]({{%relref "features/audio-cpp" %}})) - **Text-to-Speech Backends**: For speech synthesis (e.g., piper, Kokoro, VibeVoice, Qwen3-TTS, [audio.cpp]({{%relref "features/audio-cpp" %}})) - **Sound Generation Backends**: For music and audio generation (e.g., ACE-Step, [audio.cpp]({{%relref "features/audio-cpp" %}})) - **Sound Classification Backends**: For sound-event classification / audio tagging - identifying everyday sounds like baby cry, glass breaking, alarms (e.g., ced.cpp) diff --git a/docs/content/features/whisper-medusa.md b/docs/content/features/whisper-medusa.md new file mode 100644 index 000000000000..83ce2f0bb08b --- /dev/null +++ b/docs/content/features/whisper-medusa.md @@ -0,0 +1,48 @@ +--- +title: "Whisper-Medusa" +description: "Run Whisper-Medusa speech-to-text checkpoints with LocalAI" +--- + +The `whisper-medusa` backend serves aiola's Whisper-Medusa checkpoints through +LocalAI's OpenAI-compatible audio transcription endpoint. Whisper-Medusa uses +multiple decoding heads to predict several tokens per decoding step. + +## Model configuration + +Install the `whisper-medusa` backend from the backend gallery, then create a +model configuration such as: + +```yaml +name: whisper-medusa +backend: whisper-medusa +parameters: + model: aiola/whisper-medusa-linear-libri +options: + - language:en + - regulation_start:140 + - regulation_factor:1.01 +``` + +Transcribe a clip with the standard endpoint: + +```bash +curl http://localhost:8080/v1/audio/transcriptions \ + -F file=@audio.wav \ + -F model=whisper-medusa \ + -F language=en +``` + +`language` in the request overrides the configured default. The regulation +options control the exponential length penalty passed to upstream generation. + +{{% notice warning %}} +The upstream Whisper-Medusa repository is archived. Its implementation accepts +clips of at most 30 seconds and expects 16 kHz audio; LocalAI resamples input to +16 kHz and rejects longer clips. The LibriSpeech checkpoints are optimized for +English. Use `aiola/whisper-medusa-multilingual` for supported multilingual +audio. +{{% /notice %}} + +The backend currently ships Linux CPU and NVIDIA CUDA 12 images. It is not +advertised for macOS, ROCm, Intel GPU, or Jetson because upstream pins PyTorch +2.2.2 and has not published compatibility for those platforms. diff --git a/docs/superpowers/plans/2026-08-02-whisper-medusa-backend.md b/docs/superpowers/plans/2026-08-02-whisper-medusa-backend.md new file mode 100644 index 000000000000..3e35636e9168 --- /dev/null +++ b/docs/superpowers/plans/2026-08-02-whisper-medusa-backend.md @@ -0,0 +1,63 @@ +# Whisper-Medusa Backend Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add a dedicated LocalAI speech-to-text backend for aiola Whisper-Medusa checkpoints. + +**Architecture:** A Python gRPC backend owns model loading, audio normalization, and Medusa generation. LocalAI's existing `AudioTranscription` RPC remains unchanged; build, gallery, and documentation surfaces follow the existing Python ASR backend pattern. + +**Tech Stack:** Python 3.11, PyTorch, torchaudio, transformers, whisper-medusa, gRPC, YAML, Make. + +## Global Constraints + +- Accept local model paths and Hugging Face model identifiers through `resolve_model_reference`. +- Normalize input audio to mono 16 kHz before generation. +- Default the language to `en` and expose upstream generation regulation options. +- Document upstream's archived status, 30-second clip limit, and checkpoint language limitations. +- Build Linux CPU and NVIDIA CUDA 12 images; do not claim unsupported Darwin or ROCm coverage. + +--- + +### Task 1: Backend behavior + +**Files:** +- Create: `backend/python/whisper-medusa/backend.py` +- Create: `backend/python/whisper-medusa/test_unit.py` + +**Interfaces:** +- Consumes: LocalAI `LoadModel` and `AudioTranscription` protobuf requests. +- Produces: `BackendServicer`, `_parse_options`, and `_prepare_audio`. + +- [ ] Write unit tests for option parsing, mono conversion, resampling, load failure, and transcription. +- [ ] Run `python -m unittest test_unit.py` and confirm it fails because the backend does not exist. +- [ ] Implement the minimal gRPC backend and rerun the unit tests to green. + +### Task 2: Packaging and registration + +**Files:** +- Create: `backend/python/whisper-medusa/{Makefile,install.sh,protogen.sh,run.sh,test.sh,requirements.txt,requirements-cpu.txt,requirements-cublas12.txt}` +- Modify: `Makefile` +- Modify: `.github/backend-matrix.yml` +- Modify: `backend/index.yaml` + +**Interfaces:** +- Consumes: the Python backend Docker build conventions. +- Produces: `whisper-medusa` install/build targets and CPU/CUDA backend images. + +- [ ] Add packaging scripts and pinned upstream dependency. +- [ ] Register the backend in Make and backend metadata. +- [ ] Add Linux amd64 CPU and CUDA 12 CI matrix entries. +- [ ] Validate Make and YAML parsing. + +### Task 3: Documentation and verification + +**Files:** +- Create: `docs/content/features/whisper-medusa.md` +- Modify: `docs/content/features/backends.md` + +**Interfaces:** +- Produces: user-facing model YAML and limitation guidance. + +- [ ] Document setup, configuration, options, and upstream constraints. +- [ ] Run unit tests, syntax compilation, registration checks, and diff review. +- [ ] Commit with the required `Assisted-by` trailer, push, and open a PR closing issue #3127.