ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

code cleanup

Browse source
This commit is contained in:
Manan17 2026-03-01 08:04:38 +00:00
commit 8cdeb006b6
14 changed files with 277 additions and 275 deletions
Download patch file Download diff file Expand all files Collapse all files

71
studio/TESTING.md Normal file
View file

@ -0,0 +1,71 @@
# Testing Matrix: Audio, Text & VLM
## Text Model Tests
| # | Test | Model | Steps | Expected |
|---|------|-------|-------|----------|
| 1 | **Text inference — basic chat** | `unsloth/Qwen3-4B` or any text model | Load model → send "Hello, how are you?" | Streaming text response, no errors |
| 2 | **Text inference — system prompt** | Any text model | Set system prompt to "You are a pirate" → send "Tell me about the ocean" | Response in pirate persona |
| 3 | **Text training — LoRA** | Any text model | Load model → pick Alpaca dataset → set max_steps=10 → start training | Training completes, checkpoint saved, loss decreases |
| 4 | **Text compare mode** | Any text model with trained LoRA | Open compare view → send message | Both Base and LoRA columns respond, responses differ |
## VLM (Vision) Tests
| # | Test | Model | Steps | Expected |
|---|------|-------|-------|----------|
| 5 | **VLM inference — image description** | `unsloth/Llama-3.2-11B-Vision` or Gemma-3 vision | Load → attach image via paperclip → "What's in this image?" | Describes the image content accurately |
| 6 | **VLM inference — text only (no image)** | Same VLM | Send text message without image | Normal text response (no crash) |
| 7 | **VLM training — vision LoRA** | Any VLM | Load → pick image-text dataset → check finetune_vision_layers is ON → train max_steps=10 | Training completes with vision+language LoRA |
| 8 | **VLM compare mode with image** | VLM with trained LoRA | Open compare view → upload image → send | Both panels describe the image, LoRA panel should differ |
| 9 | **VLM dataset mapping** | Any VLM | Pick a dataset that needs manual column mapping | Mapping card shows correctly, VLM-specific labels appear |
## Audio TTS (Text-to-Speech) Tests
| # | Test | Model | Steps | Expected |
|---|------|-------|-------|----------|
| 10 | **TTS inference — SNAC/Orpheus** | `canopylabs/orpheus-3b-0.1-ft` | Load → send text message | AudioPlayer renders with playable WAV audio |
| 11 | **TTS inference — SparkTTS (BiCodec)** | `SparkAudio/Spark-TTS-0.5B` | Load → send text "Hello world" | AudioPlayer with synthesized speech (not raw bicodec tokens) |
| 12 | **TTS inference — OuteTTS (DAC)** | `OuteAI/Llama-OuteTTS-1.0-1B` | Load → send text | AudioPlayer with audio output |
| 13 | **TTS inference — CSM/Sesame** | `sesame/csm-1b` | Load → send text | AudioPlayer with audio output |
| 14 | **TTS training — LoRA** | Any TTS model (e.g., Orpheus) | Load → pick audio dataset → train max_steps=10 | Training completes, checkpoint saved |
| 15 | **TTS compare mode** | TTS with LoRA adapter | Open compare → send text | Both panels show AudioPlayer, LoRA should sound different |
| 16 | **SparkTTS LoRA inference** | SparkTTS with trained adapter | Load LoRA checkpoint → send text | Plays audio without 404 error (bicodec loads from local path) |
## Audio ASR (Speech-to-Text) Tests
| # | Test | Model | Steps | Expected |
|---|------|-------|-------|----------|
| 17 | **Whisper inference — audio transcription** | `unsloth/whisper-large-v3` | Load → upload audio via headphones button → send | Transcribed text appears (no "only accepts audio" error) |
| 18 | **Whisper inference — no audio error** | `unsloth/whisper-large-v3` | Load → send text without audio | Clear error: "Whisper models require audio input" |
| 19 | **Gemma 3n inference — audio ASR** | `unsloth/gemma-3n-E4B-it` | Load → upload audio → "Transcribe this audio" | Accurate transcription, uses greedy decoding |
| 20 | **Gemma 3n inference — text only** | `unsloth/gemma-3n-E4B-it` | Load → send text without audio | Normal text response (model is also a text/vision model) |
| 21 | **Gemma 3n inference — image** | `unsloth/gemma-3n-E4B-it` | Load → attach image → "What's in this image?" | Describes image (Gemma 3n supports vision too) |
| 22 | **Gemma 3n training — audio dataset** | `unsloth/gemma-3n-E4B-it` | Load → pick speech dataset → train max_steps=10 | Training completes, audio mapping card shows "audio and text" |
| 23 | **Audio chip in user message** | Any ASR model | Upload audio → send message | Audio filename chip appears in user message bubble |
## Cross-Cutting / Edge Case Tests
| # | Test | Model | Steps | Expected |
|---|------|-------|-------|----------|
| 24 | **Model switch — text to TTS** | Text model → TTS model | Load text model → chat → switch to TTS → chat | First gives text, second gives AudioPlayer — no leftover state |
| 25 | **Model switch — TTS to ASR** | TTS → Whisper or Gemma 3n | Load TTS → generate audio → switch to ASR → upload audio | TTS gives audio, ASR gives text — clean transition |
| 26 | **Model switch — VLM to text** | VLM → text model | Load VLM → send image → switch to text model → send text | No vision errors on text model, image attachment ignored |
| 27 | **Abort mid-generation** | Any streaming model | Send message → click stop button mid-stream | Generation stops cleanly, partial response visible, no crash |
| 28 | **Large audio file rejection** | Any ASR model | Try uploading audio > 50MB | Upload rejected (MAX_AUDIO_SIZE), no crash |
| 29 | **No model loaded error** | No model | Open chat → send message | Toast: "No model loaded — Pick model in top bar, then retry" |
| 30 | **Training then inference** | Any model | Train LoRA → load checkpoint → chat | Trained checkpoint responds (different from base) |
## Quick Smoke Test Order (prioritized)
If running a fast subset, do these 10 in order:
1. **#1** — Text basic chat (sanity check)
2. **#5** — VLM image description
3. **#10** — TTS audio generation (SNAC)
4. **#19** — Gemma 3n audio ASR
5. **#17** — Whisper transcription
6. **#4** — Text compare mode
7. **#3** — Text LoRA training
8. **#23** — Audio chip in user message
9. **#24** — Model switch text→TTS
10. **#29** — No model loaded error

1
studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml
View file

@ -3,6 +3,7 @@
# Also applies to: unsloth/whisper-large-v3, openai/whisper-large-v3
audio_type: whisper
audio_input: true
training:
eval_steps: 5

1
studio/backend/core/inference/OuteTTS

@ -1 +0,0 @@
Subproject commit 59d896747aa0a6a207837e7da2d6921805eae684

351
studio/backend/core/inference/inference.py
View file

@ -118,22 +118,45 @@ class InferenceBackend:
elif audio_type == "bicodec":
import os
from unsloth import FastModel
from huggingface_hub import snapshot_download
# Spark-TTS: download full repo, then load from /LLM subfolder
hf_repo = config.path
local_dir = hf_repo.split("/")[-1]
repo_path = snapshot_download(hf_repo, local_dir=local_dir)
abs_repo_path = os.path.abspath(repo_path)
llm_path = os.path.join(abs_repo_path, "LLM")
logger.info(f"Spark-TTS: downloaded repo to {repo_path}, loading LLM from {llm_path}")
if config.is_lora and config.base_model:
# LoRA adapter: load from local adapter path.
# base_model is e.g. /home/.../Spark-TTS-0.5B/LLM
# The BiCodec weights are in the parent dir (Spark-TTS-0.5B/).
base_path = config.base_model
if os.path.isdir(base_path):
abs_repo_path = os.path.abspath(os.path.dirname(base_path))
else:
# base_model is an HF ID — download it
from huggingface_hub import snapshot_download
local_dir = base_path.split("/")[-1]
repo_path = snapshot_download(base_path, local_dir=local_dir)
abs_repo_path = os.path.abspath(repo_path)
logger.info(f"Spark-TTS LoRA: loading adapter from {config.path}, BiCodec from {abs_repo_path}")
model, tokenizer = FastModel.from_pretrained(
config.path,
dtype=torch.float32,
load_in_4bit=False,
token=hf_token if hf_token and hf_token.strip() else None,
)
else:
# Base model: download full HF repo, then load from /LLM subfolder
from huggingface_hub import snapshot_download
hf_repo = config.path
local_dir = hf_repo.split("/")[-1]
repo_path = snapshot_download(hf_repo, local_dir=local_dir)
abs_repo_path = os.path.abspath(repo_path)
llm_path = os.path.join(abs_repo_path, "LLM")
logger.info(f"Spark-TTS: downloaded repo to {repo_path}, loading LLM from {llm_path}")
model, tokenizer = FastModel.from_pretrained(
llm_path,
dtype=torch.float32,
load_in_4bit=False,
token=hf_token if hf_token and hf_token.strip() else None,
)
model, tokenizer = FastModel.from_pretrained(
llm_path,
dtype=torch.float32,
load_in_4bit=False,
token=hf_token if hf_token and hf_token.strip() else None,
)
FastModel.for_inference(model)
self.models[model_name]["model"] = model
self.models[model_name]["tokenizer"] = tokenizer
@ -150,6 +173,35 @@ class InferenceBackend:
FastModel.for_inference(model)
self.models[model_name]["model"] = model
self.models[model_name]["tokenizer"] = tokenizer
elif audio_type == "whisper":
# Whisper ASR — uses FastModel with WhisperForConditionalGeneration
from unsloth import FastModel
from transformers import WhisperForConditionalGeneration
model, tokenizer = FastModel.from_pretrained(
config.path,
auto_model=WhisperForConditionalGeneration,
whisper_language="English",
whisper_task="transcribe",
load_in_4bit=False,
token=hf_token if hf_token and hf_token.strip() else None,
)
FastModel.for_inference(model)
model.eval()
# Create ASR pipeline (per notebook)
from transformers import pipeline as hf_pipeline
whisper_pipe = hf_pipeline(
"automatic-speech-recognition",
model=model,
tokenizer=tokenizer.tokenizer,
feature_extractor=tokenizer.feature_extractor,
processor=tokenizer,
return_language=True,
torch_dtype=torch.float16,
)
self.models[model_name]["model"] = model
self.models[model_name]["tokenizer"] = tokenizer
self.models[model_name]["whisper_pipeline"] = whisper_pipe
else:
# SNAC (Orpheus) uses FastLanguageModel
model, tokenizer = FastLanguageModel.from_pretrained(
@ -162,9 +214,10 @@ class InferenceBackend:
self.models[model_name]["model"] = model
self.models[model_name]["tokenizer"] = tokenizer
# Load the external codec for this audio type
model_repo_path = self.models[model_name].get("model_repo_path")
self._audio_codec_manager.load_codec(audio_type, self.device, model_repo_path=model_repo_path)
# Load the external codec for TTS audio types (Whisper is ASR, no codec needed)
if audio_type != "whisper":
model_repo_path = self.models[model_name].get("model_repo_path")
self._audio_codec_manager.load_codec(audio_type, self.device, model_repo_path=model_repo_path)
self.active_model_name = model_name
self.loading_models.discard(model_name)
@ -259,9 +312,7 @@ class InferenceBackend:
self.loading_models.discard(model_name)
raise Exception(error_msg)
pass
# Add this new function
def unload_model(self, model_name: str) -> bool:
"""
Completely removes a model from the registry and clears GPU memory.
@ -295,7 +346,6 @@ class InferenceBackend:
else:
logger.warning(f"Attempted to unload model '{model_name}', but it was not found in the registry.")
return True
pass
def revert_to_base_model(self, base_model_name: str) -> bool:
"""
@ -331,151 +381,6 @@ class InferenceBackend:
logger.error(traceback.format_exc())
return False
def activate_lora_adapter(self, base_model_name: str, lora_path: str) -> Tuple[bool, Optional[str]]:
"""
Activates a specific LoRA adapter on what is assumed to be a clean base model.
Uses PeftModel.from_pretrained() which correctly wraps the base model.
"""
model = self.models[base_model_name].get("model")
adapter_name_to_load = lora_path.split("/")[-1].replace(".", "_")
try:
# Use PeftModel.from_pretrained to wrap the clean base model with the adapter.
# This is the correct approach after model.unload() + del peft_config.
logger.info(f"Loading LoRA adapter '{adapter_name_to_load}' from '{lora_path}'...")
model = PeftModel.from_pretrained(model, lora_path, adapter_name=adapter_name_to_load)
self.models[base_model_name]["model"] = model
logger.info(f"LoRA adapter '{adapter_name_to_load}' activated successfully.")
return True, adapter_name_to_load
except Exception as e:
logger.error(f"Failed to activate LoRA adapter '{adapter_name_to_load}': {e}")
import traceback
logger.error(traceback.format_exc())
return False, None
def load_adapter(self, base_model_name: str, adapter_path: str, adapter_name: str = None) -> bool:
"""
Load a LoRA adapter onto the base model if it's not already registered.
This method is idempotent.
"""
if base_model_name not in self.models:
logger.error(f"Base model {base_model_name} not loaded")
return False
model = self.models[base_model_name].get("model")
if model is None:
logger.error(f"Model object for {base_model_name} is None.")
return False
if adapter_name is None:
adapter_name = adapter_path.split("/")[-1].replace(".", "_")
# If we've loaded this adapter before, we don't need to do anything.
if adapter_name in self.models[base_model_name].get("loaded_adapters", {}):
logger.info(f"Adapter '{adapter_name}' is already registered. Skipping.")
return True
try:
logger.info(f"Loading new adapter '{adapter_name}' from '{adapter_path}' onto {base_model_name}")
# Unsloth modifies the model in-place and returns None. Do NOT re-assign.
model.load_adapter(adapter_path, adapter_name=adapter_name)
# Update our internal registry so we don't load it again.
self.models[base_model_name]["loaded_adapters"][adapter_name] = adapter_path
total_adapters = len(getattr(model, 'peft_config', {}))
logger.info(f"Adapter '{adapter_name}' loaded successfully. (Total adapters on model: {total_adapters})")
return True
except Exception as e:
logger.error(f"Failed to load adapter '{adapter_name}': {e}")
import traceback
logger.error(traceback.format_exc())
return False
pass
def enable_adapter(self, base_model_name: str, adapter_name: str) -> bool:
"""Enable specific adapter (for generation)"""
if base_model_name not in self.models:
return False
model = self.models[base_model_name]["model"]
try:
logger.info(f"Enabling adapter: {adapter_name}")
model.set_adapter(adapter_name)
self.models[base_model_name]["active_adapter"] = adapter_name
return True
except Exception as e:
logger.error(f"Failed to enable adapter: {e}")
return False
def disable_adapters(self, base_model_name: str) -> bool:
"""Disable all adapters (back to pure base model)"""
if base_model_name not in self.models:
return False
model = self.models[base_model_name]["model"]
try:
logger.info(f"Disabling all adapters on {base_model_name}")
model.disable_adapters()
self.models[base_model_name]["active_adapter"] = None
return True
except Exception as e:
logger.error(f"Failed to disable adapters: {e}")
return False
# In backend/inference.py
def load_for_eval(self, lora_path: str, max_seq_length: int = 2048,
dtype = None, load_in_4bit: bool = True,
hf_token: Optional[str] = None) -> Tuple[bool, Optional[str], Optional[str]]:
"""
Prepare for eval: ensure base model and the specified adapter are loaded.
"""
try:
from utils.models import ModelConfig
lora_config = ModelConfig.from_lora_path(lora_path, hf_token)
if not lora_config:
return False, None, None
base_model_name = lora_config.base_model
# 1. Load the base model if it's not already in memory (this logic is correct)
if base_model_name not in self.models or not self.models[base_model_name].get("model"):
logger.info(f"Base model '{base_model_name}' not loaded, loading now.")
base_config = ModelConfig.from_ui_selection(base_model_name, None, is_lora=False)
if not self.load_model(base_config, max_seq_length, dtype, load_in_4bit, hf_token):
return False, None, None
else:
logger.info(f"Base model '{base_model_name}' is already in memory.")
self.active_model_name = base_model_name
# 2. Delegate to our now-idempotent load_adapter function.
# It will handle all cases: first adapter, or subsequent adapters.
adapter_name = lora_path.split("/")[-1].replace(".", "_")
adapter_success = self.load_adapter(
base_model_name=base_model_name,
adapter_path=lora_path,
adapter_name=adapter_name
)
if not adapter_success:
return False, base_model_name, None
return True, base_model_name, adapter_name
except Exception as e:
logger.error(f"Error during load_for_eval: {e}")
import traceback
logger.error(traceback.format_exc())
return False, None, None
pass
def load_for_eval(self, lora_path: str, max_seq_length: int = 2048,
dtype = None, load_in_4bit: bool = True,
hf_token: Optional[str] = None) -> Tuple[bool, Optional[str], Optional[str]]:
@ -522,7 +427,6 @@ class InferenceBackend:
import traceback
logger.error(traceback.format_exc())
return False, None, None
pass
def load_adapter(self, base_model_name: str, adapter_path: str, adapter_name: str) -> bool:
"""
@ -550,7 +454,6 @@ class InferenceBackend:
except Exception as e:
logger.error(f"Failed to load adapter '{adapter_name}': {e}")
return False
pass
def set_active_adapter(self, base_model_name: str, adapter_name: str) -> bool:
"""
@ -566,7 +469,6 @@ class InferenceBackend:
# This will catch the "adapter not found" error if something goes wrong.
logger.error(f"Failed to set active adapter to '{adapter_name}': {e}")
return False
pass
def _apply_adapter_state(self, use_adapter: Optional[Union[bool, str]]) -> None:
"""
@ -885,7 +787,6 @@ class InferenceBackend:
except Exception as e:
logger.error(f"Vision generation error: {e}")
yield f"Error: {str(e)}"
pass
def generate_audio_input_response(self, messages, system_prompt, audio_array,
temperature, top_p, top_k, min_p,
@ -903,25 +804,29 @@ class InferenceBackend:
processor = model_info.get("processor") or model_info.get("tokenizer")
raw_tokenizer = getattr(processor, "tokenizer", processor)
# Extract last user text
user_text = "Transcribe this audio."
# Extract last user text — default matches notebook prompt
user_text = "Please transcribe this audio."
if messages:
for msg in reversed(messages):
if msg["role"] == "user" and msg.get("content"):
user_text = msg["content"]
break
# Use ASR-specific system prompt if user hasn't set a custom one
if not system_prompt or system_prompt == "You are a helpful AI assistant.":
system_prompt = "You are an assistant that transcribes speech accurately."
# Build messages in Gemma 3n format — audio goes INTO apply_chat_template
audio_messages = []
if system_prompt:
audio_messages.append({"role": "system", "content": [{"type": "text", "text": system_prompt}]})
audio_messages.append({
"role": "user",
"content": [
{"type": "audio", "audio": audio_array},
{"type": "text", "text": user_text},
],
})
audio_messages = [
{"role": "system", "content": [{"type": "text", "text": system_prompt}]},
{
"role": "user",
"content": [
{"type": "audio", "audio": audio_array},
{"type": "text", "text": user_text},
],
},
]
# apply_chat_template handles audio embedding + tokenization in one step
inputs = processor.apply_chat_template(
@ -943,16 +848,13 @@ class InferenceBackend:
timeout=0.2,
)
# Notebook uses do_sample=False for ASR (greedy decoding for accuracy)
generation_kwargs = dict(
**inputs,
streamer=streamer,
max_new_tokens=max_new_tokens,
use_cache=True,
do_sample=temperature > 0,
temperature=temperature,
top_p=top_p,
top_k=top_k,
min_p=min_p,
do_sample=False,
)
err: dict[str, str] = {}
@ -1003,6 +905,28 @@ class InferenceBackend:
logger.error(f"Audio input generation error: {e}")
yield f"Error: {str(e)}"
def generate_whisper_response(self, audio_array, cancel_event=None) -> Generator[str, None, None]:
"""Whisper ASR — takes audio numpy array, yields transcribed text.
Uses the pre-built transformers pipeline (created during model loading).
"""
model_info = self.models[self.active_model_name]
whisper_pipe = model_info.get("whisper_pipeline")
if not whisper_pipe:
yield "Error: Whisper pipeline not initialized"
return
try:
with self._generation_lock:
result = whisper_pipe({"raw": audio_array, "sampling_rate": 16000})
text = result.get("text", "") if isinstance(result, dict) else str(result)
if text:
yield text
except Exception as e:
logger.error(f"Whisper ASR error: {e}")
yield f"Error: {str(e)}"
def generate_stream(self,
prompt: str,
temperature: float = 0.7,
@ -1124,9 +1048,6 @@ class InferenceBackend:
logger.error(f"Error during generation: {e}")
yield f"Error: {str(e)}"
# ... other helper methods (format_chat_prompt, _clean_generated_text, etc.)
pass
# ── Audio (TTS) Generation ────────────────────────────────────
def generate_audio_response(
@ -1507,7 +1428,6 @@ class InferenceBackend:
logger.debug(f"Reset generation state for model: {model_name}")
except Exception as e:
logger.warning(f"Could not fully reset model state for {model_name}: {e}")
pass
def reset_generation_state(self):
"""Reset any cached generation state to prevent hanging after errors"""
@ -1673,57 +1593,6 @@ class InferenceBackend:
return False
def load_model_simple(self,
model_path: str,
hf_token: Optional[str] = None,
max_seq_length: int = 2048,
load_in_4bit: bool = True) -> bool:
"""
Simple model loading wrapper for chat interface.
Accepts model path as string and handles ModelConfig creation internally.
Args:
model_path: Model name or path (e.g., "unsloth/llama-3-8b")
hf_token: HuggingFace token for gated models
max_seq_length: Maximum sequence length
load_in_4bit: Whether to use 4-bit quantization
Returns:
bool: True if successful, False otherwise
"""
try:
from backend.model_config import ModelConfig
logger.info(f"load_model_simple called with: {model_path}")
# Create config from string path
config = ModelConfig.from_ui_selection(
model_path,
lora_path=None, # No LoRA for chat
is_lora=False
)
logger.info(f"Created ModelConfig with identifier: {config.identifier}")
# Call existing load_model with config
return self.load_model(
config=config,
max_seq_length=max_seq_length,
dtype=None, # Auto-detect
load_in_4bit=load_in_4bit,
hf_token=hf_token
)
except Exception as e:
logger.error(f"Error in load_model_simple: {e}")
import traceback
traceback.print_exc()
return False
pass
# Global inference backend instance
inference_backend = InferenceBackend()

1
studio/backend/core/training/inference/OuteTTS

@ -1 +0,0 @@
Subproject commit 59d896747aa0a6a207837e7da2d6921805eae684

2
studio/backend/core/training/trainer.py
View file

@ -1048,7 +1048,7 @@ class UnslothTrainer:
return formatted
self._update_progress(status_message="Formatting audio VLM dataset...")
dataset = dataset.map(format_messages, batched=True, batch_size=4, num_proc=4)
dataset = dataset.map(format_messages, batched=True, batch_size=4, num_proc=safe_num_proc(4))
print(f"Audio VLM dataset formatted: {len(dataset)} examples\n")
return dataset

15
studio/backend/routes/inference.py
View file

@ -569,10 +569,18 @@ async def openai_chat_completions(
model_name = backend.active_model_name or payload.model
# ── Audio TTS path: auto-route to audio generation ────
# (Whisper is ASR not TTS — handled below in audio input path)
model_info = backend.models.get(backend.active_model_name, {})
if model_info.get("is_audio"):
if model_info.get("is_audio") and model_info.get("audio_type") != "whisper":
return await generate_audio(payload, request)
# ── Whisper without audio: return clear error ──
if model_info.get("audio_type") == "whisper" and not payload.audio_base64:
raise HTTPException(
status_code=400,
detail="Whisper models require audio input. Please upload an audio file.",
)
# ── Audio INPUT path: decode WAV and route to audio input generation ──
if payload.audio_base64 and model_info.get("has_audio_input"):
audio_array = _decode_audio_base64(payload.audio_base64)
@ -582,6 +590,11 @@ async def openai_chat_completions(
created = int(time.time())
def audio_input_generate():
if model_info.get("audio_type") == "whisper":
return backend.generate_whisper_response(
audio_array=audio_array,
cancel_event=cancel_event,
)
return backend.generate_audio_input_response(
messages=chat_messages,
system_prompt=system_prompt,

15
studio/backend/utils/hardware/hardware.py
View file

@ -423,13 +423,14 @@ def safe_num_proc(desired: Optional[int] = None) -> int:
"""
Return a safe ``num_proc`` for ``dataset.map()`` calls.
Fork-based multiprocessing deadlocks when CUDA has already been
initialized (e.g. after inference). This helper detects that case
and forces ``num_proc=1``.
On multi-GPU machines the NVIDIA driver spawns extra background threads,
making ``os.fork()`` prone to deadlocks when many workers are created.
This helper caps ``num_proc`` to 4 on such machines.
On single-GPU (or CPU-only) machines the original value is returned
unchanged.
Args:
desired: The num_proc you *want*. If None, auto-computes from
``os.cpu_count()``.
@ -442,6 +443,14 @@ def safe_num_proc(desired: Optional[int] = None) -> int:
if desired is None or not isinstance(desired, int):
desired = max(1, os.cpu_count() // 3)
# After inference, CUDA is initialized — forking will deadlock.
try:
import torch
if torch.cuda.is_initialized():
return 1
except ImportError:
pass
if get_physical_gpu_count() > 1:
capped = min(4, desired)
print(

25
studio/backend/utils/models/model_config.py
View file

@ -443,7 +443,7 @@ def is_audio_model(model_name: str) -> Optional[str]:
try:
defaults = load_model_defaults(model_name)
audio_type = defaults.get('audio_type')
if audio_type and isinstance(audio_type, str) and audio_type in ('snac', 'csm', 'bicodec', 'dac'):
if audio_type and isinstance(audio_type, str) and audio_type in ('snac', 'csm', 'bicodec', 'dac', 'whisper'):
logger.info(f"Model {model_name} detected as audio model: audio_type={audio_type}")
return audio_type
return None
@ -914,6 +914,23 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]:
logger.info(f"Loaded model defaults from {config_path} (via mapping)")
return config
# If model_name is a local path (e.g. /home/.../Spark-TTS-0.5B/LLM from
# adapter_config.json), try matching the last 1-2 path components against
# the registry (e.g. "Spark-TTS-0.5B/LLM").
if model_name not in _REVERSE_MODEL_MAPPING and (model_name.startswith("/") or model_name.startswith(".")):
parts = Path(model_name).parts
for depth in [2, 1]:
if len(parts) >= depth:
suffix = "/".join(parts[-depth:])
if suffix in _REVERSE_MODEL_MAPPING:
canonical_file = _REVERSE_MODEL_MAPPING[suffix]
for config_path in defaults_dir.rglob(canonical_file):
if config_path.is_file():
with open(config_path, 'r', encoding='utf-8') as f:
config = yaml.safe_load(f) or {}
logger.info(f"Loaded model defaults from {config_path} (via path suffix '{suffix}')")
return config
# Try exact model name match (for backward compatibility)
model_filename = model_name.replace("/", "_") + ".yaml"
# Search in subfolders and root
@ -1163,6 +1180,12 @@ class ModelConfig:
display_name = Path(path).name if is_local else identifier.split("/")[-1]
# Audio models are never vision models (e.g. WhisperForConditionalGeneration
# and CsmForConditionalGeneration match the ForConditionalGeneration suffix
# but are not VLMs).
if audio_type_val is not None:
vision = False
return cls(
identifier=identifier,
display_name=display_name,

32
studio/frontend/src/components/assistant-ui/thread.tsx
View file

@ -8,6 +8,8 @@ import { Reasoning, ReasoningGroup } from "@/components/assistant-ui/reasoning";
import { ToolFallback } from "@/components/assistant-ui/tool-fallback";
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
import { Button } from "@/components/ui/button";
import { sentAudioNames } from "@/features/chat/api/chat-adapter";
import { AUDIO_ACCEPT, MAX_AUDIO_SIZE, fileToBase64 } from "@/lib/audio-utils";
import { copyToClipboard } from "@/lib/copy-to-clipboard";
import { cn } from "@/lib/utils";
import {
@ -165,22 +167,6 @@ const ComposerAnimated: FC = () => {
);
};
const AUDIO_ACCEPT = "audio/wav,audio/mpeg,audio/webm,audio/ogg,audio/flac,audio/mp4";
const MAX_AUDIO_SIZE = 50 * 1024 * 1024;
function fileToBase64(file: File): Promise<string> {
return new Promise((resolve, reject) => {
const reader = new FileReader();
reader.onload = () => {
const result = reader.result as string;
const commaIndex = result.indexOf(",");
resolve(commaIndex >= 0 ? result.slice(commaIndex + 1) : result);
};
reader.onerror = () => reject(new Error("Failed to read file"));
reader.readAsDataURL(file);
});
}
const PendingAudioChip: FC = () => {
const audioName = useChatRuntimeStore((s) => s.pendingAudioName);
const clearPendingAudio = useChatRuntimeStore((s) => s.clearPendingAudio);
@ -438,6 +424,19 @@ const AssistantActionBar: FC = () => {
);
};
const UserMessageAudio: FC = () => {
const audioName = useAuiState(({ message }) => sentAudioNames.get(message.id));
if (!audioName) return null;
return (
<div className="col-start-2 flex justify-end">
<div className="flex items-center gap-2 rounded-lg border border-foreground/20 bg-muted px-3 py-1.5 text-xs">
<HeadphonesIcon className="size-3.5 text-muted-foreground" />
<span className="max-w-48 truncate">{audioName}</span>
</div>
</div>
);
};
const UserMessage: FC = () => {
return (
<MessagePrimitive.Root
@ -445,6 +444,7 @@ const UserMessage: FC = () => {
data-role="user"
>
<UserMessageAttachments />
<UserMessageAudio />
<div className="aui-user-message-content-wrapper relative col-start-2 min-w-0">
<div className="aui-user-message-content wrap-break-word rounded-2xl bg-muted px-4 py-2.5 text-foreground">

10
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -11,6 +11,9 @@ import {
type RunMessages = Parameters<ChatModelAdapter["run"]>[0]["messages"];
type RunMessage = RunMessages[number];
/** Tracks which user messages were sent with an audio file (messageId → filename). */
export const sentAudioNames = new Map<string, string>();
function collectTextParts(message: RunMessage): string[] {
const textParts = message.content
.filter((part) => part.type === "text")
@ -157,6 +160,11 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
const audioBase64 = findLatestUserAudioBase64(messages);
// Clear pending audio from store after extracting (consumed on send)
if (audioBase64) {
const audioName = runtime.pendingAudioName;
if (audioName) {
const lastUserMsg = [...messages].reverse().find((m) => m.role === "user");
if (lastUserMsg) sentAudioNames.set(lastUserMsg.id, audioName);
}
runtime.clearPendingAudio();
}
const useAdapter = await resolveUseAdapter(unstable_threadId);
@ -165,7 +173,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
const activeModel = runtime.models.find(
(m) => m.id === params.checkpoint,
);
if (activeModel?.isAudio) {
if (activeModel?.isAudio && !activeModel?.hasAudioInput) {
const threadKey = unstable_threadId || "__default";
runtime.setThreadRunning(threadKey, true);
try {

12
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -1,5 +1,6 @@
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
import { Button } from "@/components/ui/button";
import { AUDIO_ACCEPT, MAX_AUDIO_SIZE, fileToBase64 } from "@/lib/audio-utils";
import { useAui } from "@assistant-ui/react";
import { ArrowUpIcon, HeadphonesIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react";
import { useChatRuntimeStore } from "./stores/chat-runtime-store";
@ -29,8 +30,6 @@ export interface CompareHandle {
const IMAGE_ACCEPT = "image/jpeg,image/png,image/webp,image/gif";
const MAX_IMAGE_SIZE = 20 * 1024 * 1024;
const AUDIO_ACCEPT = "audio/wav,audio/mpeg,audio/webm,audio/ogg,audio/flac,audio/mp4";
const MAX_AUDIO_SIZE = 50 * 1024 * 1024;
function fileToBase64DataURL(file: File): Promise<string> {
return new Promise((resolve, reject) => {
@ -220,15 +219,10 @@ export function SharedComposer({
if (!file) continue;
// Handle audio files
if (file.type.match(/^audio\//i) && file.size <= MAX_AUDIO_SIZE) {
const reader = new FileReader();
reader.onload = () => {
const result = reader.result as string;
const commaIndex = result.indexOf(",");
const base64 = commaIndex >= 0 ? result.slice(commaIndex + 1) : result;
fileToBase64(file).then((base64) => {
setPendingAudio({ name: file.name, base64 });
setPendingAudioStore(base64, file.name);
};
reader.readAsDataURL(file);
});
continue;
}
// Handle image files

1
studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
View file

@ -91,6 +91,7 @@ export function DatasetMappingCard({
mappingOk,
autoDetected = false,
isVlm = false,
isAudio = false,
format,
}: {
mapping: Record<string, string>;

15
studio/frontend/src/lib/audio-utils.ts Normal file
View file

@ -0,0 +1,15 @@
export const AUDIO_ACCEPT = "audio/wav,audio/mpeg,audio/webm,audio/ogg,audio/flac,audio/mp4";
export const MAX_AUDIO_SIZE = 50 * 1024 * 1024;
export function fileToBase64(file: File): Promise<string> {
return new Promise((resolve, reject) => {
const reader = new FileReader();
reader.onload = () => {
const result = reader.result as string;
const commaIndex = result.indexOf(",");
resolve(commaIndex >= 0 ? result.slice(commaIndex + 1) : result);
};
reader.onerror = () => reject(new Error("Failed to read file"));
reader.readAsDataURL(file);
});
}