diff --git a/studio/backend/core/inference/audio_codecs.py b/studio/backend/core/inference/audio_codecs.py index 3a418d921d..bcf3ec2937 100644 --- a/studio/backend/core/inference/audio_codecs.py +++ b/studio/backend/core/inference/audio_codecs.py @@ -302,6 +302,28 @@ class AudioCodecManager: waveform = audio.squeeze().cpu().numpy() return _numpy_to_wav_bytes(waveform, 24000), 24000 + def decode( + self, + audio_type: str, + device: str, + token_ids: Optional[list] = None, + text: Optional[str] = None, + ) -> Tuple[bytes, int]: + """Unified decode — dispatches to the right codec decoder.""" + if audio_type == "snac": + if not token_ids: + raise ValueError("SNAC decoding requires token_ids") + return self.decode_snac(torch.tensor([token_ids], dtype = torch.long), device) + elif audio_type == "bicodec": + if not text: + raise ValueError("BiCodec decoding requires text") + return self.decode_bicodec(text, device) + elif audio_type == "dac": + if not text: + raise ValueError("DAC decoding requires text") + return self.decode_dac(text, device) + raise ValueError(f"Cannot decode audio_type: {audio_type}") + # ── Cleanup ────────────────────────────────────────────────── def unload(self) -> None: diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 42cce3c4d8..439e142141 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -10,6 +10,7 @@ through its OpenAI-compatible /v1/chat/completions endpoint. import atexit import json +import struct import structlog from loggers import get_logger import shutil @@ -45,6 +46,8 @@ class LlamaCppBackend: self._hf_variant: Optional[str] = None self._is_vision: bool = False self._healthy = False + self._context_length: Optional[int] = None + self._chat_template: Optional[str] = None self._lock = threading.Lock() self._stdout_lines: list[str] = [] self._stdout_thread: Optional[threading.Thread] = None @@ -80,6 +83,14 @@ class LlamaCppBackend: def hf_variant(self) -> Optional[str]: return self._hf_variant + @property + def context_length(self) -> Optional[int]: + return self._context_length + + @property + def chat_template(self) -> Optional[str]: + return self._chat_template + # ── Binary discovery ────────────────────────────────────────── @staticmethod @@ -371,6 +382,99 @@ class LlamaCppBackend: # Pipe closed — process is terminating pass + # GGUF KV type sizes for fast skipping + _GGUF_TYPE_SIZE = { + 0: 1, + 1: 1, + 2: 2, + 3: 2, + 4: 4, + 5: 4, + 6: 4, + 7: 1, + 10: 8, + 11: 8, + 12: 8, + } + + @staticmethod + def _gguf_skip_value(f, vtype: int) -> None: + """Skip a GGUF KV value without reading it.""" + sz = LlamaCppBackend._GGUF_TYPE_SIZE.get(vtype) + if sz is not None: + f.seek(sz, 1) + elif vtype == 8: # STRING + slen = struct.unpack(" None: + """Read context_length and chat_template from a GGUF file's KV header. + + Parses only the KV pairs we need (~30ms even for multi-GB files). + For split GGUFs, metadata is always in shard 1. + """ + try: + WANTED = {"general.architecture", "tokenizer.chat_template"} + arch = None + ctx_key = None + + with open(gguf_path, "rb") as f: + magic = struct.unpack(" Optional[str]: + """Detect audio/TTS codec by probing the loaded model's vocabulary.""" + if not self.is_loaded: + return None + try: + with httpx.Client(timeout = 10) as client: + + def _detok(tid: int) -> str: + r = client.post( + f"{self.base_url}/detokenize", json = {"tokens": [tid]} + ) + return r.json().get("content", "") if r.status_code == 200 else "" + + def _tok(text: str) -> list[int]: + r = client.post( + f"{self.base_url}/tokenize", + json = {"content": text, "add_special": False}, + ) + return r.json().get("tokens", []) if r.status_code == 200 else [] + + # Check codec-specific tokens (not generic ones that may exist in non-audio models) + if "")) == 1 and len(_tok("<|audio_eos|>")) == 1: + return "csm" + if len(_tok("<|startoftranscript|>")) == 1: + return "whisper" + if ( + len(_tok("<|bicodec_semantic_0|>")) == 1 + and len(_tok("<|bicodec_global_0|>")) == 1 + ): + return "bicodec" + if len(_tok("<|c1_0|>")) == 1 and len(_tok("<|c2_0|>")) == 1: + return "dac" + except Exception as e: + logger.debug(f"Audio type detection failed: {e}") + return None + + # Prompt format per codec: (template, stop_tokens, needs_token_ids) + # Matches prompts in InferenceBackend._generate_snac/bicodec/dac + _TTS_PROMPTS = { + "snac": ( + "{text}<|eot_id|>", + [""], + True, + ), + "bicodec": ( + "<|task_tts|><|start_content|>{text}<|end_content|><|start_global_token|>", + ["<|im_end|>", ""], + False, + ), + "dac": ( + "<|im_start|>\n<|text_start|>{text}<|text_end|>\n<|audio_start|><|global_features_start|>\n", + ["<|im_end|>", "<|audio_end|>"], + False, + ), + } + + _codec_mgr = None # Shared AudioCodecManager instance + + def init_audio_codec(self, audio_type: str) -> None: + """Load the audio codec at model load time (mirrors non-GGUF path).""" + import torch + from core.inference.audio_codecs import AudioCodecManager + + if LlamaCppBackend._codec_mgr is None: + LlamaCppBackend._codec_mgr = AudioCodecManager() + + device = "cuda" if torch.cuda.is_available() else "cpu" + model_repo_path = None + + # BiCodec needs a repo with BiCodec/ weights — download canonical SparkTTS + if audio_type == "bicodec": + from huggingface_hub import snapshot_download + import os + + repo_path = snapshot_download( + "unsloth/Spark-TTS-0.5B", local_dir = "Spark-TTS-0.5B" + ) + model_repo_path = os.path.abspath(repo_path) + + LlamaCppBackend._codec_mgr.load_codec( + audio_type, device, model_repo_path = model_repo_path + ) + logger.info(f"Loaded audio codec for GGUF TTS: {audio_type}") + + def generate_audio_response( + self, + text: str, + audio_type: str, + temperature: float = 0.6, + top_p: float = 0.95, + top_k: int = 50, + min_p: float = 0.0, + max_new_tokens: int = 2048, + repetition_penalty: float = 1.1, + ) -> tuple: + """ + Generate TTS audio via llama-server /completion + codec decoding. + Returns (wav_bytes, sample_rate). + """ + if audio_type not in self._TTS_PROMPTS: + raise RuntimeError(f"GGUF TTS does not support '{audio_type}' codec.") + + tpl, stop, need_ids = self._TTS_PROMPTS[audio_type] + + payload: dict = { + "prompt": tpl.format(text = text), + "stream": False, + "n_predict": max_new_tokens, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k if top_k >= 0 else 0, + "min_p": min_p, + "repeat_penalty": repetition_penalty, + } + if stop: + payload["stop"] = stop + if need_ids: + payload["n_probs"] = 1 + + with httpx.Client(timeout = httpx.Timeout(300, connect = 10)) as client: + resp = client.post(f"{self.base_url}/completion", json = payload) + if resp.status_code != 200: + raise RuntimeError( + f"llama-server returned {resp.status_code}: {resp.text}" + ) + + data = resp.json() + token_ids = ( + [p["id"] for p in data.get("completion_probabilities", []) if "id" in p] + if need_ids + else None + ) + + import torch + + device = "cuda" if torch.cuda.is_available() else "cpu" + return LlamaCppBackend._codec_mgr.decode( + audio_type, device, token_ids = token_ids, text = data.get("content", "") + ) diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 620304f37a..f45ec5890d 100644 --- a/studio/backend/models/inference.py +++ b/studio/backend/models/inference.py @@ -119,6 +119,9 @@ class LoadResponse(BaseModel): inference: dict = Field( ..., description = "Inference parameters (temperature, top_p, top_k, min_p)" ) + context_length: Optional[int] = Field( + None, description = "Model's native context length (from GGUF metadata)" + ) class UnloadResponse(BaseModel): diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 55bd9710a5..60936cc2eb 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -155,6 +155,17 @@ async def load_model( logger.info(f"Loaded GGUF model via llama-server: {config.identifier}") + # Detect TTS audio by probing the loaded model's vocabulary + from utils.models import is_audio_input_type + + _gguf_audio = llama_backend.detect_audio_type() + _gguf_is_audio = _gguf_audio in ("snac", "bicodec", "dac") + llama_backend._is_audio = _gguf_is_audio + llama_backend._audio_type = _gguf_audio + if _gguf_is_audio: + logger.info(f"GGUF model detected as audio: audio_type={_gguf_audio}") + await asyncio.to_thread(llama_backend.init_audio_codec, _gguf_audio) + inference_config = load_inference_config(config.identifier) return LoadResponse( @@ -164,7 +175,11 @@ async def load_model( is_vision = config.is_vision, is_lora = False, is_gguf = True, + is_audio = _gguf_is_audio, + audio_type = _gguf_audio, + has_audio_input = is_audio_input_type(_gguf_audio), inference = inference_config, + context_length = llama_backend.context_length, ) # ── Standard path: load via Unsloth/transformers ────────── @@ -473,6 +488,8 @@ async def get_status( is_vision = llama_backend.is_vision, is_gguf = True, gguf_variant = llama_backend.hf_variant, + is_audio = getattr(llama_backend, "_is_audio", False), + audio_type = getattr(llama_backend, "_audio_type", None), loading = [], loaded = [llama_backend.model_identifier], ) @@ -521,78 +538,84 @@ async def generate_audio( """ Generate audio (TTS) from the latest user message. Returns a JSON response with base64-encoded WAV audio. - Only works when an audio model is loaded. + Works with both GGUF (llama-server) and Unsloth/transformers backends. """ import base64 - backend = get_inference_backend() - if not backend.active_model_name: - raise HTTPException(status_code = 400, detail = "No model loaded.") - - model_info = backend.models.get(backend.active_model_name, {}) - if not model_info.get("is_audio"): - raise HTTPException( - status_code = 400, detail = "Active model is not an audio model." - ) - # Extract text from the last user message _, chat_messages, _ = _extract_content_parts(payload.messages) if not chat_messages: raise HTTPException(status_code = 400, detail = "No messages provided.") - last_user_msg = next( (m for m in reversed(chat_messages) if m["role"] == "user"), None ) if not last_user_msg: raise HTTPException(status_code = 400, detail = "No user message found.") - text = last_user_msg["content"] + # Pick backend — both return (wav_bytes, sample_rate) + llama_backend = get_llama_cpp_backend() + if llama_backend.is_loaded and getattr(llama_backend, "_is_audio", False): + model_name = llama_backend.model_identifier + gen = lambda: llama_backend.generate_audio_response( + text = text, + audio_type = llama_backend._audio_type, + temperature = payload.temperature, + top_p = payload.top_p, + top_k = payload.top_k, + min_p = payload.min_p, + max_new_tokens = payload.max_tokens or 2048, + repetition_penalty = payload.repetition_penalty, + ) + else: + backend = get_inference_backend() + if not backend.active_model_name: + raise HTTPException(status_code = 400, detail = "No model loaded.") + model_info = backend.models.get(backend.active_model_name, {}) + if not model_info.get("is_audio"): + raise HTTPException( + status_code = 400, detail = "Active model is not an audio model." + ) + model_name = backend.active_model_name + gen = lambda: backend.generate_audio_response( + text = text, + temperature = payload.temperature, + top_p = payload.top_p, + top_k = payload.top_k, + min_p = payload.min_p, + max_new_tokens = payload.max_tokens or 2048, + repetition_penalty = payload.repetition_penalty, + use_adapter = payload.use_adapter, + ) + try: wav_bytes, sample_rate = await asyncio.get_event_loop().run_in_executor( - None, - lambda: backend.generate_audio_response( - text = text, - temperature = payload.temperature, - top_p = payload.top_p, - top_k = payload.top_k, - min_p = payload.min_p, - max_new_tokens = payload.max_tokens or 2048, - repetition_penalty = payload.repetition_penalty, - use_adapter = payload.use_adapter, - ), + None, gen ) - - audio_b64 = base64.b64encode(wav_bytes).decode("ascii") - completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" - - return JSONResponse( - content = { - "id": completion_id, - "object": "chat.completion.audio", - "model": backend.active_model_name, - "audio": { - "data": audio_b64, - "format": "wav", - "sample_rate": sample_rate, - }, - "choices": [ - { - "index": 0, - "message": { - "role": "assistant", - "content": f'[Generated audio from: "{text[:100]}"]', - }, - "finish_reason": "stop", - } - ], - } - ) - except Exception as e: logger.error(f"Audio generation error: {e}", exc_info = True) raise HTTPException(status_code = 500, detail = str(e)) + audio_b64 = base64.b64encode(wav_bytes).decode("ascii") + return JSONResponse( + content = { + "id": f"chatcmpl-{uuid.uuid4().hex[:12]}", + "object": "chat.completion.audio", + "model": model_name, + "audio": {"data": audio_b64, "format": "wav", "sample_rate": sample_rate}, + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": f'[Generated audio from: "{text[:100]}"]', + }, + "finish_reason": "stop", + } + ], + } + ) + # ===================================================================== # OpenAI-Compatible Chat Completions (/chat/completions) @@ -714,6 +737,8 @@ async def openai_chat_completions( # ── Determine which backend is active ───────────────────── if using_gguf: model_name = llama_backend.model_identifier or payload.model + if getattr(llama_backend, "_is_audio", False): + return await generate_audio(payload, request) else: backend = get_inference_backend() if not backend.active_model_name: diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 868f3aaf33..fa0daea7fd 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -304,6 +304,22 @@ async def list_models( ) loaded_models.append(model_info) + # Include active GGUF model (loaded via llama-server) + from routes.inference import get_llama_cpp_backend + + llama_backend = get_llama_cpp_backend() + if llama_backend.is_loaded and llama_backend.model_identifier: + loaded_models.append( + ModelDetails( + id = llama_backend.model_identifier, + name = llama_backend.model_identifier.split("/")[-1], + is_gguf = True, + is_vision = llama_backend.is_vision, + is_audio = getattr(llama_backend, "_is_audio", False), + audio_type = getattr(llama_backend, "_audio_type", None), + ) + ) + # Combine default and loaded models all_models = [] seen_ids = set() diff --git a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx index c8f8881808..2284cd194f 100644 --- a/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx +++ b/studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx @@ -340,6 +340,14 @@ function isGgufRepo(id: string): boolean { return id.toUpperCase().includes("-GGUF"); } +/** Extract param count label from model name (e.g. "Qwen3-0.6B" -> "0.6B"). */ +function extractParamLabel(id: string): string | undefined { + // Match patterns like "0.6B", "1B", "4B", "3.5B", "70B", "1.5B" etc. + const name = id.split("/").pop() ?? id; + const match = name.match(/(?:^|[-_])(\d+(?:\.\d+)?)[Bb](?:[-_]|$)/); + return match ? `${match[1]}B` : undefined; +} + // Module-level caches so re-mounting the popover shows results instantly let _cachedGgufCache: CachedGgufRepo[] = []; let _cachedModelsCache: CachedModelRepo[] = []; @@ -558,7 +566,7 @@ export function HubModelPicker({ meta={ isGgufRepo(id) ? "GGUF" - : vram?.detail ?? undefined + : vram?.detail ?? extractParamLabel(id) } selected={value === id} onClick={() => handleModelClick(id)} @@ -593,7 +601,7 @@ export function HubModelPicker({ meta={ isGgufRepo(id) ? "GGUF" - : metricsById.get(id) + : metricsById.get(id) ?? extractParamLabel(id) } selected={value === id} onClick={() => handleModelClick(id)} diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 30a8d97461..6833babc32 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -86,10 +86,17 @@ function toOpenAIMessage(message: RunMessage): { return null; } - return { - role: message.role, - content: collectTextParts(message).join("\n"), - }; + let content = collectTextParts(message).join("\n"); + // Strip inline audio base64 from prior assistant messages to avoid + // inflating token counts (e.g. audio-player responses with embedded WAV). + if (message.role === "assistant") { + content = content.replace( + /data:audio\/[a-z0-9.+-]+;base64,[A-Za-z0-9+/=]+/g, + "[audio]", + ); + } + + return { role: message.role, content }; } function extractImageBase64(input: string): string | undefined { @@ -194,7 +201,8 @@ function waitForModelReady(abortSignal?: AbortSignal): Promise { async function autoLoadSmallestModel(): Promise { const toastId = toast("Loading a model…", { description: "Auto-selecting the smallest downloaded model.", - duration: Infinity, + duration: 5000, + closeButton: true, }); try { const [ggufRepos, modelRepos] = await Promise.all([ @@ -214,7 +222,7 @@ async function autoLoadSmallestModel(): Promise { .sort((a, b) => a.size_bytes - b.size_bytes); if (downloaded.length > 0) { const variant = downloaded[0]; - await loadModel({ + const loadResp = await loadModel({ model_path: repo.repo_id, hf_token: null, max_seq_length: 4096, @@ -223,7 +231,9 @@ async function autoLoadSmallestModel(): Promise { gguf_variant: variant.quant, trust_remote_code: false, }); - useChatRuntimeStore.getState().setCheckpoint(repo.repo_id, variant.quant); + const store = useChatRuntimeStore.getState(); + store.setCheckpoint(repo.repo_id, variant.quant); + store.setParams({ ...store.params, maxTokens: loadResp.context_length ?? 131072 }); toast.success(`Loaded ${repo.repo_id} (${variant.quant})`, { id: toastId }); return true; } @@ -247,7 +257,9 @@ async function autoLoadSmallestModel(): Promise { gguf_variant: null, trust_remote_code: false, }); - useChatRuntimeStore.getState().setCheckpoint(repo.repo_id); + const store = useChatRuntimeStore.getState(); + store.setCheckpoint(repo.repo_id); + store.setParams({ ...store.params, maxTokens: 4096 }); toast.success(`Loaded ${repo.repo_id}`, { id: toastId }); return true; } catch { diff --git a/studio/frontend/src/features/chat/chat-settings-sheet.tsx b/studio/frontend/src/features/chat/chat-settings-sheet.tsx index bb2b05f9da..302bc17335 100644 --- a/studio/frontend/src/features/chat/chat-settings-sheet.tsx +++ b/studio/frontend/src/features/chat/chat-settings-sheet.tsx @@ -340,10 +340,10 @@ export function ChatSettingsPanel({ label="Max Tokens" value={params.maxTokens} min={64} - max={isGguf ? 131072 : 32768} - step={64} + max={isGguf ? params.maxTokens : 32768} + step={isGguf ? params.maxTokens : 64} onChange={set("maxTokens")} - displayValue={isGguf && params.maxTokens >= 131072 ? "Max" : undefined} + displayValue={isGguf ? "Max" : undefined} /> 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 47bd0698a8..b2ec407d4c 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 @@ -132,9 +132,11 @@ function mergeRecommendedInference( modelId: string, ): InferenceParams { const inference = response.inference; - // GGUF: max tokens = 131072 (effectively unlimited, model decides) - // Non-GGUF: max tokens = 4096 - const defaultMaxTokens = response.is_gguf ? 131072 : 4096; + // GGUF: use actual context length from GGUF metadata, fallback to 131072 + // Non-GGUF: 4096 + const defaultMaxTokens = response.is_gguf + ? (response.context_length ?? 131072) + : 4096; return { ...current, checkpoint: modelId, diff --git a/studio/frontend/src/features/chat/types/api.ts b/studio/frontend/src/features/chat/types/api.ts index 4046deb4c2..54ce25c1e2 100644 --- a/studio/frontend/src/features/chat/types/api.ts +++ b/studio/frontend/src/features/chat/types/api.ts @@ -82,6 +82,7 @@ export interface LoadModelResponse { min_p?: number; trust_remote_code?: boolean; }; + context_length?: number | null; } export interface UnloadModelRequest {