Compare commits

...
Sign in to create a new pull request.

5 commits

Author SHA1 Message Date
Daniel Han
7b33cd70ba Remove package-lock.json from tracking 2026-03-16 11:44:23 +00:00
pre-commit-ci[bot]
4c29c5e9d3 [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
2026-03-16 11:43:28 +00:00
Manan Shah
cbb4929139 studio: GGUF TTS audio support (from PR #4318)
Add GGUF TTS audio generation via llama-server. When a GGUF model
loads, the backend probes its vocabulary to detect audio codecs
(SNAC/BiCodec/DAC/CSM/Whisper). If detected, the codec is pre-loaded
and the model is reported as audio to the frontend.

During chat, TTS models route to the audio generation path which sends
a per-codec prompt to llama-server's /completion endpoint, extracts
generated tokens/text, and decodes to WAV using AudioCodecManager.

Also strips base64 audio data from prior assistant messages to prevent
context overflow.

Co-authored-by: Manan Shah <mananshah511@gmail.com>
2026-03-16 11:43:09 +00:00
Daniel Han
58523dc4a9 studio: read GGUF context_length via fast header parser, set max tokens
- Fast GGUF metadata reader (~30-55ms) parses only KV header, skips
  tensor data and large arrays (tokenizer vocab etc)
- Extracts context_length and chat_template from GGUF metadata
- Returns context_length in LoadResponse for frontend to use
- Frontend sets maxTokens to actual context_length for GGUFs (e.g.
  262144 for Qwen3.5-9B, 131072 for Qwen2.5-7B)
- Max Tokens slider shows "Max" and is locked for GGUFs
- Auto-load path also uses actual context_length from load response
- Toast auto-dismiss (5s) and close button for auto-load toast
2026-03-16 11:41:09 +00:00
Daniel Han
51321be31a studio: extract param count from model name as fallback
When HuggingFace API doesn't return totalParams for a model,
extract the param count from the model name (e.g. "Qwen3-0.6B"
-> "0.6B", "Llama-3.2-1B-Instruct" -> "1B"). Applied to both
the recommended list and HF search results.
2026-03-16 09:54:46 +00:00
10 changed files with 421 additions and 67 deletions

View file

@ -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:

View file

@ -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("<Q", f.read(8))[0]
f.seek(slen, 1)
elif vtype == 9: # ARRAY
atype = struct.unpack("<I", f.read(4))[0]
alen = struct.unpack("<Q", f.read(8))[0]
elem_sz = LlamaCppBackend._GGUF_TYPE_SIZE.get(atype)
if elem_sz is not None:
f.seek(elem_sz * alen, 1)
elif atype == 8:
for _ in range(alen):
slen = struct.unpack("<Q", f.read(8))[0]
f.seek(slen, 1)
else:
for _ in range(alen):
LlamaCppBackend._gguf_skip_value(f, atype)
def _read_gguf_metadata(self, gguf_path: str) -> 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("<I", f.read(4))[0]
if magic != 0x46554747: # b"GGUF" as little-endian u32
return
_version = struct.unpack("<I", f.read(4))[0]
_tensor_count, kv_count = struct.unpack("<QQ", f.read(16))
for _ in range(kv_count):
key_len = struct.unpack("<Q", f.read(8))[0]
key = f.read(key_len).decode("utf-8")
vtype = struct.unpack("<I", f.read(4))[0]
if key in WANTED or (ctx_key and key == ctx_key):
# Read this value
if vtype == 8: # STRING
slen = struct.unpack("<Q", f.read(8))[0]
val_s = f.read(slen).decode("utf-8")
if key == "general.architecture":
arch = val_s
ctx_key = f"{arch}.context_length"
elif key == "tokenizer.chat_template":
self._chat_template = val_s
elif vtype == 4: # UINT32
val_i = struct.unpack("<I", f.read(4))[0]
if ctx_key and key == ctx_key:
self._context_length = val_i
elif vtype == 10: # UINT64
val_i = struct.unpack("<Q", f.read(8))[0]
if ctx_key and key == ctx_key:
self._context_length = val_i
else:
self._gguf_skip_value(f, vtype)
else:
self._gguf_skip_value(f, vtype)
if self._context_length:
logger.info(f"GGUF metadata: context_length={self._context_length}")
if self._chat_template:
logger.info(
f"GGUF metadata: chat_template={len(self._chat_template)} chars"
)
except Exception as e:
logger.warning(f"Failed to read GGUF metadata: {e}")
# ── HF download (no lock held) ───────────────────────────────
def _download_gguf(
@ -642,6 +746,9 @@ class LlamaCppBackend:
else:
raise ValueError("Either gguf_path or hf_repo must be provided")
# Read GGUF metadata (context_length, chat_template) -- fast, header only
self._read_gguf_metadata(model_path)
# Check cancel after download
if self._cancel_event.is_set():
logger.info("Load cancelled after download phase")
@ -779,8 +886,20 @@ class LlamaCppBackend:
self._hf_repo = None
self._hf_variant = None
self._is_vision = False
self._is_audio = False
self._audio_type = None
self._port = None
self._healthy = False
self._context_length = None
self._chat_template = None
# Free audio codec GPU memory
if LlamaCppBackend._codec_mgr is not None:
LlamaCppBackend._codec_mgr.unload()
LlamaCppBackend._codec_mgr = None
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
return True
def _kill_process(self):
@ -1028,3 +1147,149 @@ class LlamaCppBackend:
if cancel_event is not None and cancel_event.is_set():
return
raise
# ── TTS support ────────────────────────────────────────────
def detect_audio_type(self) -> 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 "<custom_token_" in _detok(128258) and "<custom_token_" in _detok(
128259
):
return "snac"
if len(_tok("<|AUDIO|>")) == 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": (
"<custom_token_3>{text}<|eot_id|><custom_token_4>",
["<custom_token_2>"],
True,
),
"bicodec": (
"<|task_tts|><|start_content|>{text}<|end_content|><|start_global_token|>",
["<|im_end|>", "</s>"],
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", "")
)

View file

@ -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):

View file

@ -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:

View file

@ -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()

View file

@ -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)}

View file

@ -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<void> {
async function autoLoadSmallestModel(): Promise<boolean> {
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<boolean> {
.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<boolean> {
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<boolean> {
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 {

View file

@ -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}
/>
</div>
</CollapsibleSection>

View file

@ -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,

View file

@ -82,6 +82,7 @@ export interface LoadModelResponse {
min_p?: number;
trust_remote_code?: boolean;
};
context_length?: number | null;
}
export interface UnloadModelRequest {