diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 96f24887bd..c83c8ecc1b 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -631,7 +631,8 @@ class ChatCompletionRequest(BaseModel): None, description = "[x-unsloth] Base64-encoded image for vision models" ) audio_base64: Optional[str] = Field( - None, description = "[x-unsloth] Base64-encoded WAV for audio-input models (ASR)" + None, + description = "[x-unsloth] Base64-encoded audio (wav/mp3/ogg/flac/m4a) for audio-input models", ) use_adapter: Optional[Union[bool, str]] = Field( None, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 418206e938..70a5a3efe7 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -2124,6 +2124,172 @@ def _decode_audio_base64(b64: str) -> np.ndarray: return waveform.squeeze(0).numpy() +# Reject oversized audio before decoding. base64 inflates raw bytes by ~4/3, so +# cap the encoded length to bound the upload. _MAX_AUDIO_SECONDS additionally +# bounds the *decoded* length, since a small compressed file (opus/flac/etc.) +# can expand to a far larger PCM array than the encoded-size cap implies. +_MAX_AUDIO_RAW_BYTES = 25 * 1024 * 1024 +_MAX_AUDIO_B64_CHARS = _MAX_AUDIO_RAW_BYTES * 4 // 3 +_MAX_AUDIO_SECONDS = 30 * 60 +_WAV_HEADER_BYTES = 44 +_MIN_TRANSCODE_AUDIO_SAMPLE_RATE = 8000 + + +def _sniff_audio_container(raw: bytes) -> Optional[str]: + """Return 'wav' or 'mp3' if the bytes are a container llama-server accepts + directly (so we can forward them untouched), else None (needs transcoding).""" + if len(raw) >= 12 and raw[:4] == b"RIFF" and raw[8:12] == b"WAVE": + return "wav" + # mp3: ID3 tag, or an MPEG audio frame sync (no other accepted format leads + # with 0xFF, so the simple sync check doesn't collide). + if raw[:3] == b"ID3" or (len(raw) >= 2 and raw[0] == 0xFF and (raw[1] & 0xE0) == 0xE0): + return "mp3" + return None + + +def _mono_f32_to_wav_bytes(arr: np.ndarray, sample_rate: int) -> bytes: + """Encode a mono float32 array as 16-bit PCM WAV bytes. + + Torch-free (numpy + stdlib only) so it works on no-torch GGUF-only installs; + the shared audio_codecs helper pulls in torch at import time. + """ + import io + import wave + + arr = np.nan_to_num(np.asarray(arr, dtype = np.float32).flatten(), posinf = 0.0, neginf = 0.0) + if arr.size == 0: + raise ValueError("decoded audio is empty") + peak = float(np.abs(arr).max()) + if peak > 1.0: + arr = arr / peak + pcm = (arr * 32767.0).astype(np.int16) + + buf = io.BytesIO() + with wave.open(buf, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(int(sample_rate)) + wf.writeframes(pcm.tobytes()) + return buf.getvalue() + + +def _resample_mono_linear(arr: np.ndarray, source_rate: int, target_rate: int) -> np.ndarray: + """Small numpy-only resampler for upload size limiting.""" + if source_rate <= 0 or target_rate <= 0 or source_rate == target_rate: + return arr + duration = len(arr) / float(source_rate) + target_len = max(1, int(round(duration * target_rate))) + if target_len == len(arr): + return arr + source_x = np.linspace(0.0, duration, num = len(arr), endpoint = False) + target_x = np.linspace(0.0, duration, num = target_len, endpoint = False) + return np.interp(target_x, source_x, arr).astype(np.float32) + + +def _fit_transcoded_audio_to_wav_cap(arr: np.ndarray, sample_rate: int) -> tuple[np.ndarray, int]: + """Downsample only when needed so transcoded WAV stays within the upload cap.""" + if sample_rate <= 0: + raise ValueError("decoded audio has an invalid sample rate") + wav_bytes = _WAV_HEADER_BYTES + len(arr) * 2 + if wav_bytes <= _MAX_AUDIO_RAW_BYTES: + return arr, sample_rate + + duration = len(arr) / float(sample_rate) + max_samples = max(1, (_MAX_AUDIO_RAW_BYTES - _WAV_HEADER_BYTES) // 2) + target_rate = int(max_samples // duration) + if target_rate < _MIN_TRANSCODE_AUDIO_SAMPLE_RATE: + raise ValueError("decoded audio exceeds the transcoded WAV size limit") + target_rate = min(sample_rate, target_rate) + fitted = _resample_mono_linear(arr, sample_rate, target_rate) + if _WAV_HEADER_BYTES + len(fitted) * 2 > _MAX_AUDIO_RAW_BYTES: + raise ValueError("decoded audio exceeds the transcoded WAV size limit") + return fitted, target_rate + + +def _decode_audio_mono(raw: bytes) -> tuple[np.ndarray, int]: + """Decode audio bytes to (mono float32 array, native sample_rate). + + soundfile (libsndfile) reads wav/mp3/ogg/flac straight from memory. librosa + (ffmpeg-backed) additionally covers m4a/webm but needs a real path and is + absent on no-torch GGUF-only installs. Both imports are inside the fallback + so a missing decoder degrades to the next one (and finally a clear error) + rather than crashing. + """ + import io + + try: + import soundfile as sf + arr, sr = sf.read(io.BytesIO(raw), dtype = "float32") + except Exception: + try: + import librosa + except ModuleNotFoundError as e: + raise RuntimeError( + "this audio format needs librosa, which is not installed in " + "GGUF-only environments; use wav, mp3, ogg or flac" + ) from e + import os + import tempfile + from utils.paths import ensure_dir, tmp_root + + with tempfile.NamedTemporaryFile( + suffix = ".audio", + delete = False, + dir = str(ensure_dir(tmp_root())), + ) as tmp: + tmp.write(raw) + tmp_path = tmp.name + try: + arr, sr = librosa.load(tmp_path, sr = None, mono = True) + finally: + os.unlink(tmp_path) + if arr.ndim > 1: + arr = arr.mean(axis = 1) + if sr > 0 and len(arr) > sr * _MAX_AUDIO_SECONDS: + raise ValueError(f"decoded audio exceeds the {_MAX_AUDIO_SECONDS // 60}-minute limit") + return arr, sr + + +def _prepare_audio_for_llama(b64: str) -> tuple[str, str]: + """Return (base64, format) ready for llama-server's input_audio part. + + llama-server's API only accepts wav/mp3, and decodes/resamples/down-mixes + them itself, so wav and mp3 uploads are forwarded untouched (no decode, no + PCM payload inflation). Other containers (m4a/ogg/webm/flac) are decoded to + a mono WAV. Blocking; call via a thread from async paths. + """ + if b64.startswith("data:"): + b64 = b64.split(",", 1)[1] if "," in b64 else "" + raw = base64.b64decode(b64) + passthrough = _sniff_audio_container(raw) + if passthrough is not None: + return b64, passthrough + + arr, sr = _decode_audio_mono(raw) + arr, sr = _fit_transcoded_audio_to_wav_cap(arr, sr) + return base64.b64encode(_mono_f32_to_wav_bytes(arr, sr)).decode("ascii"), "wav" + + +def _inject_audio_part(messages: list[dict], audio_b64: str, audio_format: str) -> None: + """Append an input_audio part to the last user message, in place. + + Audio rides in the message list like image_url parts do, so it flows through + both the plain and tool-calling generation paths. + """ + part = { + "type": "input_audio", + "input_audio": {"data": audio_b64, "format": audio_format}, + } + for msg in reversed(messages): + if msg.get("role") == "user": + content = msg.get("content") + if isinstance(content, list): + content.append(part) + else: + msg["content"] = [{"type": "text", "text": content or ""}, part] + return + + def _extract_content_parts(messages: list) -> tuple[str, list[dict], "Optional[str]"]: """ Parse OpenAI-format messages into components the inference backend expects. @@ -3081,9 +3247,12 @@ async def openai_chat_completions( if _wants_multiple_choices(payload): _raise_unsupported_n("GGUF tool or response_format passthrough") if payload.audio_base64: + # This path forwards the request verbatim, so the transcoded audio + # never gets injected. (The agentic tool loop below does support + # audio.) raise HTTPException( status_code = 400, - detail = "Audio input is not supported for GGUF chat models yet.", + detail = "Audio input is not supported together with guided decoding or client-supplied tools yet.", ) # Preserve the vision guard from the non-passthrough path below: @@ -3134,11 +3303,34 @@ async def openai_chat_completions( # ── GGUF path: proxy to llama-server /v1/chat/completions ── if using_gguf: + # Forward uploaded audio as an input_audio part. wav/mp3 pass through + # untouched (llama-server decodes and resamples them via the mmproj + # audio encoder); other containers are transcoded to WAV here. The part + # is injected into the message list below so it rides through both the + # plain and tool-calling paths, exactly like image_url parts. + audio_b64 = None + audio_format = "wav" if payload.audio_base64: - raise HTTPException( - status_code = 400, - detail = "Audio input is not supported for GGUF chat models yet.", - ) + if not getattr(llama_backend, "_has_audio_input", False): + raise HTTPException( + status_code = 400, + detail = "Audio provided but current GGUF model does not support audio input.", + ) + if len(payload.audio_base64) > _MAX_AUDIO_B64_CHARS: + raise HTTPException( + status_code = 413, + detail = "Audio file is too large (max ~25 MB).", + ) + try: + audio_b64, audio_format = await asyncio.to_thread( + _prepare_audio_for_llama, payload.audio_base64 + ) + except Exception as e: + logger.warning("Audio decode failed: %s", e, exc_info = True) + raise HTTPException( + status_code = 400, + detail = "Could not decode the provided audio file.", + ) gguf_messages, _ = _openai_messages_for_gguf_chat( payload, @@ -3146,6 +3338,8 @@ async def openai_chat_completions( ) gguf_messages = _set_or_prepend_system_message(gguf_messages, system_prompt) image_b64 = None + if audio_b64: + _inject_audio_part(gguf_messages, audio_b64, audio_format) cancel_event = threading.Event() diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 3763370dcd..0a165724d0 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -75,7 +75,6 @@ import { DocumentPreviewMount } from "@/features/rag/components/document-preview import { useUserProfileStore } from "@/features/profile/stores/user-profile-store"; import { applyQwenThinkingParams } from "@/features/chat/utils/qwen-params"; import { isTauri } from "@/lib/api-base"; -import { AUDIO_ACCEPT, MAX_AUDIO_SIZE, fileToBase64 } from "@/lib/audio-utils"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; import { toast } from "@/lib/toast"; import { cn } from "@/lib/utils"; @@ -1388,63 +1387,6 @@ function useImeComposerInputHandlers() { }; } -// Audio upload row, only for audio-input models. -const ComposerAudioMenuItem: FC = () => { - const setPendingAudio = useChatRuntimeStore((s) => s.setPendingAudio); - const activeModel = useChatRuntimeStore((s) => { - const checkpoint = s.params.checkpoint; - return s.models.find((m) => m.id === checkpoint); - }); - - const handleAudioFile = useCallback( - async (file: File) => { - if (file.size > MAX_AUDIO_SIZE) { - return; - } - try { - const base64 = await fileToBase64(file); - setPendingAudio(base64, file.name); - } catch { - // skip - } - }, - [setPendingAudio], - ); - - // Build the input on document.body, not in the menu: selecting the item - // closes the dropdown, unmounting a menu-rendered input before the OS picker - // returns and dropping the file. - const pickAudio = useCallback(() => { - const input = document.createElement("input"); - input.type = "file"; - input.accept = AUDIO_ACCEPT; - input.hidden = true; - document.body.appendChild(input); - input.onchange = () => { - const file = input.files?.[0]; - if (file) handleAudioFile(file); - document.body.removeChild(input); - }; - input.oncancel = () => { - if (!input.files || input.files.length === 0) { - document.body.removeChild(input); - } - }; - input.click(); - }, [handleAudioFile]); - - if (!activeModel?.hasAudioInput) { - return null; - } - - return ( - pickAudio()}> - - Upload audio - - ); -}; - // Phosphor microphone. Inlined to avoid a new icon dependency. const MicIcon: FC<{ className?: string }> = ({ className }) => ( { }; // Plus menu: attachment and workflow actions. Opens downward in the welcome // composer; the docked composer passes side="top" to open upward. +const AUDIO_ACCEPT_TOKEN_RE = + /^(audio\/|\.(?:wav|mp3|m4a|ogg|oga|flac)$)/i; + +function attachmentAcceptForPicker(accept: string, audioEnabled: boolean): string { + if (audioEnabled || accept === "*") { + return accept; + } + const filtered = accept + .split(",") + .map((token) => token.trim()) + .filter((token) => token && !AUDIO_ACCEPT_TOKEN_RE.test(token)) + .join(","); + return filtered || accept; +} + const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({ side = "bottom", }) => { @@ -2023,6 +1980,14 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({ const modelLoaded = useChatRuntimeStore( (s) => !!s.params.checkpoint && !s.modelLoading, ); + const audioAttachmentsEnabled = useChatRuntimeStore((s) => { + const activeCheckpoint = s.params.checkpoint; + if (!activeCheckpoint || s.modelLoading) { + return false; + } + const activeModel = s.models.find((m) => m.id === activeCheckpoint); + return Boolean(activeModel?.hasAudioInput); + }); const checkpoint = useChatRuntimeStore((s) => s.params.checkpoint); const supportsTools = useChatRuntimeStore((s) => s.supportsTools); const supportsBuiltinWebSearch = useChatRuntimeStore( @@ -2085,6 +2050,40 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({ const [promptStorageOpen, setPromptStorageOpen] = useState(false); const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId); const aui = useAui(); + const composerCanAddAttachments = useAuiState( + ({ composer }) => composer.isEditing, + ); + const pickAttachment = useCallback(() => { + const input = document.createElement("input"); + input.type = "file"; + input.multiple = true; + input.hidden = true; + + const attachmentAccept = attachmentAcceptForPicker( + aui.composer().getState().attachmentAccept, + audioAttachmentsEnabled, + ); + if (attachmentAccept !== "*") { + input.accept = attachmentAccept; + } + + document.body.appendChild(input); + input.onchange = (event) => { + const files = (event.target as HTMLInputElement).files; + if (files) { + for (const file of files) { + void aui.composer().addAttachment(file); + } + } + document.body.removeChild(input); + }; + input.oncancel = () => { + if (!input.files || input.files.length === 0) { + document.body.removeChild(input); + } + }; + input.click(); + }, [aui, audioAttachmentsEnabled]); // Disable Export chat until the thread has content. const messageCount = useAuiState(({ thread }) => thread.messages.length); const { startQueue } = useContext(PromptQueueContext); @@ -2136,13 +2135,13 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({ // Don't refocus the + on close; restored focus showed a stray ring. onCloseAutoFocus={(event) => event.preventDefault()} > - - - - Add photos & files - - - + pickAttachment()} + > + + Add photos & files + = 0; i -= 1) { const message = messages[i]; if (!message || message.role !== "user") continue; + // Message content parts (from compare view's CompareMessagePart with type: "audio") for (const part of message.content ?? []) { - if (part.type === "audio" && "audio" in part) { - const audioPart = ( - part as unknown as { - type: "audio"; - audio: string | { data: string; format: string }; - } - ).audio; - const raw = typeof audioPart === "string" ? audioPart : audioPart?.data; - if (raw) return raw.startsWith("data:") ? raw.split(",")[1] : raw; + const base64 = extractAudioPartBase64(part); + if (base64) return base64; + } + + // Attachment content parts (from AudioAttachmentAdapter) + if ("attachments" in message) { + for (const attachment of message.attachments ?? []) { + for (const part of attachment.content ?? []) { + const base64 = extractAudioPartBase64(part); + if (base64) return base64; + } } } + + // Only the newest user message counts. audio_base64 switches the + // backend onto the audio generation path, so replaying audio from an + // older turn would hijack text follow-ups (Whisper would retranscribe + // the stale clip). Matches the consumed-on-send semantics of the + // legacy pendingAudio path. + break; } // Runtime store (main composer's audio upload). diff --git a/studio/frontend/src/features/chat/audio-attachment-adapter.ts b/studio/frontend/src/features/chat/audio-attachment-adapter.ts new file mode 100644 index 0000000000..148f025e57 --- /dev/null +++ b/studio/frontend/src/features/chat/audio-attachment-adapter.ts @@ -0,0 +1,96 @@ +// SPDX-License-Identifier: AGPL-3.0-only +// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +import { + AUDIO_ACCEPT, + MAX_AUDIO_SIZE, + fileToBase64, +} from "@/lib/audio-utils"; +import type { + Attachment, + AttachmentAdapter, + CompleteAttachment, + PendingAttachment, +} from "@assistant-ui/react"; +import { toast } from "sonner"; +import { useChatRuntimeStore } from "./stores/chat-runtime-store"; + +// crypto.randomUUID is undefined in non-secure contexts (HTTP over a LAN IP). +function newAttachmentId(): string { + if (typeof globalThis.crypto?.randomUUID === "function") { + return globalThis.crypto.randomUUID(); + } + return `${Date.now()}-${Math.random().toString(36).slice(2, 10)}`; +} + +// Audio shares the "Add photos & files" picker. Like VisionImageAdapter, +// unsupported models are rejected at add() time with a toast. +export class AudioAttachmentAdapter implements AttachmentAdapter { + // MIME is unreliable for some containers (m4a), so also match by + // extension. No .webm extension: it would claim video/webm files; real + // audio webm (MediaRecorder) always reports the audio/webm MIME. + accept = `${AUDIO_ACCEPT},audio/x-m4a,.wav,.mp3,.m4a,.ogg,.oga,.flac`; + private readonly attachmentIds = new Set(); + + async add({ file }: { file: File }): Promise { + const state = useChatRuntimeStore.getState(); + const checkpoint = state.params.checkpoint; + const activeModel = state.models.find((m) => m.id === checkpoint); + const modelLoaded = !!checkpoint && !state.modelLoading; + let unavailableReason: string | null = null; + if (!modelLoaded) { + unavailableReason = "Load a model before adding audio files."; + } else if (!activeModel?.hasAudioInput) { + const label = activeModel?.name || checkpoint || "Current model"; + unavailableReason = `${label} cannot accept audio. Load an audio-input model before attaching audio files.`; + } + if (unavailableReason) { + toast.error(unavailableReason); + throw new Error(unavailableReason); + } + if (file.size > MAX_AUDIO_SIZE) { + const sizeReason = "Audio size exceeds 50MB limit"; + toast.error(sizeReason); + throw new Error(sizeReason); + } + if (this.attachmentIds.size > 0 || state.pendingAudioBase64) { + const duplicateReason = "Only one audio file can be attached per message."; + toast.error(duplicateReason); + throw new Error(duplicateReason); + } + + const id = newAttachmentId(); + this.attachmentIds.add(id); + return { + id, + type: "file", + name: file.name, + contentType: file.type, + file, + status: { type: "requires-action", reason: "composer-send" }, + }; + } + + async send(attachment: PendingAttachment): Promise { + try { + const base64 = await fileToBase64(attachment.file); + // Backend takes raw base64; format only satisfies the part type. + const format = attachment.contentType === "audio/mpeg" ? "mp3" : "wav"; + return { + id: attachment.id, + type: "file", + name: attachment.name, + contentType: attachment.contentType, + content: [{ type: "audio", audio: { data: base64, format } }], + status: { type: "complete" }, + }; + } finally { + this.attachmentIds.delete(attachment.id); + } + } + + remove(attachment: Attachment): Promise { + this.attachmentIds.delete(attachment.id); + return Promise.resolve(); + } +} diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index a152cbb358..66ee2d04e4 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts @@ -142,6 +142,50 @@ function toChatModelSummary(model: { }; } +// Merge capability flags from a load/status response into the matching +// models[] entry. /api/models/list omits audio capability for default and +// active-GGUF entries, so the attach gates (`activeModel?.hasAudioInput`) +// would otherwise stay false. Mirrors the compare composer's sync. +// Exported for tests. +export function syncModelCapabilities( + modelId: string, + resp: { + display_name?: string | null; + is_vision?: boolean; + is_lora?: boolean; + is_gguf?: boolean; + is_audio?: boolean; + audio_type?: string | null; + has_audio_input?: boolean; + }, +): void { + const store = useChatRuntimeStore.getState(); + const models = store.models; + const synced = { + isVision: Boolean(resp.is_vision), + isGguf: Boolean(resp.is_gguf), + isAudio: Boolean(resp.is_audio), + audioType: resp.audio_type ?? null, + hasAudioInput: Boolean(resp.has_audio_input), + }; + const idx = models.findIndex((m) => m.id === modelId); + if (idx === -1) { + store.setModels([ + ...models, + { + id: modelId, + name: resp.display_name || modelId, + isLora: Boolean(resp.is_lora), + ...synced, + }, + ]); + } else { + const next = [...models]; + next[idx] = { ...next[idx], ...synced }; + store.setModels(next); + } +} + function toLoraSummary(lora: { display_name: string; adapter_path: string; @@ -390,6 +434,9 @@ export function useChatModelRuntime() { loadedChatTemplateOverride: statusRes.chat_template_override, }), }); + // setModels(listRes...) above used catalog data, which omits audio + // capability. Re-apply live status so attach gates survive a refresh. + syncModelCapabilities(statusRes.active_model, statusRes); // Set reasoning default for Qwen3.5/3.6 small models if ( @@ -736,6 +783,8 @@ export function useChatModelRuntime() { loadedIsMultimodal: isMultimodalResponse(loadResponse), activeNativePathToken: nativePathToken ?? null, }); + // Unlock attach menus for capabilities the catalog entry lacked. + syncModelCapabilities(modelId, loadResponse); // Qwen3/3.5/3.6: apply thinking-mode-specific params after load if ( modelId.toLowerCase().includes("qwen3") && diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index 8d9e5a10e3..359301e182 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/studio/frontend/src/features/chat/runtime-provider.tsx @@ -51,6 +51,7 @@ import { readActiveOpenDocumentAttachmentContent, readOpenDocumentAttachmentContent, } from "./open-document"; +import { AudioAttachmentAdapter } from "./audio-attachment-adapter"; import { useChatRuntimeStore } from "./stores/chat-runtime-store"; import type { MessageRecord, ModelType, ThreadRecord } from "./types"; import { @@ -983,6 +984,7 @@ function useStudioRuntimeAdapters( () => new CompositeAttachmentAdapter([ new VisionImageAdapter(), + new AudioAttachmentAdapter(), new TextAttachmentAdapter(), new HtmlAttachmentAdapter(), new PDFAttachmentAdapter(),