diff --git a/vibevoice-asr/README.md b/vibevoice-asr/README.md new file mode 100644 index 00000000..9e6a8060 --- /dev/null +++ b/vibevoice-asr/README.md @@ -0,0 +1,74 @@ +# VibeVoice-ASR + +[Microsoft VibeVoice-ASR](https://huggingface.co/microsoft/VibeVoice-ASR) is a multimodal speech-to-text model deployed via [vLLM](https://github.com/vllm-project/vllm) with an OpenAI-compatible chat completions API. It transcribes audio into JSON-shaped segments with speaker labels and timestamps. + +## Hardware + +- **GPU:** H100 (single) +- **System memory:** 32 GiB +- **Cold start:** ~3 minutes (HF snapshot + tokenizer file generation + vLLM warmup) + +## Deployment + +```bash +truss push --remote --publish +``` + +A Hugging Face access token is required as a Baseten secret named `hf_access_token` to pull the model weights. + +## API + +The deployment exposes the OpenAI chat completions endpoint. Point any OpenAI-compatible client at it: + +```python +from openai import OpenAI + +client = OpenAI( + api_key="", + base_url="https://model-.api.baseten.co/environments/production/sync/v1", +) + +response = client.chat.completions.create( + model="vibevoice", + messages=[ + {"role": "system", "content": "You are a helpful assistant that transcribes audio input into text output in JSON format."}, + {"role": "user", "content": [ + {"type": "audio_url", "audio_url": {"url": "https://github.com/ggerganov/whisper.cpp/raw/master/samples/jfk.wav"}}, + {"type": "text", "text": "Transcribe this audio."}, + ]}, + ], + max_tokens=64, + temperature=0, +) + +print(response.choices[0].message.content) +``` + +The model returns a JSON-stringified array of segments in the assistant message content: + +```json +[ + {"Start": 0.0, "End": 10.46, "Speaker": 0, + "Content": "And so, my fellow Americans, ask not what your country can do for you, ask what you can do for your country."} +] +``` + +You can also pass audio inline as a base64 data URI: `data:audio/wav;base64,<...>`. + +## Known limitations + +- **WAV (or FLAC) audio only.** vLLM 0.14.1's audio loader can't decode m4a from a buffer; m4a requests return a 400. ffmpeg-convert client-side first. +- **Size `max_tokens` proportional to audio length.** VibeVoice-ASR doesn't reliably emit an end-of-sequence token, so an over-generous `max_tokens` produces a repetition loop after the real transcript ends. A safe rule: `max_tokens ≈ 20 * audio_length_sec + 200`. +- **One audio clip per request.** The model is trained on single-audio chat messages; multiple `audio_url` items in one user turn are not supported. + +## About the patch.py + +`data/patch.py` applies three small in-place fixes to the installed `vllm_plugin/model.py` at container boot: + +1. **KV-cache delegation forwarders** on `VibeVoiceForCausalLM` — required for vLLM 0.14.1. Without these, vLLM sees zero attention layers and crashes at startup with `IndexError: list index out of range` on `available_gpu_memory[0]`. +2. **`get_data_parser` on the Info class** — forward-compatibility shim for vLLM 0.21+ which calls the parser on the Info class instead of the Processor class. +3. **`mm_data_items` rename try/except** — forward-compatibility shim for vLLM 0.15+ which renamed the `ProcessorInputs.mm_data` field. + +These are upstream plugin issues. When Microsoft fixes them, delete `data/patch.py` and remove the `python3 /app/data/patch.py` line from `config.yaml`'s `start_command`. + +The patch script is idempotent and uses `assert`-checked anchor strings, so a plugin upstream change that breaks the anchors causes a clear failure at startup rather than silent miscompile. diff --git a/vibevoice-asr/config.yaml b/vibevoice-asr/config.yaml new file mode 100644 index 00000000..dc6545ac --- /dev/null +++ b/vibevoice-asr/config.yaml @@ -0,0 +1,107 @@ +model_name: VibeVoice-ASR +description: Microsoft VibeVoice-ASR — OpenAI-compatible audio transcription via vLLM +python_version: py310 + +base_image: + image: vllm/vllm-openai:v0.14.1 + python_executable_path: /usr/bin/python3 + +environment_variables: + HF_HOME: /cache/org + HF_HUB_CACHE: /cache/org + TRANSFORMERS_CACHE: /cache/org + VIBEVOICE_FFMPEG_MAX_CONCURRENCY: "64" + VLLM_MEDIA_LOADING_THREAD_COUNT: "16" + PYTORCH_ALLOC_CONF: "expandable_segments:True" + +requirements: + - transformers==4.57.6 + - accelerate>=0.30.0 + - safetensors + - huggingface-hub>=0.23.0 + - librosa>=0.10.0 + - soundfile + - scipy + - pydub + - diffusers + - git+https://github.com/microsoft/VibeVoice.git@main + +resources: + accelerator: H100 + cpu: "4" + memory: 32Gi + use_gpu: true + +runtime: + predict_concurrency: 32 + +secrets: + hf_access_token: null + +system_packages: + - ffmpeg + - git + +# Weights are pre-downloaded at build time and mounted at /models/vibevoice-asr, +# so cold starts skip the 9.2 GB HF download entirely. +weights: + - source: hf://microsoft/VibeVoice-ASR@main + mount_location: /models/vibevoice-asr + auth: + auth_method: CUSTOM_SECRET + auth_secret_name: hf_access_token + +# Pass-through mode: no model.py, Truss just runs vllm serve and proxies +# /predict requests to /v1/chat/completions on the container's localhost. +docker_server: + server_port: 8000 + predict_endpoint: /v1/chat/completions + readiness_endpoint: /v1/models + liveness_endpoint: /v1/models + start_command: | + bash -c ' + set -e + echo "[entrypoint] Applying microsoft/VibeVoice plugin patches..." + python3 /app/data/patch.py + echo "[entrypoint] Generating tokenizer files..." + python3 -m vllm_plugin.tools.generate_tokenizer_files --output /models/vibevoice-asr + echo "[entrypoint] Starting vLLM serve..." + exec vllm serve /models/vibevoice-asr \ + --served-model-name vibevoice \ + --trust-remote-code \ + --dtype bfloat16 \ + --max-num-seqs 16 \ + --max-model-len 32768 \ + --gpu-memory-utilization 0.85 \ + --num-gpu-blocks-override 4096 \ + --no-enable-prefix-caching \ + --enable-chunked-prefill \ + --chat-template-content-format openai \ + --allowed-local-media-path /app \ + --media-io-kwargs "{\"audio\": {\"target_sr\": 24000}}" \ + --enforce-eager \ + --skip-mm-profiling \ + --host 0.0.0.0 \ + --port 8000 + ' + +model_metadata: + tags: + - openai-compatible + - audio + - asr + - speech-to-text + example_model_input: + model: vibevoice + messages: + - role: system + content: You are a helpful assistant that transcribes audio input into text output in JSON format. + - role: user + content: + - type: audio_url + audio_url: + url: https://github.com/ggerganov/whisper.cpp/raw/master/samples/jfk.wav + - type: text + text: Transcribe this audio. + max_tokens: 64 + temperature: 0.0 diff --git a/vibevoice-asr/data/patch.py b/vibevoice-asr/data/patch.py new file mode 100644 index 00000000..9fad8860 --- /dev/null +++ b/vibevoice-asr/data/patch.py @@ -0,0 +1,119 @@ +"""Apply microsoft/VibeVoice plugin patches at container boot. + +The Microsoft VibeVoice vLLM plugin (installed via `pip install +git+https://github.com/microsoft/VibeVoice.git`) has three known issues +against current vLLM releases. This script patches the installed +`vllm_plugin/model.py` file in-place at boot, before `vllm serve` starts. + +Patches applied: + 1. KV-cache delegation forwarders on VibeVoiceForCausalLM (REQUIRED for vLLM 0.14.1) + Without this, vLLM 0.14.1 sees zero attention layers and startup crashes + with `IndexError: list index out of range` on `available_gpu_memory[0]`. + + 2. get_data_parser method on VibeVoiceProcessingInfo (forward-compat for vLLM 0.21+) + v0.14.1 calls `_get_data_parser` on the Processor; v0.21+ calls `get_data_parser` + on the Info class. No-op on v0.14.1. + + 3. mm_data_items field rename try/except (forward-compat for vLLM 0.15+) + v0.14.1 ProcessorInputs uses `mm_data`; v0.15+ uses `mm_data_items`. + No-op on v0.14.1. + +Idempotent: re-running is a no-op once the marker is present. +Anchor-checked: each insertion `assert`s its target exists, so an upstream +plugin update that breaks the anchors causes a loud startup error instead +of silent miscompile. + +Once Microsoft fixes the plugin upstream, delete this file and remove the +patch.py call from config.yaml's start_command. +""" + +import os +import sys +import vllm_plugin + +TARGET = os.path.join(os.path.dirname(vllm_plugin.__file__), "model.py") +MARKER = "[Truss local patch]" + +src = open(TARGET).read() + +if MARKER in src: + print(f"[patch.py] {TARGET} already patched, skipping") + sys.exit(0) + + +# ── Patch 1: KV-cache delegation forwarders (REQUIRED on vLLM 0.14.1) ──────── +ANCHOR_1 = ( + " def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:\n" + " return self.language_model.compute_logits(hidden_states)" +) +INSERT_1 = """ + + # [Truss local patch] vLLM 0.14.1 KV-cache discovery looks for + # `.get_kv_cache_spec()` and `.model` on the top-level registered class. + # VibeVoiceForCausalLM is a wrapper around self.language_model — forward + # these so vLLM sees the real attention layers. Without this, startup + # crashes with IndexError on `available_gpu_memory[0]`. + def get_kv_cache_spec(self, *args, **kwargs): + return self.language_model.get_kv_cache_spec(*args, **kwargs) + + @property + def model(self): + return self.language_model.model""" + +assert ANCHOR_1 in src, ( + "Patch 1 anchor (compute_logits signature) not found in plugin — upstream may have changed" +) +src = src.replace(ANCHOR_1, ANCHOR_1 + INSERT_1) + + +# ── Patch 2: get_data_parser on Info class (forward-compat for vLLM 0.21+) ── +ANCHOR_2 = ( + "class VibeVoiceProcessingInfo(BaseProcessingInfo):\n" + ' """Processing info for VibeVoice multimodal model."""' +) +INSERT_2 = """ + + # [Truss local patch] vLLM 0.21+ calls info.get_data_parser() on this + # class (the plugin only defines _get_data_parser on the Processor). + # No-op on v0.14.1 which uses the older API. + def get_data_parser(self) -> MultiModalDataParser: + return MultiModalDataParser(target_sr=24000)""" + +assert ANCHOR_2 in src, ( + "Patch 2 anchor (VibeVoiceProcessingInfo class) not found in plugin" +) +src = src.replace(ANCHOR_2, ANCHOR_2 + INSERT_2) + + +# ── Patch 3: mm_data_items rename try/except (forward-compat for vLLM 0.15+) ─ +ANCHOR_3 = ( + ' """Build ProcessorInputs for dummy profiling."""\n' + " return ProcessorInputs(\n" + " prompt=self.get_dummy_text(mm_counts),\n" + " mm_data=self.get_dummy_mm_data(seq_len, mm_counts, mm_options),\n" + " )" +) +REPLACE_3 = ''' """[Truss local patch] vLLM 0.15+ renamed mm_data → mm_data_items. + Try the new signature first, fall back to legacy for v0.14.1.""" + mm_data_dict = self.get_dummy_mm_data(seq_len, mm_counts, mm_options) + try: + mm_data_items = MultiModalDataParser().parse_mm_data(mm_data_dict) + return ProcessorInputs( + prompt=self.get_dummy_text(mm_counts), + mm_data_items=mm_data_items, + ) + except TypeError: + return ProcessorInputs( + prompt=self.get_dummy_text(mm_counts), + mm_data=mm_data_dict, + )''' + +assert ANCHOR_3 in src, ( + "Patch 3 anchor (get_dummy_processor_inputs body) not found in plugin" +) +src = src.replace(ANCHOR_3, REPLACE_3) + + +# Write back +open(TARGET, "w").write(src) +print(f"[patch.py] applied 3 patches to {TARGET}")