diff --git a/studio/TESTING.md b/studio/TESTING.md new file mode 100644 index 0000000000..cc2b5a9525 --- /dev/null +++ b/studio/TESTING.md @@ -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 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml index b7cb5830f4..1906ecda51 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml @@ -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 diff --git a/studio/backend/core/inference/OuteTTS b/studio/backend/core/inference/OuteTTS deleted file mode 160000 index 59d896747a..0000000000 --- a/studio/backend/core/inference/OuteTTS +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 59d896747aa0a6a207837e7da2d6921805eae684 diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 284c4fa392..a2449c534c 100644 --- a/studio/backend/core/inference/inference.py +++ b/studio/backend/core/inference/inference.py @@ -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() diff --git a/studio/backend/core/training/inference/OuteTTS b/studio/backend/core/training/inference/OuteTTS deleted file mode 160000 index 59d896747a..0000000000 --- a/studio/backend/core/training/inference/OuteTTS +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 59d896747aa0a6a207837e7da2d6921805eae684 diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index b9cfa6d130..8b760c7599 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -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 diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 3fa4e43da5..2008ab9469 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -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, diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index b885e130d5..cb2f097b92 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -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( diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 6144bbe09c..42f561ca14 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -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, diff --git a/studio/frontend/src/components/assistant-ui/thread.tsx b/studio/frontend/src/components/assistant-ui/thread.tsx index 6e09a4e8fe..d959305c63 100644 --- a/studio/frontend/src/components/assistant-ui/thread.tsx +++ b/studio/frontend/src/components/assistant-ui/thread.tsx @@ -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 { - 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 ( +
+
+ + {audioName} +
+
+ ); +}; + const UserMessage: FC = () => { return ( { data-role="user" > +
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index 1c2d46f2db..f8aefda438 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -11,6 +11,9 @@ import { type RunMessages = Parameters[0]["messages"]; type RunMessage = RunMessages[number]; +/** Tracks which user messages were sent with an audio file (messageId → filename). */ +export const sentAudioNames = new Map(); + 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 { diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index ea5be5ea3f..fb4eb8fa69 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -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 { 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 diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx index a0938f1338..70f0cf5231 100644 --- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx @@ -91,6 +91,7 @@ export function DatasetMappingCard({ mappingOk, autoDetected = false, isVlm = false, + isAudio = false, format, }: { mapping: Record; diff --git a/studio/frontend/src/lib/audio-utils.ts b/studio/frontend/src/lib/audio-utils.ts new file mode 100644 index 0000000000..5adde372f2 --- /dev/null +++ b/studio/frontend/src/lib/audio-utils.ts @@ -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 { + 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); + }); +}