Merge pull request #223 from unslothai/feature/support-for-audio-models

Adding support for audio llms

install_python_stack.py

@@ -171,6 +171,13 @@ def install_python_stack() -> int:
         req=REQ_ROOT / "extras.txt",
     )
 
+    # 3b. Extra dependencies (no-deps) — audio model support etc.
+    pip_install(
+        "Installing extras (no-deps)",
+        "--no-deps", "--no-cache-dir",
+        req=REQ_ROOT / "extras-no-deps.txt",
+    )
+
     # 4. Overrides (torchao, transformers) — force-reinstall
     pip_install(
         "Installing torchao + transformers overrides",
@@ -234,8 +241,11 @@ def install_python_stack() -> int:
         [sys.executable, str(SINGLE_ENV / "patch_metadata.py")],
     )
 
-    # 12. Final check
-    run("Running pip check", [sys.executable, "-m", "pip", "check"], quiet=False)
+    # 12. Final check (silent — third-party conflicts are expected)
+    subprocess.run(
+        [sys.executable, "-m", "pip", "check"],
+        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
+    )
 
     print(_green("✅ Python dependencies installed"))
     return 0

studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml

@@ -41,6 +41,8 @@ logging:
   tensorboard_dir: "runs"
   log_frequency: 10
 
+audio_input: true
+
 inference:
   temperature: 1.0
   top_k: 64

studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml

@@ -41,6 +41,8 @@ logging:
   tensorboard_dir: "runs"
   log_frequency: 10
 
+audio_input: true
+
 inference:
   temperature: 1.0
   top_k: 64

studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml

@@ -3,7 +3,10 @@
 # Also applies to: OuteAI/Llama-OuteTTS-1.0-1B
 # added inference parameters from unsloth notebook
 
+audio_type: dac
+
 training:
+  eval_steps: 0
   max_seq_length: 2048
   # num_epochs: 4
   num_epochs: 0

studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml

@@ -3,7 +3,10 @@
 # Also applies to: Spark-TTS-0.5B/LLM
 # added inference parameters from unsloth notebook
 
+audio_type: bicodec
+
 training:
+  eval_steps: 0
   max_seq_length: 2048
   # num_epochs: 4
   num_epochs: 0

studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml

@@ -2,7 +2,10 @@
 # Based on Sesame_CSM_(1B)-TTS.ipynb
 # Also applies to: sesame/csm-1b
 
+audio_type: csm
+
 training:
+  eval_steps: 0
   max_seq_length: 2048
   # num_epochs: 4
   num_epochs: 0

studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml

@@ -3,7 +3,10 @@
 # Also applies to: unsloth/orpheus-3b-0.1-ft-unsloth-bnb-4bit, canopylabs/orpheus-3b-0.1-ft, unsloth/orpheus-3b-0.1-ft-bnb-4bit
 # added inference parameters from unsloth notebook
 
+audio_type: snac
+
 training:
+  eval_steps: 0
   max_seq_length: 2048
   # num_epochs: 4
   num_epochs: 0

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

@@ -2,7 +2,11 @@
 # Based on Whisper.ipynb
 # Also applies to: unsloth/whisper-large-v3, openai/whisper-large-v3
 
+audio_type: whisper
+audio_input: true
+
 training:
+  eval_steps: 5
   max_seq_length: 448
   # num_epochs: 4
   num_epochs: 0

studio/backend/core/export/export.py

@@ -17,6 +17,7 @@ import torch
 from utils.hardware import clear_gpu_cache
 
 from utils.models import is_vision_model, get_base_model_from_lora
+from utils.models.model_config import detect_audio_type
 from core.inference import get_inference_backend
 
 logger = logging.getLogger(__name__)
@@ -96,6 +97,7 @@ class ExportBackend:
         self.current_tokenizer = None
         self.is_vision = False
         self.is_peft = False
+        self._audio_type = None
 
     def cleanup_memory(self):
         """Offload and delete all models from memory"""
@@ -111,6 +113,7 @@ class ExportBackend:
             self.current_model = None
             self.current_tokenizer = None
             self.current_checkpoint = None
+            self._audio_type = None
 
             # Clear GPU memory cache (handles gc + backend-specific cleanup)
             clear_gpu_cache()
@@ -148,24 +151,75 @@ class ExportBackend:
             # First, cleanup existing models
             self.cleanup_memory()
 
-            # Detect if vision model
             checkpoint_path_obj = Path(checkpoint_path)
 
-            # Check if it's a LoRA adapter
+            # Determine the model identity for type detection
             adapter_config = checkpoint_path_obj / "adapter_config.json"
+            base_model = None
             if adapter_config.exists():
-                # It's a LoRA - get base model to check vision
                 base_model = get_base_model_from_lora(checkpoint_path)
-                if base_model:
-                    self.is_vision = is_vision_model(base_model)
-                else:
+                if not base_model:
                     return False, "Could not determine base model for adapter"
-            else:
-                # Check the model itself
-                self.is_vision = is_vision_model(checkpoint_path)
+
+            model_id = base_model or checkpoint_path
+
+            # Detect audio type and vision
+            self._audio_type = detect_audio_type(model_id)
+            self.is_vision = not self._audio_type and is_vision_model(model_id)
 
             # Load model based on type
-            if self.is_vision:
+            if self._audio_type == 'csm':
+                from unsloth import FastModel
+                from transformers import CsmForConditionalGeneration
+                logger.info("Loading as CSM audio model...")
+                model, tokenizer = FastModel.from_pretrained(
+                    model_name=checkpoint_path,
+                    max_seq_length=max_seq_length,
+                    dtype=None,
+                    auto_model=CsmForConditionalGeneration,
+                    load_in_4bit=False,
+                )
+
+            elif self._audio_type == 'whisper':
+                from unsloth import FastModel
+                from transformers import WhisperForConditionalGeneration
+                logger.info("Loading as Whisper audio model...")
+                model, tokenizer = FastModel.from_pretrained(
+                    model_name=checkpoint_path,
+                    dtype=None,
+                    load_in_4bit=False,
+                    auto_model=WhisperForConditionalGeneration,
+                )
+
+            elif self._audio_type == 'snac':
+                logger.info("Loading as SNAC (Orpheus) audio model...")
+                model, tokenizer = FastLanguageModel.from_pretrained(
+                    model_name=checkpoint_path,
+                    max_seq_length=max_seq_length,
+                    dtype=None,
+                    load_in_4bit=load_in_4bit,
+                )
+
+            elif self._audio_type == 'bicodec':
+                from unsloth import FastModel
+                logger.info("Loading as BiCodec (Spark-TTS) audio model...")
+                model, tokenizer = FastModel.from_pretrained(
+                    model_name=checkpoint_path,
+                    max_seq_length=max_seq_length,
+                    dtype=torch.float32,
+                    load_in_4bit=False,
+                )
+
+            elif self._audio_type == 'dac':
+                from unsloth import FastModel
+                logger.info("Loading as DAC (OuteTTS) audio model...")
+                model, tokenizer = FastModel.from_pretrained(
+                    model_name=checkpoint_path,
+                    max_seq_length=max_seq_length,
+                    load_in_4bit=False,
+                )
+
+            elif self.is_vision:
                 logger.info("Loading as vision model...")
                 model, processor = FastVisionModel.from_pretrained(
                     model_name=checkpoint_path,
@@ -174,6 +228,7 @@ class ExportBackend:
                     load_in_4bit=load_in_4bit,
                 )
                 tokenizer = processor  # For vision models, processor acts as tokenizer
+
             else:
                 logger.info("Loading as text model...")
                 model, tokenizer = FastLanguageModel.from_pretrained(
@@ -191,7 +246,12 @@ class ExportBackend:
             self.current_tokenizer = tokenizer
             self.current_checkpoint = checkpoint_path
 
-            model_type = "Vision" if self.is_vision else "Text"
+            if self._audio_type:
+                model_type = f"Audio ({self._audio_type})"
+            elif self.is_vision:
+                model_type = "Vision"
+            else:
+                model_type = "Text"
             peft_info = " (PEFT Adapter)" if self.is_peft else " (Merged Model)"
 
             logger.info(f"Successfully loaded {model_type} model{peft_info}")
@@ -246,6 +306,9 @@ class ExportBackend:
             # Determine save method
             if format_type == "4-bit (FP4)":
                 save_method = "merged_4bit_forced"
+            elif self._audio_type == 'whisper':
+                # Whisper uses save_method=None for local 16-bit merged save
+                save_method = None
             else:  # 16-bit (FP16)
                 save_method = "merged_16bit"
 
@@ -271,10 +334,12 @@ class ExportBackend:
 
                 logger.info(f"Pushing merged model to Hub: {repo_id}")
 
+                # Whisper uses save_method=None for local but "merged_16bit" for hub push
+                hub_save_method = save_method if save_method is not None else "merged_16bit"
                 self.current_model.push_to_hub_merged(
                     repo_id,
                     self.current_tokenizer,
-                    save_method=save_method,
+                    save_method=hub_save_method,
                     token=hf_token,
                     private=private
                 )

studio/backend/core/inference/audio_codecs.py

@@ -0,0 +1,280 @@
+"""
+Audio codec loading and decoding for TTS inference.
+Supports: SNAC (Orpheus), CSM (Sesame), BiCodec (Spark), DAC (OuteTTS)
+"""
+import io
+import re
+import wave
+import logging
+from typing import Optional, Tuple
+
+import numpy as np
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+def _numpy_to_wav_bytes(waveform: np.ndarray, sample_rate: int) -> bytes:
+    """Convert a float32 numpy waveform to WAV bytes (16-bit PCM)."""
+    waveform = waveform.flatten()
+    peak = max(abs(waveform.max()), abs(waveform.min()))
+    if peak > 1.0:
+        waveform = waveform / peak
+    pcm = (waveform * 32767).astype(np.int16)
+
+    buf = io.BytesIO()
+    with wave.open(buf, "wb") as wf:
+        wf.setnchannels(1)
+        wf.setsampwidth(2)
+        wf.setframerate(sample_rate)
+        wf.writeframes(pcm.tobytes())
+
+    return buf.getvalue()
+
+
+class AudioCodecManager:
+    """Manages loading and caching of audio codec models for TTS decoding."""
+
+    def __init__(self):
+        self._snac_model = None
+        self._bicodec_tokenizer = None
+        self._bicodec_repo_path = None
+        self._dac_audio_codec = None
+
+    def load_codec(self, audio_type: str, device: str = "cuda", model_repo_path: Optional[str] = None) -> None:
+        """Load the appropriate codec for the given audio type."""
+        if audio_type == "snac":
+            self._load_snac(device)
+        elif audio_type == "bicodec":
+            self._load_bicodec(device, model_repo_path)
+        elif audio_type == "dac":
+            self._load_dac(device)
+        elif audio_type == "csm":
+            pass  # CSM decoding is built into the model (output_audio=True)
+        else:
+            raise ValueError(f"Unknown audio_type: {audio_type}")
+
+    # ── Lazy loaders ─────────────────────────────────────────────
+
+    def _load_snac(self, device: str) -> None:
+        if self._snac_model is not None:
+            return
+        from snac import SNAC
+        self._snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval()
+        logger.info("Loaded SNAC codec (24kHz)")
+
+    def _load_bicodec(self, device: str, model_repo_path: Optional[str] = None) -> None:
+        if self._bicodec_tokenizer is not None:
+            return
+        import os
+        import sys
+        import subprocess
+
+        # Clone SparkAudio/Spark-TTS GitHub repo for the sparktts Python package
+        # (same approach as training — the HF model repos don't contain the package)
+        spark_code_dir = os.path.join(os.path.dirname(model_repo_path or "."), "Spark-TTS")
+        sparktts_pkg = os.path.join(spark_code_dir, "sparktts")
+        if not os.path.isdir(sparktts_pkg):
+            logger.info(f"Cloning SparkAudio/Spark-TTS to {spark_code_dir}...")
+            subprocess.run(
+                ["git", "clone", "--depth", "1", "https://github.com/SparkAudio/Spark-TTS", spark_code_dir],
+                check=True,
+            )
+
+        if spark_code_dir not in sys.path:
+            sys.path.insert(0, spark_code_dir)
+
+        from sparktts.models.audio_tokenizer import BiCodecTokenizer
+
+        # BiCodecTokenizer needs the MODEL repo path (contains BiCodec/ weights)
+        tokenizer_path = model_repo_path or spark_code_dir
+        self._bicodec_repo_path = tokenizer_path
+        self._bicodec_tokenizer = BiCodecTokenizer(tokenizer_path, device)
+        logger.info(f"Loaded BiCodec tokenizer from {tokenizer_path}")
+
+    def _load_dac(self, device: str) -> None:
+        if self._dac_audio_codec is not None:
+            return
+        import os
+        import sys
+        import subprocess
+
+        # Clone OuteTTS repo (same pattern as Spark-TTS / BiCodec)
+        # The pip package has problematic dependencies; the notebook clones and
+        # removes gguf_model.py, interface.py, __init__.py before importing.
+        base_dir = os.path.dirname(os.path.abspath(__file__))
+        outetts_code_dir = os.path.join(base_dir, "OuteTTS")
+        outetts_pkg = os.path.join(outetts_code_dir, "outetts")
+        if not os.path.isdir(outetts_pkg):
+            logger.info(f"Cloning edwko/OuteTTS to {outetts_code_dir}...")
+            subprocess.run(
+                ["git", "clone", "--depth", "1", "https://github.com/edwko/OuteTTS", outetts_code_dir],
+                check=True,
+            )
+            # Remove files that pull in heavy / incompatible dependencies
+            # (matches notebook: gguf_model.py is under models/, others under outetts/)
+            remove_paths = [
+                os.path.join(outetts_pkg, "models", "gguf_model.py"),
+                os.path.join(outetts_pkg, "interface.py"),
+                os.path.join(outetts_pkg, "__init__.py"),
+            ]
+            for fpath in remove_paths:
+                if os.path.exists(fpath):
+                    os.remove(fpath)
+                    logger.info(f"Removed {fpath}")
+
+        if outetts_code_dir not in sys.path:
+            sys.path.insert(0, outetts_code_dir)
+
+        from outetts.version.v3.audio_processor import AudioProcessor
+        from outetts.models.config import ModelConfig as OuteTTSModelConfig
+
+        dummy_config = OuteTTSModelConfig(
+            tokenizer_path="OuteAI/Llama-OuteTTS-1.0-1B",
+            device=device,
+            audio_codec_path=None,
+        )
+        processor = AudioProcessor(config=dummy_config)
+        self._dac_audio_codec = processor.audio_codec
+        logger.info("Loaded DAC audio codec")
+
+    # ── Decoders ─────────────────────────────────────────────────
+
+    def decode_snac(self, generated_ids: torch.Tensor, device: str) -> Tuple[bytes, int]:
+        """
+        Decode SNAC tokens (Orpheus) into WAV bytes.
+
+        generated_ids: full model output including prompt tokens.
+        Looks for START_OF_SPEECH (128257) marker, extracts codes after it,
+        strips EOS (128258), redistributes 7-per-frame codes into 3 SNAC layers.
+
+        Returns (wav_bytes, 24000).
+        """
+        # Find START_OF_SPEECH token (128257)
+        token_indices = (generated_ids == 128257).nonzero(as_tuple=True)
+        if len(token_indices[1]) > 0:
+            cropped = generated_ids[:, token_indices[1][-1] + 1:]
+        else:
+            # Gracefully fall back to using entire output if marker not found
+            logger.warning("No START_OF_SPEECH token (128257) found — using full generated output")
+            cropped = generated_ids
+        row = cropped[0]
+
+        # Remove EOS tokens (128258)
+        row = row[row != 128258]
+
+        # Trim to multiple of 7
+        row = row[: (len(row) // 7) * 7]
+        if len(row) == 0:
+            raise ValueError("No valid audio codes found after START_OF_SPEECH token")
+
+        codes = [t.item() - 128266 for t in row]
+
+        # Redistribute into 3 SNAC layers (7 codes per frame → 1+2+4)
+        layer_1, layer_2, layer_3 = [], [], []
+        for i in range(len(codes) // 7):
+            layer_1.append(codes[7 * i])
+            layer_2.append(codes[7 * i + 1] - 4096)
+            layer_3.append(codes[7 * i + 2] - 8192)
+            layer_3.append(codes[7 * i + 3] - 12288)
+            layer_2.append(codes[7 * i + 4] - 16384)
+            layer_3.append(codes[7 * i + 5] - 20480)
+            layer_3.append(codes[7 * i + 6] - 24576)
+
+        snac_codes = [
+            torch.tensor(layer).unsqueeze(0).to(device)
+            for layer in [layer_1, layer_2, layer_3]
+        ]
+
+        with torch.no_grad():
+            audio = self._snac_model.decode(snac_codes)
+
+        waveform = audio.squeeze().cpu().numpy()
+        return _numpy_to_wav_bytes(waveform, 24000), 24000
+
+    def decode_csm(self, audio_values: torch.Tensor) -> Tuple[bytes, int]:
+        """
+        Decode CSM output (already a waveform from model.generate(output_audio=True)).
+        Returns (wav_bytes, 24000).
+        """
+        waveform = audio_values[0].to(torch.float32).cpu().numpy()
+        return _numpy_to_wav_bytes(waveform, 24000), 24000
+
+    def decode_bicodec(self, generated_text: str, device: str) -> Tuple[bytes, int]:
+        """
+        Decode BiCodec tokens (Spark-TTS) from generated text.
+        Extracts bicodec_semantic_N and bicodec_global_N tokens via regex.
+        Returns (wav_bytes, sample_rate).
+        """
+        semantic_matches = re.findall(r"<\|bicodec_semantic_(\d+)\|>", generated_text)
+        global_matches = re.findall(r"<\|bicodec_global_(\d+)\|>", generated_text)
+
+        logger.info(f"BiCodec decode: {len(global_matches)} global tokens, {len(semantic_matches)} semantic tokens")
+        if len(global_matches) < 10:
+            logger.info(f"BiCodec generated text (first 500 chars): {generated_text[:500]}")
+
+        if not semantic_matches:
+            raise ValueError("No bicodec_semantic tokens found in generated output")
+
+        semantic_ids = torch.tensor([int(t) for t in semantic_matches]).long().unsqueeze(0)
+
+        # Speaker encoder expects exactly 32 global tokens (token_num=32 in BiCodec config).
+        # Pad with zeros or truncate to 32.
+        GLOBAL_TOKEN_NUM = 32
+        if global_matches:
+            raw = [int(t) for t in global_matches]
+        else:
+            raw = []
+        if len(raw) < GLOBAL_TOKEN_NUM:
+            raw = raw + [0] * (GLOBAL_TOKEN_NUM - len(raw))
+        raw = raw[:GLOBAL_TOKEN_NUM]
+        global_ids = torch.tensor(raw).long().unsqueeze(0)  # (1, 32)
+
+        self._bicodec_tokenizer.device = device
+        self._bicodec_tokenizer.model.to(device)
+
+        wav_np = self._bicodec_tokenizer.detokenize(
+            global_ids.to(device),
+            semantic_ids.to(device),
+        )
+        sr = self._bicodec_tokenizer.config.get("sample_rate", 16000)
+        return _numpy_to_wav_bytes(wav_np, sr), sr
+
+    def decode_dac(self, generated_text: str, device: str) -> Tuple[bytes, int]:
+        """
+        Decode DAC tokens (OuteTTS) from generated text.
+        Extracts c1_N and c2_N codec code tokens via regex.
+        Returns (wav_bytes, 24000).
+        """
+        c1 = list(map(int, re.findall(r"<\|c1_(\d+)\|>", generated_text)))
+        c2 = list(map(int, re.findall(r"<\|c2_(\d+)\|>", generated_text)))
+
+        if not c1 or not c2:
+            raise ValueError("No DAC code tokens (c1/c2) found in generated output")
+
+        t = min(len(c1), len(c2))
+        c1 = c1[:t]
+        c2 = c2[:t]
+
+        codes = torch.tensor([[c1, c2]], dtype=torch.int64).to(device)
+        with torch.no_grad():
+            audio = self._dac_audio_codec.decode(codes)
+
+        waveform = audio.squeeze().cpu().numpy()
+        return _numpy_to_wav_bytes(waveform, 24000), 24000
+
+    # ── Cleanup ──────────────────────────────────────────────────
+
+    def unload(self) -> None:
+        """Release all codec models from memory."""
+        if self._snac_model is not None:
+            del self._snac_model
+            self._snac_model = None
+        if self._bicodec_tokenizer is not None:
+            del self._bicodec_tokenizer
+            self._bicodec_tokenizer = None
+            self._bicodec_repo_path = None
+        if self._dac_audio_codec is not None:
+            del self._dac_audio_codec
+            self._dac_audio_codec = None
+        logger.info("Unloaded all audio codecs")

studio/backend/core/inference/inference.py

@@ -15,6 +15,7 @@ from utils.models import ModelConfig, get_base_model_from_lora
 from utils.paths import is_model_cached
 from utils.utils import format_error_message
 from utils.hardware import get_device, clear_gpu_cache, log_gpu_memory
+from core.inference.audio_codecs import AudioCodecManager
 from io import StringIO
 import logging
 
@@ -39,6 +40,7 @@ class InferenceBackend:
             "unsloth/Qwen2-VL-2B-Instruct-bnb-4bit",
         ]
         self.device = get_device().value
+        self._audio_codec_manager = AudioCodecManager()
 
         # Thread safety — _generation_lock serializes model.generate() calls.
         # Must be a regular Lock (NOT RLock) because in async FastAPI, multiple
@@ -84,12 +86,146 @@ class InferenceBackend:
             self.models[model_name] = {
                 "is_vision": config.is_vision,
                 "is_lora": config.is_lora,
+                "is_audio": config.is_audio,
+                "audio_type": config.audio_type,
+                "has_audio_input": config.has_audio_input,
                 "model_path": config.path,
                 "base_model": config.base_model if config.is_lora else None,
                 "loaded_adapters": {},
                 "active_adapter": None,
             }
 
+            # ── Audio model loading path ──────────────────────────
+            if config.is_audio:
+                audio_type = config.audio_type
+                adapter_info = " (LoRA adapter)" if config.is_lora else ""
+                logger.info(f"Loading audio ({audio_type}) model{adapter_info}: {model_name}")
+                log_gpu_memory(f"Before loading {model_name}")
+
+                if audio_type == "csm":
+                    from unsloth import FastModel
+                    from transformers import CsmForConditionalGeneration
+                    model, processor = FastModel.from_pretrained(
+                        config.path,
+                        auto_model=CsmForConditionalGeneration,
+                        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"] = processor
+                    self.models[model_name]["processor"] = processor
+                elif audio_type == "bicodec":
+                    import os
+                    from unsloth import FastModel
+
+                    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,
+                        )
+
+                    FastModel.for_inference(model)
+                    self.models[model_name]["model"] = model
+                    self.models[model_name]["tokenizer"] = tokenizer
+                    self.models[model_name]["model_repo_path"] = abs_repo_path
+                elif audio_type == "dac":
+                    # OuteTTS uses FastModel (not FastLanguageModel)
+                    from unsloth import FastModel
+                    model, tokenizer = FastModel.from_pretrained(
+                        config.path,
+                        max_seq_length=max_seq_length,
+                        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
+                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(
+                        model_name=config.path,
+                        max_seq_length=max_seq_length,
+                        load_in_4bit=False,
+                        token=hf_token if hf_token and hf_token.strip() else None,
+                    )
+                    FastLanguageModel.for_inference(model)
+                    self.models[model_name]["model"] = model
+                    self.models[model_name]["tokenizer"] = tokenizer
+
+                # Load the external codec for TTS audio types
+                # (Whisper is ASR, audio_vlm is audio input — neither needs a codec)
+                if audio_type not in ("whisper", "audio_vlm"):
+                    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)
+                logger.info(f"Successfully loaded audio model: {model_name}")
+                log_gpu_memory(f"After loading {model_name}")
+                return True
+
             model_type = "vision" if config.is_vision else "text"
             adapter_info = " (LoRA adapter)" if self.models[model_name]["is_lora"] else ""
             logger.info(f"Loading {model_type} model{adapter_info}: {model_name}")
@@ -177,15 +313,17 @@ 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.
         """
         if model_name in self.models:
             try:
+                # If this was an audio model, clean up codecs
+                if self.models[model_name].get("is_audio"):
+                    self._audio_codec_manager.unload()
+
                 logger.info(f"Unloading model '{model_name}' from memory.")
                 # Delete the model entry from our registry
                 del self.models[model_name]
@@ -209,7 +347,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:
         """
@@ -245,61 +382,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 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
-
     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]]:
@@ -346,7 +428,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:
         """
@@ -374,7 +455,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:
         """
@@ -390,7 +470,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:
         """
@@ -709,7 +788,146 @@ 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,
+                                       max_new_tokens, repetition_penalty,
+                                       cancel_event=None) -> Generator[str, None, None]:
+        """Handle audio input (ASR) generation — accepts audio numpy array, streams text output.
+
+        Uses processor.apply_chat_template with audio embedded in messages (Gemma 3n pattern).
+        """
+        import threading
+        import numpy as np
+
+        model_info = self.models[self.active_model_name]
+        model = model_info["model"]
+        processor = model_info.get("processor") or model_info.get("tokenizer")
+        raw_tokenizer = getattr(processor, "tokenizer", processor)
+
+        # 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 = [
+            {"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(
+            audio_messages,
+            add_generation_prompt=True,
+            tokenize=True,
+            return_dict=True,
+            return_tensors="pt",
+            truncation=False,
+        ).to(self.device)
+
+        try:
+            from transformers import TextIteratorStreamer
+            from queue import Empty
+
+            streamer = TextIteratorStreamer(
+                raw_tokenizer,
+                skip_prompt=True,
+                skip_special_tokens=True,
+                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=False,
+            )
+
+            err: dict[str, str] = {}
+
+            def generate_fn():
+                with self._generation_lock:
+                    try:
+                        model.generate(**generation_kwargs)
+                    except Exception as e:
+                        err["msg"] = str(e)
+                        logger.error(f"Audio input generation error in thread: {e}")
+                    finally:
+                        try:
+                            streamer.end()
+                        except Exception:
+                            pass
+
+            thread = threading.Thread(target=generate_fn)
+            thread.start()
+
+            output = ""
+            try:
+                while True:
+                    if cancel_event is not None and cancel_event.is_set():
+                        break
+                    try:
+                        new_token = next(streamer)
+                    except StopIteration:
+                        break
+                    except Empty:
+                        if not thread.is_alive():
+                            break
+                        continue
+                    if new_token:
+                        output += new_token
+                        yield new_token
+            finally:
+                if cancel_event is not None:
+                    cancel_event.set()
+                thread.join(timeout=10)
+                if thread.is_alive():
+                    logger.warning("Audio input generation thread did not exit after cancel/join timeout")
+
+            if err.get("msg"):
+                yield f"Error: {err['msg']}"
+
+        except Exception as e:
+            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,
@@ -832,8 +1050,164 @@ 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(
+        self,
+        text: 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,
+        use_adapter: Optional[Union[bool, str]] = None,
+    ) -> Tuple[bytes, int]:
+        """
+        Generate audio from text for TTS models.
+        Returns (wav_bytes, sample_rate).
+        Blocking — generates complete audio before returning.
+        """
+        if not self.active_model_name:
+            raise RuntimeError("No active model")
+
+        model_info = self.models[self.active_model_name]
+        audio_type = model_info.get("audio_type")
+        model = model_info["model"]
+        tokenizer = model_info.get("tokenizer")
+
+        if not audio_type:
+            raise RuntimeError(f"Model {self.active_model_name} is not an audio model")
+
+        top_k = self._normalize_top_k(top_k)
+
+        with self._generation_lock:
+            if use_adapter is not None:
+                self._apply_adapter_state(use_adapter)
+
+            if audio_type == "snac":
+                return self._generate_snac(model, tokenizer, text, temperature, top_p, max_new_tokens, repetition_penalty)
+            elif audio_type == "csm":
+                processor = model_info.get("processor", tokenizer)
+                return self._generate_csm(model, processor, text, max_new_tokens)
+            elif audio_type == "bicodec":
+                return self._generate_bicodec(model, tokenizer, text, temperature, top_k, max_new_tokens)
+            elif audio_type == "dac":
+                return self._generate_dac(model, tokenizer, text, temperature, top_k, top_p, min_p, max_new_tokens, repetition_penalty)
+            else:
+                raise RuntimeError(f"Unknown audio_type: {audio_type}")
+
+    def _generate_snac(self, model, tokenizer, text, temperature, top_p, max_new_tokens, repetition_penalty):
+        """Generate audio using SNAC codec (Orpheus)."""
+        device = model.device
+        start_token = torch.tensor([[128259]], device=device)   # START_OF_HUMAN
+        end_tokens = torch.tensor([[128009, 128260]], device=device)  # EOT, END_OF_HUMAN
+        text_ids = tokenizer(text, return_tensors="pt").input_ids.to(device)
+        input_ids = torch.cat([start_token, text_ids, end_tokens], dim=1)
+        attention_mask = torch.ones_like(input_ids)
+
+        generated = model.generate(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=temperature,
+            top_p=top_p,
+            repetition_penalty=repetition_penalty,
+            eos_token_id=128258,  # END_OF_SPEECH
+            use_cache=True,
+        )
+        return self._audio_codec_manager.decode_snac(generated, str(device))
+
+    def _generate_csm(self, model, processor, text, max_new_tokens):
+        """Generate audio using CSM (Sesame)."""
+        speaker_id = 0
+        inputs = processor(f"[{speaker_id}]{text}", add_special_tokens=True, return_tensors="pt").to(model.device)
+        audio_values = model.generate(**inputs, max_new_tokens=max_new_tokens, output_audio=True)
+        return self._audio_codec_manager.decode_csm(audio_values)
+
+    def _generate_bicodec(self, model, tokenizer, text, temperature, top_k, max_new_tokens):
+        """Generate audio using BiCodec (Spark-TTS)."""
+        prompt = "<|task_tts|><|start_content|>" + text + "<|end_content|><|start_global_token|>"
+        inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
+        generated = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=True,
+            temperature=temperature,
+            top_k=top_k,
+            eos_token_id=tokenizer.eos_token_id,
+            pad_token_id=tokenizer.pad_token_id,
+        )
+        new_tokens = generated[:, inputs.input_ids.shape[1]:]
+        decoded_text = tokenizer.batch_decode(new_tokens, skip_special_tokens=False)[0]
+        return self._audio_codec_manager.decode_bicodec(decoded_text, str(model.device))
+
+    def _generate_dac(self, model, tokenizer, text, temperature, top_k, top_p, min_p, max_new_tokens, repetition_penalty):
+        """Generate audio using DAC (OuteTTS). Follows Oute_TTS_(1B).ipynb exactly."""
+        # Monkey-patch RepetitionPenaltyLogitsProcessor with a 64-token penalty
+        # window (same as the OuteTTS notebook) to avoid degenerate repetition.
+        self._patch_repetition_penalty_processor()
+
+        prompt = "<|im_start|>\n<|text_start|>" + text + "<|text_end|>\n<|audio_start|><|global_features_start|>\n"
+        with torch.inference_mode():
+            with torch.amp.autocast('cuda', dtype=model.dtype):
+                inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
+                generated = model.generate(
+                    **inputs,
+                    temperature=temperature,
+                    top_k=top_k,
+                    top_p=top_p,
+                    min_p=min_p,
+                    repetition_penalty=repetition_penalty,
+                    max_new_tokens=max_new_tokens,
+                )
+        decoded_text = tokenizer.batch_decode(generated, skip_special_tokens=False)[0]
+        return self._audio_codec_manager.decode_dac(decoded_text, str(model.device))
+
+    _repetition_penalty_patched = False
+
+    @classmethod
+    def _patch_repetition_penalty_processor(cls):
+        """
+        Monkey-patch transformers' RepetitionPenaltyLogitsProcessor with a
+        64-token sliding window variant (from the OuteTTS notebook).
+        Only applied once per process.
+        """
+        if cls._repetition_penalty_patched:
+            return
+        cls._repetition_penalty_patched = True
+
+        from transformers import LogitsProcessor
+        import transformers.generation.utils as generation_utils
+
+        class RepetitionPenaltyLogitsProcessorPatch(LogitsProcessor):
+            def __init__(self, penalty: float):
+                self.penalty_last_n = 64
+                if not isinstance(penalty, float) or penalty <= 0:
+                    raise ValueError(f"`penalty` has to be a positive float, but is {penalty}")
+                self.penalty = penalty
+
+            @torch.no_grad()
+            def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor:
+                if self.penalty_last_n == 0 or self.penalty == 1.0:
+                    return scores
+                batch_size, seq_len = input_ids.shape
+                vocab_size = scores.shape[-1]
+                for b in range(batch_size):
+                    start_index = max(0, seq_len - self.penalty_last_n)
+                    window_indices = input_ids[b, start_index:]
+                    if window_indices.numel() == 0:
+                        continue
+                    for token_id in set(window_indices.tolist()):
+                        if token_id >= vocab_size:
+                            continue
+                        logit = scores[b, token_id]
+                        scores[b, token_id] = logit * self.penalty if logit <= 0 else logit / self.penalty
+                return scores
+
+        generation_utils.RepetitionPenaltyLogitsProcessor = RepetitionPenaltyLogitsProcessorPatch
+        logger.info("Patched RepetitionPenaltyLogitsProcessor with 64-token window for OuteTTS")
 
     def format_chat_prompt(self, messages: list, system_prompt: str = None) -> str:
         if not self.active_model_name or self.active_model_name not in self.models:
@@ -1056,7 +1430,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"""
@@ -1183,10 +1556,10 @@ class InferenceBackend:
         return next(iter(self.loading_models)) if self.loading_models else None
 
     def load_model_simple(self,
-                     model_path: str,
-                     hf_token: Optional[str] = None,
-                     max_seq_length: int = 2048,
-                     load_in_4bit: bool = True) -> bool:
+                         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.
@@ -1201,10 +1574,6 @@ class InferenceBackend:
             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,
@@ -1212,8 +1581,6 @@ class InferenceBackend:
                 is_lora=False
             )
 
-            logger.info(f"Created ModelConfig with identifier: {config.identifier}")
-
             # Call existing load_model with config
             return self.load_model(
                 config=config,
@@ -1225,11 +1592,8 @@ class InferenceBackend:
 
         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

studio/backend/core/inference/orchestrator.py

@@ -312,6 +312,9 @@ class InferenceOrchestrator:
                     "is_vision": model_info.get("is_vision", False),
                     "is_lora": model_info.get("is_lora", False),
                     "display_name": model_info.get("display_name", model_name),
+                    "is_audio": model_info.get("is_audio", False),
+                    "audio_type": model_info.get("audio_type"),
+                    "has_audio_input": model_info.get("has_audio_input", False),
                 }
                 self.loading_models.discard(model_name)
                 logger.info("Model '%s' loaded successfully in subprocess", model_name)
@@ -545,6 +548,203 @@ class InferenceOrchestrator:
         except RuntimeError:
             pass
 
+    # ------------------------------------------------------------------
+    # Audio generation — TTS, ASR, audio input
+    # ------------------------------------------------------------------
+
+    def generate_audio_response(
+        self,
+        text: 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,
+        use_adapter: Optional[Union[bool, str]] = None,
+    ) -> Tuple[bytes, int]:
+        """Generate TTS audio. Returns (wav_bytes, sample_rate).
+
+        Blocking — sends command and waits for the complete audio response.
+        """
+        if not self._ensure_subprocess_alive():
+            raise RuntimeError("Inference subprocess is not running")
+        if not self.active_model_name:
+            raise RuntimeError("No active model")
+
+        import uuid
+        request_id = str(uuid.uuid4())
+
+        cmd = {
+            "type": "generate_audio",
+            "request_id": request_id,
+            "text": text,
+            "temperature": temperature,
+            "top_p": top_p,
+            "top_k": top_k,
+            "min_p": min_p,
+            "max_new_tokens": max_new_tokens,
+            "repetition_penalty": repetition_penalty,
+        }
+        if use_adapter is not None:
+            cmd["use_adapter"] = use_adapter
+
+        self._send_cmd(cmd)
+
+        # Wait for audio_done or audio_error
+        deadline = time.monotonic() + 120.0
+        while time.monotonic() < deadline:
+            remaining = max(0.1, deadline - time.monotonic())
+            resp = self._read_resp(timeout=min(remaining, 1.0))
+
+            if resp is None:
+                if not self._ensure_subprocess_alive():
+                    raise RuntimeError("Inference subprocess crashed during audio generation")
+                continue
+
+            rtype = resp.get("type", "")
+
+            if rtype == "audio_done":
+                wav_bytes = base64.b64decode(resp["wav_base64"])
+                sample_rate = resp["sample_rate"]
+                return wav_bytes, sample_rate
+
+            if rtype == "audio_error":
+                raise RuntimeError(resp.get("error", "Audio generation failed"))
+
+            if rtype == "error":
+                raise RuntimeError(resp.get("error", "Unknown error"))
+
+            if rtype == "status":
+                continue
+
+        raise RuntimeError("Timeout waiting for audio generation (120s)")
+
+    def generate_whisper_response(
+        self,
+        audio_array,
+        cancel_event=None,
+    ) -> Generator[str, None, None]:
+        """Whisper ASR — sends audio to subprocess, yields text."""
+        yield from self._generate_audio_input_inner(
+            audio_array=audio_array,
+            audio_type="whisper",
+            messages=[],
+            system_prompt="",
+            cancel_event=cancel_event,
+        )
+
+    def generate_audio_input_response(
+        self,
+        messages,
+        system_prompt,
+        audio_array,
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        top_k: int = 40,
+        min_p: float = 0.0,
+        max_new_tokens: int = 512,
+        repetition_penalty: float = 1.1,
+        cancel_event=None,
+    ) -> Generator[str, None, None]:
+        """Audio input generation (e.g. Gemma 3n) — streams text tokens."""
+        yield from self._generate_audio_input_inner(
+            audio_array=audio_array,
+            audio_type=None,  # worker will use generate_audio_input_response
+            messages=messages,
+            system_prompt=system_prompt,
+            temperature=temperature,
+            top_p=top_p,
+            top_k=top_k,
+            min_p=min_p,
+            max_new_tokens=max_new_tokens,
+            repetition_penalty=repetition_penalty,
+            cancel_event=cancel_event,
+        )
+
+    def _generate_audio_input_inner(
+        self,
+        audio_array,
+        audio_type: Optional[str] = None,
+        messages: list = None,
+        system_prompt: str = "",
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        top_k: int = 40,
+        min_p: float = 0.0,
+        max_new_tokens: int = 512,
+        repetition_penalty: float = 1.1,
+        cancel_event=None,
+    ) -> Generator[str, None, None]:
+        """Shared inner logic for audio input generation (Whisper + ASR)."""
+        if not self._ensure_subprocess_alive():
+            yield "Error: Inference subprocess is not running"
+            return
+        if not self.active_model_name:
+            yield "Error: No active model"
+            return
+
+        with self._gen_lock:
+            import uuid
+            request_id = str(uuid.uuid4())
+
+            # Convert numpy array to list for mp.Queue serialization
+            audio_data = audio_array.tolist() if hasattr(audio_array, 'tolist') else list(audio_array)
+
+            cmd = {
+                "type": "generate_audio_input",
+                "request_id": request_id,
+                "audio_data": audio_data,
+                "audio_type": audio_type,
+                "messages": messages or [],
+                "system_prompt": system_prompt,
+                "temperature": temperature,
+                "top_p": top_p,
+                "top_k": top_k,
+                "min_p": min_p,
+                "max_new_tokens": max_new_tokens,
+                "repetition_penalty": repetition_penalty,
+            }
+
+            try:
+                self._send_cmd(cmd)
+            except RuntimeError as exc:
+                yield f"Error: {exc}"
+                return
+
+            # Yield tokens — same pattern as _generate_locked
+            while True:
+                resp = self._read_resp(timeout=30.0)
+
+                if resp is None:
+                    if not self._ensure_subprocess_alive():
+                        yield "Error: Inference subprocess crashed during audio input generation"
+                        return
+                    continue
+
+                rtype = resp.get("type", "")
+
+                if rtype == "status":
+                    continue
+
+                if rtype == "error" and not resp.get("request_id"):
+                    yield f"Error: {resp.get('error', 'Unknown error')}"
+                    return
+
+                if rtype == "token":
+                    if cancel_event is not None and cancel_event.is_set():
+                        self._cancel_generation()
+                        self._drain_until_gen_done(timeout=5.0)
+                        return
+                    yield resp.get("text", "")
+
+                elif rtype == "gen_done":
+                    return
+
+                elif rtype == "gen_error":
+                    yield f"Error: {resp.get('error', 'Unknown error')}"
+                    return
+
     # ------------------------------------------------------------------
     # Local helpers (no subprocess needed)
     # ------------------------------------------------------------------

studio/backend/core/inference/worker.py

@@ -165,6 +165,9 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
                 "is_vision": mc.is_vision,
                 "is_lora": mc.is_lora,
                 "is_gguf": False,
+                "is_audio": getattr(mc, "is_audio", False),
+                "audio_type": getattr(mc, "audio_type", None),
+                "has_audio_input": getattr(mc, "has_audio_input", False),
             }
             _send_response(resp_queue, {
                 "type": "loaded",
@@ -267,6 +270,110 @@ def _handle_generate(
         })
 
 
+def _handle_generate_audio(
+    backend,
+    cmd: dict,
+    resp_queue: Any,
+) -> None:
+    """Handle TTS audio generation — returns WAV bytes + sample_rate."""
+    request_id = cmd.get("request_id", "")
+    try:
+        wav_bytes, sample_rate = backend.generate_audio_response(
+            text=cmd["text"],
+            temperature=cmd.get("temperature", 0.6),
+            top_p=cmd.get("top_p", 0.95),
+            top_k=cmd.get("top_k", 50),
+            min_p=cmd.get("min_p", 0.0),
+            max_new_tokens=cmd.get("max_new_tokens", 2048),
+            repetition_penalty=cmd.get("repetition_penalty", 1.1),
+            use_adapter=cmd.get("use_adapter"),
+        )
+
+        # Send WAV bytes as base64 (bytes can't go through mp.Queue directly)
+        _send_response(resp_queue, {
+            "type": "audio_done",
+            "request_id": request_id,
+            "wav_base64": base64.b64encode(wav_bytes).decode("ascii"),
+            "sample_rate": sample_rate,
+            "ts": time.time(),
+        })
+
+    except Exception as exc:
+        logger.error("Audio generation error: %s", exc, exc_info=True)
+        _send_response(resp_queue, {
+            "type": "audio_error",
+            "request_id": request_id,
+            "error": str(exc),
+            "stack": traceback.format_exc(limit=20),
+            "ts": time.time(),
+        })
+
+
+def _handle_generate_audio_input(
+    backend,
+    cmd: dict,
+    resp_queue: Any,
+    cancel_event,
+) -> None:
+    """Handle audio input generation (ASR/Whisper) — streams text tokens back."""
+    request_id = cmd.get("request_id", "")
+
+    try:
+        import numpy as np
+
+        # Decode audio array from list (numpy arrays can't go through mp.Queue)
+        audio_array = np.array(cmd["audio_data"], dtype=np.float32)
+
+        audio_type = cmd.get("audio_type")
+
+        if audio_type == "whisper":
+            generator = backend.generate_whisper_response(
+                audio_array=audio_array,
+                cancel_event=cancel_event,
+            )
+        else:
+            generator = backend.generate_audio_input_response(
+                messages=cmd.get("messages", []),
+                system_prompt=cmd.get("system_prompt", ""),
+                audio_array=audio_array,
+                temperature=cmd.get("temperature", 0.7),
+                top_p=cmd.get("top_p", 0.9),
+                top_k=cmd.get("top_k", 40),
+                min_p=cmd.get("min_p", 0.0),
+                max_new_tokens=cmd.get("max_new_tokens", 512),
+                repetition_penalty=cmd.get("repetition_penalty", 1.1),
+                cancel_event=cancel_event,
+            )
+
+        for text_chunk in generator:
+            if cancel_event.is_set():
+                logger.info("Audio input generation cancelled for request %s", request_id)
+                break
+
+            _send_response(resp_queue, {
+                "type": "token",
+                "request_id": request_id,
+                "text": text_chunk,
+                "ts": time.time(),
+            })
+
+        _send_response(resp_queue, {
+            "type": "gen_done",
+            "request_id": request_id,
+            "ts": time.time(),
+        })
+
+    except Exception as exc:
+        logger.error("Audio input generation error: %s", exc, exc_info=True)
+        _send_response(resp_queue, {
+            "type": "gen_error",
+            "request_id": request_id,
+            "error": str(exc),
+            "stack": traceback.format_exc(limit=20),
+            "ts": time.time(),
+        })
+
+
 def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
     """Handle an unload command."""
     model_name = cmd.get("model_name", "")
@@ -414,6 +521,14 @@ def run_inference_process(
                     backend.unload_model(backend.active_model_name)
                 _handle_load(backend, cmd, resp_queue)
 
+            elif cmd_type == "generate_audio":
+                cancel_event.clear()
+                _handle_generate_audio(backend, cmd, resp_queue)
+
+            elif cmd_type == "generate_audio_input":
+                cancel_event.clear()
+                _handle_generate_audio_input(backend, cmd, resp_queue, cancel_event)
+
             elif cmd_type == "unload":
                 _handle_unload(backend, cmd, resp_queue)
 

studio/backend/core/training/training.py

@@ -143,7 +143,8 @@ class TrainingBackend:
             "dataset_slice_start": kwargs.get("dataset_slice_start"),
             "dataset_slice_end": kwargs.get("dataset_slice_end"),
             "custom_format_mapping": kwargs.get("custom_format_mapping"),
-            "is_dataset_multimodal": kwargs.get("is_dataset_multimodal", False),
+            "is_dataset_image": kwargs.get("is_dataset_image", False),
+            "is_dataset_audio": kwargs.get("is_dataset_audio", False),
             "num_epochs": kwargs.get("num_epochs", 3),
             "learning_rate": kwargs.get("learning_rate", "2e-4"),
             "batch_size": kwargs.get("batch_size", 2),

studio/backend/core/training/worker.py

@@ -190,7 +190,8 @@ def run_training_process(
             max_seq_length=config["max_seq_length"],
             load_in_4bit=config["load_in_4bit"],
             hf_token=hf_token,
-            is_dataset_multimodal=config.get("is_dataset_multimodal", False),
+            is_dataset_image=config.get("is_dataset_image", False),
+            is_dataset_audio=config.get("is_dataset_audio", False),
         )
         if not success or trainer.should_stop:
             if trainer.should_stop:

studio/backend/models/datasets.py

@@ -27,11 +27,14 @@ class CheckFormatResponse(BaseModel):
     requires_manual_mapping: bool
     detected_format: str
     columns: List[str]
-    is_multimodal: bool = False
+    is_image: bool = False
+    is_audio: bool = False
     multimodal_columns: Optional[List[str]] = None
     suggested_mapping: Optional[Dict[str, str]] = None
     detected_image_column: Optional[str] = None
+    detected_audio_column: Optional[str] = None
     detected_text_column: Optional[str] = None
+    detected_speaker_column: Optional[str] = None
     preview_samples: Optional[List[Dict]] = None
     total_rows: Optional[int] = None
     warning: Optional[str] = None

studio/backend/models/inference.py

@@ -45,6 +45,9 @@ class LoadResponse(BaseModel):
     is_vision: bool = Field(False, description="Whether model is a vision model")
     is_lora: bool = Field(False, description="Whether model is a LoRA adapter")
     is_gguf: bool = Field(False, description="Whether model is a GGUF model (llama.cpp)")
+    is_audio: bool = Field(False, description="Whether model is a TTS audio model")
+    audio_type: Optional[str] = Field(None, description="Audio codec type: snac, csm, bicodec, dac")
+    has_audio_input: bool = Field(False, description="Whether model accepts audio input (ASR)")
     inference: dict = Field(..., description="Inference parameters (temperature, top_p, top_k, min_p)")
 
 
@@ -60,6 +63,9 @@ class InferenceStatusResponse(BaseModel):
     is_vision: bool = Field(False, description="Whether the active model is a vision model")
     is_gguf: bool = Field(False, description="Whether the active model is a GGUF model (llama.cpp)")
     gguf_variant: Optional[str] = Field(None, description="GGUF quantization variant (e.g. Q4_K_M)")
+    is_audio: bool = Field(False, description="Whether the active model is a TTS audio model")
+    audio_type: Optional[str] = Field(None, description="Audio codec type: snac, csm, bicodec, dac")
+    has_audio_input: bool = Field(False, description="Whether model accepts audio input (ASR)")
     loading: List[str] = Field(default_factory=list, description="Models currently being loaded")
     loaded: List[str] = Field(default_factory=list, description="Models currently loaded")
 
@@ -136,6 +142,7 @@ class ChatCompletionRequest(BaseModel):
     min_p: float = Field(0.0, ge=0.0, le=1.0, description="[x-unsloth] Min-p sampling threshold")
     repetition_penalty: float = Field(1.1, ge=1.0, le=2.0, description="[x-unsloth] Repetition penalty")
     image_base64: Optional[str] = Field(None, description="[x-unsloth] Base64-encoded image for vision models")
+    audio_base64: Optional[str] = Field(None, description="[x-unsloth] Base64-encoded WAV for audio-input models (ASR)")
     use_adapter: Optional[Union[bool, str]] = Field(
         None,
         description=(

studio/backend/models/models.py

@@ -54,6 +54,9 @@ class ModelDetails(BaseModel):
     is_vision: bool = Field(False, description="Whether model is a vision model")
     is_lora: bool = Field(False, description="Whether model is a LoRA adapter")
     is_gguf: bool = Field(False, description="Whether model is a GGUF model (llama.cpp format)")
+    is_audio: bool = Field(False, description="Whether model is a TTS audio model")
+    audio_type: Optional[str] = Field(None, description="Audio codec type: snac, csm, bicodec, dac")
+    has_audio_input: bool = Field(False, description="Whether model accepts audio input (ASR)")
     base_model: Optional[str] = Field(None, description="Base model if this is a LoRA adapter")
 
 

studio/backend/models/training.py

@@ -67,7 +67,8 @@ class TrainingStartRequest(BaseModel):
     finetune_language_layers: bool = Field(False, description="Finetune language layers")
     finetune_attention_modules: bool = Field(False, description="Finetune attention modules")
     finetune_mlp_modules: bool = Field(False, description="Finetune MLP modules")
-    is_dataset_multimodal: bool = Field(False, description="Whether the dataset contains multimodal (image) data")
+    is_dataset_image: bool = Field(False, description="Whether the dataset contains image data")
+    is_dataset_audio: bool = Field(False, description="Whether the dataset contains audio data")
 
     # Logging parameters
     enable_wandb: bool = Field(False, description="Enable Weights & Biases logging")

studio/backend/requirements/extras-no-deps.txt

@@ -8,7 +8,7 @@ snac
 # TRL and related packages
 trl==0.23.1
 git+https://github.com/meta-pytorch/OpenEnv.git
-executorch==1.0.1
+executorch>=1.0.1
 torch-c-dlpack-ext
 sentence_transformers==5.2.0
 transformers==4.57.1

studio/backend/routes/datasets.py

@@ -187,7 +187,7 @@ def check_format(
         # Run lightweight format check on the preview slice
         result = check_dataset_format(preview_slice, is_vlm=request.is_vlm)
 
-        logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_multimodal={result.get('is_multimodal', False)}")
+        logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_image={result.get('is_image', False)}")
 
         # Generate preview samples
         preview_samples = None
@@ -230,11 +230,14 @@ def check_format(
             requires_manual_mapping=result["requires_manual_mapping"],
             detected_format=result["detected_format"],
             columns=result["columns"],
-            is_multimodal=result.get("is_multimodal", False),
+            is_image=result.get("is_image", False),
+            is_audio=result.get("is_audio", False),
             multimodal_columns=result.get("multimodal_columns"),
             suggested_mapping=result.get("suggested_mapping"),
             detected_image_column=result.get("detected_image_column"),
+            detected_audio_column=result.get("detected_audio_column"),
             detected_text_column=result.get("detected_text_column"),
+            detected_speaker_column=result.get("detected_speaker_column"),
             preview_samples=preview_samples,
             total_rows=total_rows,
             warning=warning,

studio/backend/routes/inference.py

@@ -52,6 +52,11 @@ from models.inference import (
 )
 from auth.authentication import get_current_subject
 
+import io
+import wave
+import base64
+import numpy as np
+
 router = APIRouter()
 logger = logging.getLogger(__name__)
 
@@ -241,6 +246,9 @@ async def load_model(
             is_vision=config.is_vision,
             is_lora=config.is_lora,
             is_gguf=False,
+            is_audio=config.is_audio,
+            audio_type=config.audio_type,
+            has_audio_input=config.has_audio_input,
             inference=inference_config,
         )
 
@@ -387,14 +395,23 @@ async def get_status(
         backend = get_inference_backend()
 
         is_vision = False
+        is_audio = False
+        audio_type = None
+        has_audio_input = False
         if backend.active_model_name:
             model_info = backend.models.get(backend.active_model_name, {})
             is_vision = model_info.get("is_vision", False)
+            is_audio = model_info.get("is_audio", False)
+            audio_type = model_info.get("audio_type")
+            has_audio_input = model_info.get("has_audio_input", False)
 
         return InferenceStatusResponse(
             active_model=backend.active_model_name,
             is_vision=is_vision,
             is_gguf=False,
+            is_audio=is_audio,
+            audio_type=audio_type,
+            has_audio_input=has_audio_input,
             loading=list(getattr(backend, 'loading_models', set())),
             loaded=list(backend.models.keys()),
         )
@@ -407,11 +424,118 @@ async def get_status(
         )
 
 
+# =====================================================================
+# Audio (TTS) Generation  (/audio/generate)
+# =====================================================================
+
+
+@router.post("/audio/generate")
+async def generate_audio(payload: ChatCompletionRequest, request: Request):
+    """
+    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.
+    """
+    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"]
+
+    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,
+            ),
+        )
+
+        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))
+
+
 # =====================================================================
 # OpenAI-Compatible Chat Completions  (/chat/completions)
 # =====================================================================
 
 
+def _decode_audio_base64(b64: str) -> np.ndarray:
+    """Decode base64 audio (any format) → float32 numpy array at 16kHz."""
+    import torch
+    import torchaudio
+    import tempfile
+    import os
+
+    raw = base64.b64decode(b64)
+    # torchaudio.load needs a file path or file-like object with format hint
+    # Write to a temp file so torchaudio can auto-detect the format
+    with tempfile.NamedTemporaryFile(suffix=".audio", delete=False) as tmp:
+        tmp.write(raw)
+        tmp_path = tmp.name
+    try:
+        waveform, sr = torchaudio.load(tmp_path)
+    finally:
+        os.unlink(tmp_path)
+
+    # Convert to mono if stereo
+    if waveform.shape[0] > 1:
+        waveform = waveform.mean(dim=0, keepdim=True)
+
+    # Resample to 16kHz if needed
+    if sr != 16000:
+        resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=16000)
+        waveform = resampler(waveform)
+
+    return waveform.squeeze(0).numpy()
+
+
 def _extract_content_parts(
     messages: list,
 ) -> tuple[str, list[dict], "Optional[str]"]:
@@ -501,6 +625,92 @@ 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") 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)
+            system_prompt, chat_messages, _ = _extract_content_parts(payload.messages)
+            cancel_event = threading.Event()
+            completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
+            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,
+                    audio_array=audio_array,
+                    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 512,
+                    repetition_penalty=payload.repetition_penalty,
+                    cancel_event=cancel_event,
+                )
+
+            if payload.stream:
+                async def audio_input_stream():
+                    try:
+                        first_chunk = ChatCompletionChunk(
+                            id=completion_id, created=created, model=model_name,
+                            choices=[ChunkChoice(delta=ChoiceDelta(role="assistant"), finish_reason=None)],
+                        )
+                        yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n"
+
+                        for chunk_text in audio_input_generate():
+                            if await request.is_disconnected():
+                                cancel_event.set()
+                                return
+                            if chunk_text:
+                                chunk = ChatCompletionChunk(
+                                    id=completion_id, created=created, model=model_name,
+                                    choices=[ChunkChoice(delta=ChoiceDelta(content=chunk_text), finish_reason=None)],
+                                )
+                                yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n"
+
+                        final_chunk = ChatCompletionChunk(
+                            id=completion_id, created=created, model=model_name,
+                            choices=[ChunkChoice(delta=ChoiceDelta(), finish_reason="stop")],
+                        )
+                        yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n"
+                        yield "data: [DONE]\n\n"
+                    except asyncio.CancelledError:
+                        cancel_event.set()
+                        raise
+                    except Exception as e:
+                        logger.error(f"Error during audio input streaming: {e}", exc_info=True)
+                        yield f"data: {json.dumps({'error': {'message': str(e), 'type': 'server_error'}})}\n\n"
+
+                return StreamingResponse(
+                    audio_input_stream(),
+                    media_type="text/event-stream",
+                    headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"},
+                )
+            else:
+                full_text = "".join(audio_input_generate())
+                response = ChatCompletion(
+                    id=completion_id, created=created, model=model_name,
+                    choices=[CompletionChoice(message=CompletionMessage(content=full_text), finish_reason="stop")],
+                )
+                return JSONResponse(content=response.model_dump())
+
     # ── Parse messages (handles multimodal content parts) ─────
     system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts(
         payload.messages

studio/backend/routes/models.py

@@ -26,7 +26,7 @@ try:
         list_gguf_variants,
         ModelConfig,
     )
-    from utils.models.model_config import _pick_best_gguf, _extract_quant_label
+    from utils.models.model_config import _pick_best_gguf, _extract_quant_label, is_audio_input_type
     from core.inference import get_inference_backend
 except ImportError:
     # Fallback: try to import from parent directory
@@ -43,7 +43,7 @@ except ImportError:
         list_gguf_variants,
         ModelConfig,
     )
-    from utils.models.model_config import _pick_best_gguf, _extract_quant_label
+    from utils.models.model_config import _pick_best_gguf, _extract_quant_label, is_audio_input_type
     from core.inference import get_inference_backend
 
 from models import (
@@ -225,7 +225,10 @@ async def list_models(
                 id=model_name,
                 name=model_name.split("/")[-1] if "/" in model_name else model_name,
                 is_vision=model_data.get("is_vision", False),
-                is_lora=model_data.get("is_lora", False)
+                is_lora=model_data.get("is_lora", False),
+                is_audio=model_data.get("is_audio", False),
+                audio_type=model_data.get("audio_type"),
+                has_audio_input=model_data.get("has_audio_input", False),
             )
             loaded_models.append(model_info)
         
@@ -265,40 +268,44 @@ async def list_models(
 @router.get("/config/{model_name:path}")
 async def get_model_config(
     model_name: str,
+    hf_token: Optional[str] = Query(None),
     current_subject: str = Depends(get_current_subject),
 ):
     """
     Get configuration for a specific model.
-    
+
     This endpoint wraps the backend load_model_defaults function.
     """
     try:
         logger.info(f"Getting model config for: {model_name}")
+        from utils.models.model_config import detect_audio_type
         # Load model defaults from backend
         config_dict = load_model_defaults(model_name)
 
+        # Detect model capabilities (pass HF token for gated models)
         is_vision = is_vision_model(model_name)
-        
+        audio_type = detect_audio_type(model_name, hf_token=hf_token)
+
         # Check if it's a LoRA adapter
         is_lora = False
         base_model = None
-        
-        # Try to create ModelConfig to get more info
         try:
             model_config = ModelConfig.from_identifier(model_name)
             is_lora = model_config.is_lora
             base_model = model_config.base_model if is_lora else None
         except Exception:
-            # If ModelConfig creation fails, use defaults
             pass
-        
-        logger.info(f"Model config result for {model_name}: is_vision={is_vision}, is_lora={is_lora}, base_model={base_model}")
+
+        logger.info(f"Model config result for {model_name}: is_vision={is_vision}, audio_type={audio_type}, is_lora={is_lora}")
         return ModelDetails(
             id=model_name,
             model_name=model_name,
             config=config_dict,
             is_vision=is_vision,
             is_lora=is_lora,
+            is_audio=audio_type is not None,
+            audio_type=audio_type,
+            has_audio_input=is_audio_input_type(audio_type),
             base_model=base_model,
         )
         

studio/backend/routes/training.py

@@ -184,7 +184,8 @@ async def start_training(
             "finetune_language_layers": request.finetune_language_layers,
             "finetune_attention_modules": request.finetune_attention_modules,
             "finetune_mlp_modules": request.finetune_mlp_modules,
-            "is_dataset_multimodal": request.is_dataset_multimodal,
+            "is_dataset_image": request.is_dataset_image,
+            "is_dataset_audio": request.is_dataset_audio,
             "enable_wandb": request.enable_wandb,
             "wandb_token": request.wandb_token or "",
             "wandb_project": request.wandb_project or "",

studio/backend/utils/datasets/__init__.py

@@ -45,6 +45,7 @@ from .vlm_processing import (
 
 # Data collators
 from .data_collators import (
+    DataCollatorSpeechSeq2SeqWithPadding,
     DeepSeekOCRDataCollator,
     VLMDataCollator,
 )
@@ -85,6 +86,7 @@ __all__ = [
     # VLM
     "generate_smart_vlm_instruction",
     # Collators
+    "DataCollatorSpeechSeq2SeqWithPadding",
     "DeepSeekOCRDataCollator",
     "VLMDataCollator",
     # Mappings

studio/backend/utils/datasets/data_collators.py

@@ -10,6 +10,33 @@ from dataclasses import dataclass
 from typing import Any, List, Optional, Union
 
 
+@dataclass
+class DataCollatorSpeechSeq2SeqWithPadding:
+    """
+    Data collator for Whisper speech-to-text training.
+
+    Pads input features (audio) and label sequences (text) separately,
+    masks padding in labels with -100, and strips leading BOS token.
+    Mirrors the collator from the Whisper.ipynb notebook.
+    """
+    processor: Any
+
+    def __call__(self, features: List[dict]) -> dict:
+        input_features = [{"input_features": feature["input_features"]} for feature in features]
+        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")
+
+        label_features = [{"input_ids": feature["labels"]} for feature in features]
+        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")
+
+        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)
+
+        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
+            labels = labels[:, 1:]
+
+        batch["labels"] = labels
+        return batch
+
+
 @dataclass
 class DeepSeekOCRDataCollator:
     """

studio/backend/utils/datasets/dataset_utils.py

@@ -65,13 +65,22 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
     
     # Auto-detect multimodal data regardless of is_vlm flag
     multimodal_info = detect_multimodal_dataset(dataset)
-    if multimodal_info["is_multimodal"]:
-        is_vlm = True  # Route to VLM detection automatically
-    
+    is_audio = multimodal_info.get("is_audio", False)
+
+    if multimodal_info["is_image"]:
+        is_vlm = True  # Route to VLM detection for image datasets
+
+    # Common audio fields for all return paths
+    audio_fields = {
+        "is_audio": is_audio,
+        "detected_audio_column": multimodal_info.get("detected_audio_column"),
+        "detected_speaker_column": multimodal_info.get("detected_speaker_column"),
+    }
+
     if is_vlm:
         vlm_structure = detect_vlm_dataset_structure(dataset)
         requires_mapping = vlm_structure["format"] == "unknown"
-        
+
         return {
             "requires_manual_mapping": requires_mapping,
             "detected_format": vlm_structure["format"],
@@ -79,53 +88,72 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
             "suggested_mapping": None,
             "detected_image_column": vlm_structure.get("image_column"),
             "detected_text_column": vlm_structure.get("text_column"),
-            "is_multimodal": multimodal_info["is_multimodal"],
+            "is_image": multimodal_info["is_image"],
             "multimodal_columns": multimodal_info.get("multimodal_columns"),
+            **audio_fields,
         }
-    else:
-        # LLM flow
-        detected = detect_dataset_format(dataset)
-        
-        # If format is unknown, try heuristic detection
-        if detected["format"] == "unknown":
-            heuristic_mapping = detect_custom_format_heuristic(dataset)
-            if heuristic_mapping:
-                # Heuristic succeeded - no manual mapping needed
-                return {
-                    "requires_manual_mapping": False,
-                    "detected_format": "custom_heuristic",
-                    "columns": columns,
-                    "suggested_mapping": heuristic_mapping,
-                    "detected_image_column": None,
-                    "detected_text_column": None,
-                    "is_multimodal": False,
-                    "multimodal_columns": None,
-                }
-            else:
-                # Both detection and heuristic failed
-                return {
-                    "requires_manual_mapping": True,
-                    "detected_format": "unknown",
-                    "columns": columns,
-                    "suggested_mapping": None,
-                    "detected_image_column": None,
-                    "detected_text_column": None,
-                    "is_multimodal": False,
-                    "multimodal_columns": None,
-                }
-        
-        # Known format detected
+
+    if is_audio:
+        # Audio dataset — require manual mapping only when columns can't be auto-detected
+        detected_audio = multimodal_info.get("detected_audio_column")
+        detected_text = multimodal_info.get("detected_text_column")
+        needs_mapping = not detected_audio or not detected_text
         return {
-            "requires_manual_mapping": False,
-            "detected_format": detected["format"],
+            "requires_manual_mapping": needs_mapping,
+            "detected_format": "audio",
             "columns": columns,
             "suggested_mapping": None,
             "detected_image_column": None,
-            "detected_text_column": None,
-            "is_multimodal": False,
-            "multimodal_columns": None,
+            "detected_text_column": multimodal_info.get("detected_text_column"),
+            "is_image": False,
+            "multimodal_columns": multimodal_info.get("audio_columns"),
+            **audio_fields,
         }
 
+    # LLM flow
+    detected = detect_dataset_format(dataset)
+
+    # If format is unknown, try heuristic detection
+    if detected["format"] == "unknown":
+        heuristic_mapping = detect_custom_format_heuristic(dataset)
+        if heuristic_mapping:
+            return {
+                "requires_manual_mapping": False,
+                "detected_format": "custom_heuristic",
+                "columns": columns,
+                "suggested_mapping": heuristic_mapping,
+                "detected_image_column": None,
+                "detected_text_column": None,
+                "is_image": False,
+                "multimodal_columns": None,
+                **audio_fields,
+            }
+        else:
+            return {
+                "requires_manual_mapping": True,
+                "detected_format": "unknown",
+                "columns": columns,
+                "suggested_mapping": None,
+                "detected_image_column": None,
+                "detected_text_column": None,
+                "is_image": False,
+                "multimodal_columns": None,
+                **audio_fields,
+            }
+
+    # Known format detected
+    return {
+        "requires_manual_mapping": False,
+        "detected_format": detected["format"],
+        "columns": columns,
+        "suggested_mapping": None,
+        "detected_image_column": None,
+        "detected_text_column": None,
+        "is_image": False,
+        "multimodal_columns": None,
+        **audio_fields,
+    }
+
 # Normalise any format-specific role to canonical chatml (user/assistant/system)
 _TO_CHATML = {
     "user": "user", "human": "user", "instruction": "user",
@@ -250,7 +278,7 @@ def format_dataset(
                 "chat_column": chat_column,
                 "is_standardized": True,
                 "requires_manual_mapping": False,
-                "is_multimodal": multimodal_info["is_multimodal"],
+                "is_image": multimodal_info["is_image"],
                 "multimodal_info": multimodal_info,
                 "warnings": [f"Applied user-provided column mapping ({format_type}): {custom_format_mapping}"]
             }
@@ -262,7 +290,7 @@ def format_dataset(
                 "chat_column": None,
                 "is_standardized": False,
                 "requires_manual_mapping": True,
-                "is_multimodal": multimodal_info["is_multimodal"],
+                "is_image": multimodal_info["is_image"],
                 "multimodal_info": multimodal_info,
                 "warnings": [f"Failed to apply user mapping: {e}"]
             }
@@ -273,7 +301,7 @@ def format_dataset(
     warnings = []
 
     # Add multimodal warning if detected
-    if multimodal_info["is_multimodal"]:
+    if multimodal_info["is_image"]:
         warnings.append(
             f"Multimodal dataset detected. Found columns: {multimodal_info['multimodal_columns']}"
         )
@@ -290,7 +318,7 @@ def format_dataset(
                 "chat_column": None,
                 "is_standardized": True,
                 "requires_manual_mapping": False,
-                "is_multimodal": multimodal_info["is_multimodal"],
+                "is_image": multimodal_info["is_image"],
                 "multimodal_info": multimodal_info,
                 "warnings": []
             }
@@ -310,7 +338,7 @@ def format_dataset(
                     "chat_column": detected["chat_column"],
                     "is_standardized": True,
                     "requires_manual_mapping": False,
-                    "is_multimodal": multimodal_info["is_multimodal"],
+                    "is_image": multimodal_info["is_image"],
                     "multimodal_info": multimodal_info,
                     "warnings": []
                 }
@@ -323,7 +351,7 @@ def format_dataset(
                     "chat_column": detected["chat_column"],
                     "is_standardized": False,
                     "requires_manual_mapping": True,
-                    "is_multimodal": multimodal_info["is_multimodal"],
+                    "is_image": multimodal_info["is_image"],
                     "multimodal_info": multimodal_info,
                     "warnings": warnings
                 }
@@ -336,7 +364,7 @@ def format_dataset(
                 "chat_column": detected["chat_column"],
                 "is_standardized": True,
                 "requires_manual_mapping": False,
-                "is_multimodal": multimodal_info["is_multimodal"],
+                "is_image": multimodal_info["is_image"],
                 "multimodal_info": multimodal_info,
                 "warnings": warnings
             }
@@ -387,7 +415,7 @@ def format_dataset(
                             "chat_column": "conversations",
                             "is_standardized": True,
                             "requires_manual_mapping": False,
-                            "is_multimodal": multimodal_info["is_multimodal"],
+                            "is_image": multimodal_info["is_image"],
                             "multimodal_info": multimodal_info,
                             "warnings": warnings
                         }
@@ -410,7 +438,7 @@ def format_dataset(
                         "chat_column": detected["chat_column"],
                         "is_standardized": True,
                         "requires_manual_mapping": False,
-                        "is_multimodal": multimodal_info["is_multimodal"],
+                        "is_image": multimodal_info["is_image"],
                         "multimodal_info": multimodal_info,
                         "warnings": warnings
                     }
@@ -425,7 +453,7 @@ def format_dataset(
                 "chat_column": detected["chat_column"],
                 "is_standardized": False,
                 "requires_manual_mapping": True,
-                "is_multimodal": multimodal_info["is_multimodal"],
+                "is_image": multimodal_info["is_image"],
                 "multimodal_info": multimodal_info,
                 "warnings": warnings
             }
@@ -441,7 +469,7 @@ def format_dataset(
                 "chat_column": None,
                 "is_standardized": True,
                 "requires_manual_mapping": False,
-                "is_multimodal": multimodal_info["is_multimodal"],
+                "is_image": multimodal_info["is_image"],
                 "multimodal_info": multimodal_info,
                 "warnings": []
             }
@@ -464,7 +492,7 @@ def format_dataset(
                 "chat_column": None,
                 "is_standardized": True,
                 "requires_manual_mapping": False,
-                "is_multimodal": multimodal_info["is_multimodal"],
+                "is_image": multimodal_info["is_image"],
                 "multimodal_info": multimodal_info,
                 "warnings": []
             }
@@ -478,7 +506,7 @@ def format_dataset(
                 "chat_column": detected["chat_column"],
                 "is_standardized": False,
                 "requires_manual_mapping": True,
-                "is_multimodal": multimodal_info["is_multimodal"],
+                "is_image": multimodal_info["is_image"],
                 "multimodal_info": multimodal_info,
                 "warnings": warnings
             }
@@ -495,7 +523,7 @@ def format_dataset(
                 "chat_column": "conversations",
                 "is_standardized": True,
                 "requires_manual_mapping": False,
-                "is_multimodal": multimodal_info["is_multimodal"],
+                "is_image": multimodal_info["is_image"],
                 "multimodal_info": multimodal_info,
                 "warnings": []
             }
@@ -513,7 +541,7 @@ def format_dataset(
                 "chat_column": detected["chat_column"],
                 "is_standardized": True,
                 "requires_manual_mapping": False,
-                "is_multimodal": multimodal_info["is_multimodal"],
+                "is_image": multimodal_info["is_image"],
                 "multimodal_info": multimodal_info,
                 "warnings": []
             }
@@ -526,7 +554,7 @@ def format_dataset(
                 "chat_column": detected["chat_column"],
                 "is_standardized": True,
                 "requires_manual_mapping": False,
-                "is_multimodal": multimodal_info["is_multimodal"],
+                "is_image": multimodal_info["is_image"],
                 "multimodal_info": multimodal_info,
                 "warnings": []
             }
@@ -547,7 +575,7 @@ def format_dataset(
                         "chat_column": detected["chat_column"],
                         "is_standardized": True,
                         "requires_manual_mapping": False,
-                        "is_multimodal": multimodal_info["is_multimodal"],
+                        "is_image": multimodal_info["is_image"],
                         "multimodal_info": multimodal_info,
                         "warnings": warnings
                     }
@@ -561,7 +589,7 @@ def format_dataset(
                 "chat_column": detected["chat_column"],
                 "is_standardized": False,
                 "requires_manual_mapping": True,
-                "is_multimodal": multimodal_info["is_multimodal"],
+                "is_image": multimodal_info["is_image"],
                 "multimodal_info": multimodal_info,
                 "warnings": warnings
             }
@@ -649,7 +677,7 @@ def format_and_template_dataset(
                         "final_format": "vlm_messages",
                         "chat_column": "messages",
                         "is_vlm": True,
-                        "is_multimodal": True,
+                        "is_image": True,
                         "multimodal_info": multimodal_info,
                         "success": True,
                         "requires_manual_mapping": False,
@@ -772,7 +800,7 @@ def format_and_template_dataset(
             "final_format": "vlm_messages",
             "chat_column": "messages",
             "is_vlm": True,
-            "is_multimodal": multimodal_info["is_multimodal"],
+            "is_image": multimodal_info["is_image"],
             "multimodal_info": multimodal_info,
             "vlm_structure": vlm_structure,
             "success": True,
@@ -801,7 +829,7 @@ def format_and_template_dataset(
         # Gemma emits a leading <bos> that must be stripped for text-only chatml/sharegpt.
         is_alpaca = format_type == "alpaca" or (format_type == "auto" and dataset_info["detected_format"] == "alpaca")
         is_gemma = "gemma" in model_name.lower()
-        if is_gemma and not dataset_info["is_multimodal"] and not is_alpaca:
+        if is_gemma and not dataset_info["is_image"] and not is_alpaca:
             remove_bos_prefix = True
         template_result = apply_chat_template_to_dataset(
             dataset_info=dataset_info,

studio/backend/utils/datasets/format_detection.py

@@ -326,45 +326,51 @@ def detect_custom_format_heuristic(dataset):
 
 def detect_multimodal_dataset(dataset):
     """
-    Detects if dataset contains multimodal data (images/vision).
+    Detects if dataset contains multimodal data (images and/or audio).
 
-    Two-pass approach:
-      1. Column-name heuristic (fast): checks for keywords like 'image', 'img', 'pixel'.
-      2. Value-type inspection (reliable): checks if actual values are PIL Images,
-         bytes with image headers, or HF Image-feature dicts.
+    Two-pass approach for each modality:
+      1. Column-name heuristic (fast): checks for keywords.
+      2. Value-type inspection (reliable): checks actual sample values.
 
     Returns:
         dict: {
-            "is_multimodal": bool,
+            "is_image": bool,
             "multimodal_columns": list of column names containing image data,
-            "modality_types": list of detected types (e.g., ["image", "pixel"])
+            "modality_types": list of detected types (e.g., ["image", "audio"]),
+            "is_audio": bool,
+            "audio_columns": list of column names containing audio data,
+            "detected_audio_column": str or None,
+            "detected_text_column": str or None,
         }
     """
     sample = next(iter(dataset))
     column_names = list(sample.keys())
 
-    # Keywords that indicate multimodal/image data
-    multimodal_keywords = [
+    # Keywords that indicate image data
+    image_keywords = [
         'image', 'img', 'pixel',
         'jpg', 'jpeg', 'png', 'webp', 'bmp', 'gif', 'tiff', 'svg',
         'photo', 'pic', 'picture', 'visual',
     ]
 
+    # Keywords that indicate audio data
+    audio_keywords = ['audio', 'speech', 'wav', 'waveform', 'sound']
+
     multimodal_columns = []
+    audio_columns = []
     modality_types = set()
 
-    # ── Pass 1: column-name heuristic ───────────────────────
+    # ── Image detection ─────────────────────────────────────
+    # Pass 1: column-name heuristic
     for col_name in column_names:
         col_lower = col_name.lower()
-
-        for keyword in multimodal_keywords:
+        for keyword in image_keywords:
             if keyword in col_lower:
                 multimodal_columns.append(col_name)
                 modality_types.add(keyword)
-                break  # Don't check other keywords for this column
+                break
 
-    # ── Pass 2: inspect actual values ───────────────────────
-    # Catches columns with non-obvious names (e.g. "jpg", "photo", "pic")
+    # Pass 2: inspect actual values
     already_detected = set(multimodal_columns)
     for col_name in column_names:
         if col_name in already_detected:
@@ -374,10 +380,61 @@ def detect_multimodal_dataset(dataset):
             multimodal_columns.append(col_name)
             modality_types.add("image")
 
+    # ── Audio detection ─────────────────────────────────────
+    # Pass 1: column-name heuristic
+    for col_name in column_names:
+        col_lower = col_name.lower()
+        for keyword in audio_keywords:
+            if keyword in col_lower:
+                audio_columns.append(col_name)
+                modality_types.add("audio")
+                break
+
+    # Pass 2: inspect actual values (catches non-obvious column names)
+    already_audio = set(audio_columns)
+    for col_name in column_names:
+        if col_name in already_audio:
+            continue
+        value = sample[col_name]
+        if _is_audio_value(value):
+            audio_columns.append(col_name)
+            modality_types.add("audio")
+
+    # Filter out columns that are actually audio from the image list
+    # (e.g. a column named "audio" with {"bytes", "path"} could match _is_image_value)
+    if audio_columns:
+        audio_set = set(audio_columns)
+        multimodal_columns = [c for c in multimodal_columns if c not in audio_set]
+
+    # Detect text column for audio datasets
+    detected_text_col = None
+    if audio_columns:
+        text_keywords = ['text', 'sentence', 'transcript', 'transcription', 'label']
+        for col_name in column_names:
+            if col_name.lower() in text_keywords:
+                detected_text_col = col_name
+                break
+
+    is_audio = len(audio_columns) > 0
+
+    # Detect speaker_id column for TTS datasets (CSM, Orpheus, Spark)
+    detected_speaker_col = None
+    if audio_columns:
+        speaker_keywords = ['source', 'speaker', 'speaker_id']
+        for col_name in column_names:
+            if col_name.lower() in speaker_keywords:
+                detected_speaker_col = col_name
+                break
+
     return {
-        "is_multimodal": len(multimodal_columns) > 0,
+        "is_image": len(multimodal_columns) > 0,
         "multimodal_columns": multimodal_columns,
-        "modality_types": list(modality_types)
+        "modality_types": list(modality_types),
+        "is_audio": is_audio,
+        "audio_columns": audio_columns,
+        "detected_audio_column": audio_columns[0] if audio_columns else None,
+        "detected_text_column": detected_text_col,
+        "detected_speaker_column": detected_speaker_col,
     }
 
 
@@ -395,9 +452,16 @@ def _is_image_value(value) -> bool:
         pass
 
     # HF datasets Image feature stores decoded images as PIL or dicts with
-    # {"bytes": b"...", "path": "..."} when not yet decoded
+    # {"bytes": b"...", "path": "..."} when not yet decoded.
+    # Exclude audio dicts (decoded audio has "array" + "sampling_rate").
     if isinstance(value, dict):
+        if "array" in value and "sampling_rate" in value:
+            return False  # This is audio, not image
         if "bytes" in value and "path" in value:
+            # Check path extension to exclude audio files
+            path = value.get("path") or ""
+            if isinstance(path, str) and any(path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS):
+                return False
             return True
 
     # Raw bytes with a known image magic header
@@ -407,6 +471,29 @@ def _is_image_value(value) -> bool:
     return False
 
 
+_AUDIO_EXTENSIONS = (
+    ".wav", ".mp3", ".flac", ".ogg", ".opus", ".m4a", ".aac", ".wma", ".webm",
+)
+
+
+def _is_audio_value(value) -> bool:
+    """Check if a single sample value looks like audio data."""
+    if value is None:
+        return False
+
+    # HF datasets Audio feature: decoded → {"array": np.ndarray, "sampling_rate": int}
+    if isinstance(value, dict):
+        if "array" in value and "sampling_rate" in value:
+            return True
+        # Undecoded/streaming → {"bytes": b"...", "path": "some.wav"}
+        if "bytes" in value or "path" in value:
+            path = value.get("path") or ""
+            if isinstance(path, str) and any(path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS):
+                return True
+
+    return False
+
+
 def _has_image_header(data: bytes) -> bool:
     """Quick magic-byte check for common image formats."""
     if len(data) < 4:

studio/backend/utils/models/__init__.py

@@ -5,6 +5,9 @@ from .model_config import (
     ModelConfig,
     GgufVariantInfo,
     is_vision_model,
+    detect_audio_type,
+    is_audio_input_type,
+    VALID_AUDIO_TYPES,
     scan_trained_loras,
     scan_exported_models,
     load_model_defaults,
@@ -20,6 +23,9 @@ __all__ = [
     'ModelConfig',
     'GgufVariantInfo',
     'is_vision_model',
+    'detect_audio_type',
+    'is_audio_input_type',
+    'VALID_AUDIO_TYPES',
     'scan_trained_loras',
     'scan_exported_models',
     'load_model_defaults',

studio/backend/utils/models/model_config.py

@@ -216,12 +216,18 @@ MODEL_NAME_MAPPING = {
         "unsloth/Nemotron-3-Nano-30B-A3B",
     ],
     "unsloth_orpheus-3b-0.1-ft.yaml": [
+        "unsloth/orpheus-3b-0.1-ft",
         "unsloth/orpheus-3b-0.1-ft-unsloth-bnb-4bit",
         "canopylabs/orpheus-3b-0.1-ft",
         "unsloth/orpheus-3b-0.1-ft-bnb-4bit",
     ],
     "OuteAI_Llama-OuteTTS-1.0-1B.yaml": [
         "OuteAI/Llama-OuteTTS-1.0-1B",
+        "unsloth/Llama-OuteTTS-1.0-1B",
+        "unsloth/llama-outetts-1.0-1b",
+        "OuteAI/OuteTTS-1.0-0.6B",
+        "unsloth/OuteTTS-1.0-0.6B",
+        "unsloth/outetts-1.0-0.6b",
     ],
     "unsloth_PaddleOCR-VL.yaml": [
         "unsloth/PaddleOCR-VL",
@@ -320,9 +326,11 @@ MODEL_NAME_MAPPING = {
     ],
     "sesame_csm-1b.yaml": [
         "sesame/csm-1b",
+        "unsloth/csm-1b",
     ],
     "Spark-TTS-0.5B_LLM.yaml": [
         "Spark-TTS-0.5B/LLM",
+        "unsloth/Spark-TTS-0.5B",
     ],
     "unsloth_tinyllama-bnb-4bit.yaml": [
         "unsloth/tinyllama",
@@ -507,6 +515,13 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
     try:
         config = load_model_config(model_name, use_auth=True, token=hf_token)
 
+        # Exclude audio-only models that share ForConditionalGeneration suffix
+        # (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration)
+        _audio_only_model_types = {'csm', 'whisper'}
+        model_type = getattr(config, 'model_type', None)
+        if model_type in _audio_only_model_types:
+            return False
+
         # Check 1: Architecture class name patterns
         if hasattr(config, 'architectures'):
             is_vlm = any(
@@ -545,6 +560,115 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
         return False
 
 
+VALID_AUDIO_TYPES = ('snac', 'csm', 'bicodec', 'dac', 'whisper', 'audio_vlm')
+
+# Cache detection results per session to avoid repeated API calls
+_audio_detection_cache: Dict[str, Optional[str]] = {}
+
+# Tokenizer token patterns → audio_type (all 6 types detected from tokenizer_config.json)
+_AUDIO_TOKEN_PATTERNS = {
+    'csm': lambda tokens: '<|AUDIO|>' in tokens and '<|audio_eos|>' in tokens,
+    'whisper': lambda tokens: '<|startoftranscript|>' in tokens,
+    'audio_vlm': lambda tokens: '<audio_soft_token>' in tokens,
+    'bicodec': lambda tokens: any(t.startswith('<|bicodec_') for t in tokens),
+    'dac': lambda tokens: '<|audio_start|>' in tokens and '<|audio_end|>' in tokens,
+    'snac': lambda tokens: sum(1 for t in tokens if t.startswith('<custom_token_')) > 10000,
+}
+
+
+def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Optional[str]:
+    """
+    Dynamically detect if a model is an audio model and return its type.
+
+    Fully dynamic — works for any model, not just known ones.
+    Uses tokenizer_config.json special tokens to detect all 6 audio types.
+
+    Returns: audio_type string ('snac', 'csm', 'bicodec', 'dac', 'whisper', 'audio_vlm') or None.
+    """
+    if model_name in _audio_detection_cache:
+        return _audio_detection_cache[model_name]
+
+    result = _detect_audio_from_tokenizer(model_name, hf_token)
+
+    _audio_detection_cache[model_name] = result
+    if result:
+        logger.info(f"Model {model_name} detected as audio model: audio_type={result}")
+    return result
+
+
+def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None) -> Optional[str]:
+    """Detect audio type from tokenizer special tokens (for LLM-based audio models).
+
+    First checks local HF cache, then fetches tokenizer_config.json from HuggingFace.
+    Checks added_tokens_decoder for distinctive patterns.
+    """
+    def _check_token_patterns(tok_config: dict) -> Optional[str]:
+        added = tok_config.get('added_tokens_decoder', {})
+        if not added:
+            return None
+        token_contents = [v.get('content', '') for v in added.values()]
+        for audio_type, check_fn in _AUDIO_TOKEN_PATTERNS.items():
+            if check_fn(token_contents):
+                return audio_type
+        return None
+
+    # 1) Check local HF cache first (works for gated/offline models)
+    try:
+        from huggingface_hub.constants import HF_HUB_CACHE
+        cache_dir = Path(HF_HUB_CACHE)
+        repo_dir_name = f"models--{model_name.replace('/', '--')}"
+        repo_dir = cache_dir / repo_dir_name
+        if repo_dir.exists():
+            snapshots_dir = repo_dir / "snapshots"
+            if snapshots_dir.exists():
+                for snapshot in snapshots_dir.iterdir():
+                    for tok_path in ['tokenizer_config.json', 'LLM/tokenizer_config.json']:
+                        tok_file = snapshot / tok_path
+                        if tok_file.exists():
+                            tok_config = json.loads(tok_file.read_text())
+                            result = _check_token_patterns(tok_config)
+                            if result:
+                                return result
+    except Exception as e:
+        logger.debug(f"Could not check local cache for {model_name}: {e}")
+
+    # 2) Fall back to HuggingFace API
+    try:
+        import requests
+        import os
+
+        paths_to_try = ['tokenizer_config.json', 'LLM/tokenizer_config.json']
+        # Use provided token, or fall back to env
+        token = hf_token or os.environ.get('HF_TOKEN')
+        headers = {}
+        if token:
+            headers['Authorization'] = f'Bearer {token}'
+
+        for tok_path in paths_to_try:
+            url = f"https://huggingface.co/{model_name}/resolve/main/{tok_path}"
+            resp = requests.get(url, headers=headers, timeout=15)
+            if not resp.ok:
+                continue
+
+            tok_config = resp.json()
+            result = _check_token_patterns(tok_config)
+            if result:
+                return result
+
+        return None
+    except Exception as e:
+        logger.debug(f"Could not detect audio type from tokenizer for {model_name}: {e}")
+        return None
+
+
+def is_audio_input_type(audio_type: Optional[str]) -> bool:
+    """Check if an audio_type accepts audio input (ASR/speech understanding).
+
+    Whisper (ASR) and audio_vlm (Gemma3n) accept audio input.
+    """
+    return audio_type in ('whisper', 'audio_vlm')
+
+
 def _is_mmproj(filename: str) -> bool:
     """Check if a GGUF filename is a vision projection (mmproj) file."""
     return "mmproj" in filename.lower()
@@ -1028,6 +1152,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
@@ -1064,6 +1205,9 @@ class ModelConfig:
     is_vision: bool     # Is this a vision model?
     is_lora: bool       # Is this a lora adapter?
     is_gguf: bool = False       # Is this a GGUF model?
+    is_audio: bool = False      # Is this a TTS audio model?
+    audio_type: Optional[str] = None  # Audio codec type: 'snac', 'csm', 'bicodec', 'dac'
+    has_audio_input: bool = False  # Accepts audio input (ASR/speech understanding)
     gguf_file: Optional[str] = None  # Full path to the .gguf file (local mode)
     gguf_mmproj_file: Optional[str] = None  # Full path to the mmproj .gguf file (vision projection)
     gguf_hf_repo: Optional[str] = None  # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF")
@@ -1100,6 +1244,9 @@ class ModelConfig:
             # Check if base model is vision
             is_vision = is_vision_model(base_model, hf_token=hf_token)
 
+            # Check if base model is audio
+            audio_type = detect_audio_type(base_model, hf_token=hf_token)
+
             display_name = lora_path_obj.name
             identifier = lora_path  # Use path as identifier for local LoRAs
 
@@ -1111,6 +1258,9 @@ class ModelConfig:
                 is_cached=True,  # Local LoRAs are always "cached"
                 is_vision=is_vision,
                 is_lora=True,
+                is_audio=audio_type is not None and audio_type != 'audio_vlm',
+                audio_type=audio_type,
+                has_audio_input=is_audio_input_type(audio_type),
                 base_model=base_model,
             )
 
@@ -1288,12 +1438,16 @@ class ModelConfig:
             if not base_model:
                 logger.warning(f"Could not determine base model for LoRA '{path}'")
                 return None
-            vision = is_vision_model(base_model, hf_token=hf_token)
+            check_model = base_model
         else:
-            vision = is_vision_model(identifier, hf_token=hf_token)
-        
+            check_model = identifier
+
+        vision = is_vision_model(check_model, hf_token=hf_token)
+        audio_type_val = detect_audio_type(check_model, hf_token=hf_token)
+        has_audio_in = is_audio_input_type(audio_type_val)
+
         display_name = Path(path).name if is_local else identifier.split("/")[-1]
-        
+
         return cls(
             identifier=identifier,
             display_name=display_name,
@@ -1302,6 +1456,9 @@ class ModelConfig:
             is_cached=is_model_cached(identifier) if not is_local else True,
             is_vision=vision,
             is_lora=is_lora,
+            is_audio=audio_type_val is not None and audio_type_val != 'audio_vlm',
+            audio_type=audio_type_val,
+            has_audio_input=has_audio_in,
             base_model=base_model,
         )
 
@@ -1381,4 +1538,3 @@ class ModelConfig:
             is_lora=is_lora,
             base_model=base_model, # This will be None for base models, and populated for LoRAs
         )
-    pass

studio/frontend/package.json

@@ -32,10 +32,10 @@
     "@radix-ui/react-select": "^2.2.6",
     "@radix-ui/react-separator": "^1.1.8",
     "@radix-ui/react-slot": "^1.2.4",
-    "@streamdown/cjk": "^1.0.2",
-    "@streamdown/code": "^1.0.2",
-    "@streamdown/math": "^1.0.2",
-    "@streamdown/mermaid": "^1.0.2",
+    "@streamdown/cjk": "1.0.2",
+    "@streamdown/code": "1.0.2",
+    "@streamdown/math": "1.0.2",
+    "@streamdown/mermaid": "1.0.2",
     "@tailwindcss/vite": "^4.1.18",
     "@tanstack/react-router": "^1.159.10",
     "@tanstack/react-table": "^8.21.3",
@@ -66,7 +66,7 @@
     "remark-gfm": "^4.0.1",
     "shadcn": "^3.8.4",
     "sonner": "^2.0.7",
-    "streamdown": "^2.3.0",
+    "streamdown": "2.3.0",
     "tailwind-merge": "^3.4.0",
     "tailwindcss": "^4.1.18",
     "tw-animate-css": "^1.4.0",

studio/frontend/src/components/assistant-ui/audio-player.tsx

@@ -0,0 +1,114 @@
+"use client";
+
+import { Button } from "@/components/ui/button";
+import { DownloadIcon, PauseIcon, PlayIcon } from "lucide-react";
+import { type FC, useRef, useState } from "react";
+
+interface AudioPlayerProps {
+  src: string;
+}
+
+export const AudioPlayer: FC<AudioPlayerProps> = ({ src }) => {
+  const audioRef = useRef<HTMLAudioElement>(null);
+  const [isPlaying, setIsPlaying] = useState(false);
+  const [progress, setProgress] = useState(0);
+  const [duration, setDuration] = useState(0);
+
+  const togglePlay = () => {
+    const audio = audioRef.current;
+    if (!audio) return;
+    if (isPlaying) {
+      audio.pause();
+    } else {
+      audio.play();
+    }
+    setIsPlaying(!isPlaying);
+  };
+
+  const handleTimeUpdate = () => {
+    const audio = audioRef.current;
+    if (!audio) return;
+    setProgress(audio.currentTime);
+  };
+
+  const handleLoadedMetadata = () => {
+    const audio = audioRef.current;
+    if (!audio) return;
+    setDuration(audio.duration);
+  };
+
+  const handleEnded = () => {
+    setIsPlaying(false);
+    setProgress(0);
+  };
+
+  const handleSeek = (e: React.ChangeEvent<HTMLInputElement>) => {
+    const audio = audioRef.current;
+    if (!audio) return;
+    const time = parseFloat(e.target.value);
+    audio.currentTime = time;
+    setProgress(time);
+  };
+
+  const handleDownload = () => {
+    const link = document.createElement("a");
+    link.href = src;
+    link.download = "generated-audio.wav";
+    link.click();
+  };
+
+  const formatTime = (t: number) => {
+    const mins = Math.floor(t / 60);
+    const secs = Math.floor(t % 60);
+    return `${mins}:${secs.toString().padStart(2, "0")}`;
+  };
+
+  return (
+    <div className="my-2 flex max-w-md items-center gap-3 rounded-xl border bg-muted/50 px-4 py-3">
+      <audio
+        ref={audioRef}
+        src={src}
+        onTimeUpdate={handleTimeUpdate}
+        onLoadedMetadata={handleLoadedMetadata}
+        onEnded={handleEnded}
+        preload="metadata"
+      />
+      <Button
+        variant="ghost"
+        size="icon"
+        className="size-8 shrink-0 rounded-full"
+        onClick={togglePlay}
+      >
+        {isPlaying ? (
+          <PauseIcon className="size-4" />
+        ) : (
+          <PlayIcon className="size-4" />
+        )}
+      </Button>
+      <div className="flex flex-1 flex-col gap-1">
+        <input
+          type="range"
+          min={0}
+          max={duration || 0}
+          step={0.01}
+          value={progress}
+          onChange={handleSeek}
+          className="h-1.5 w-full cursor-pointer accent-primary"
+        />
+        <div className="flex justify-between text-[10px] text-muted-foreground">
+          <span>{formatTime(progress)}</span>
+          <span>{formatTime(duration)}</span>
+        </div>
+      </div>
+      <Button
+        variant="ghost"
+        size="icon"
+        className="size-7 shrink-0 text-muted-foreground"
+        onClick={handleDownload}
+        title="Download audio"
+      >
+        <DownloadIcon className="size-3.5" />
+      </Button>
+    </div>
+  );
+};

studio/frontend/src/components/assistant-ui/markdown-text.tsx

@@ -10,6 +10,7 @@ import { mermaid } from "@streamdown/mermaid";
 import { Block, type BlockProps, Streamdown } from "streamdown";
 import { useEffect, useRef, useState } from "react";
 import "katex/dist/katex.min.css";
+import { AudioPlayer } from "./audio-player";
 
 const { withSmoothContextProvider, useSmoothStatus } = INTERNAL;
 
@@ -77,11 +78,17 @@ function StreamdownBlock(props: BlockProps) {
 
   return <Block {...props} />;
 }
+const AUDIO_PLAYER_RE = /<audio-player\s+src="([^"]+)"\s*\/>/;
 
 const MarkdownTextImpl = () => {
   const { text } = useMessagePartText();
   const status = useSmoothStatus();
 
+  const audioMatch = text.match(AUDIO_PLAYER_RE);
+  if (audioMatch) {
+    return <AudioPlayer src={audioMatch[1]} />;
+  }
+
   return (
     <div data-status={status.type}>
       <Streamdown

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 {
@@ -33,13 +35,16 @@ import {
   ChevronRightIcon,
   CopyIcon,
   DownloadIcon,
+  HeadphonesIcon,
   MicIcon,
   MoreHorizontalIcon,
   PencilIcon,
   RefreshCwIcon,
   SquareIcon,
+  XIcon,
 } from "lucide-react";
-import { type FC, useRef, useState } from "react";
+import { type FC, useCallback, useRef, useState } from "react";
+import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
 
 export const Thread: FC<{ hideComposer?: boolean; hideWelcome?: boolean }> = ({
   hideComposer,
@@ -162,11 +167,34 @@ const ComposerAnimated: FC = () => {
   );
 };
 
+const PendingAudioChip: FC = () => {
+  const audioName = useChatRuntimeStore((s) => s.pendingAudioName);
+  const clearPendingAudio = useChatRuntimeStore((s) => s.clearPendingAudio);
+  if (!audioName) return null;
+  return (
+    <div className="mb-2 flex w-full flex-row items-center gap-2 px-1.5 pt-0.5 pb-1">
+      <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>
+        <button
+          type="button"
+          onClick={clearPendingAudio}
+          className="flex size-4 items-center justify-center rounded-full hover:bg-destructive hover:text-destructive-foreground"
+          aria-label="Remove audio"
+        >
+          <XIcon className="size-3" />
+        </button>
+      </div>
+    </div>
+  );
+};
+
 const Composer: FC = () => {
   return (
     <ComposerPrimitive.Root className="aui-composer-root relative flex w-full flex-col">
       <ComposerPrimitive.AttachmentDropzone className="aui-composer-attachment-dropzone shadow-border ring-1 ring-border flex w-full flex-col rounded-2xl bg-background px-1 pt-2 outline-none transition-shadow data-[dragging=true]:ring-ring data-[dragging=true]:bg-accent/50">
         <ComposerAttachments />
+        <PendingAudioChip />
         <ComposerPrimitive.Input
           placeholder="Send a message..."
           className="aui-composer-input mb-1 max-h-32 min-h-12 w-full resize-none bg-transparent px-4 pt-2 pb-3 text-sm outline-none placeholder:text-muted-foreground focus-visible:ring-0"
@@ -180,10 +208,64 @@ const Composer: FC = () => {
   );
 };
 
+const ComposerAudioUpload: FC = () => {
+  const audioInputRef = useRef<HTMLInputElement>(null);
+  const setPendingAudio = useChatRuntimeStore((s) => s.setPendingAudio);
+  const activeModel = useChatRuntimeStore((s) => {
+    const checkpoint = s.params.checkpoint;
+    return s.models.find((m) => m.id === checkpoint);
+  });
+
+  const handleAudioFile = useCallback(
+    async (file: File) => {
+      if (file.size > MAX_AUDIO_SIZE) return;
+      try {
+        const base64 = await fileToBase64(file);
+        setPendingAudio(base64, file.name);
+      } catch {
+        // skip
+      }
+    },
+    [setPendingAudio],
+  );
+
+  if (!activeModel?.hasAudioInput) return null;
+
+  return (
+    <>
+      <input
+        ref={audioInputRef}
+        type="file"
+        accept={AUDIO_ACCEPT}
+        className="hidden"
+        onChange={(e) => {
+          const file = e.target.files?.[0];
+          if (file) handleAudioFile(file);
+          e.target.value = "";
+        }}
+      />
+      <TooltipIconButton
+        tooltip="Upload audio"
+        side="bottom"
+        variant="ghost"
+        size="icon"
+        className="size-8.5 rounded-full p-1 text-muted-foreground hover:bg-muted-foreground/15"
+        onClick={() => audioInputRef.current?.click()}
+        aria-label="Upload audio"
+      >
+        <HeadphonesIcon className="size-4.5 stroke-[1.5px]" />
+      </TooltipIconButton>
+    </>
+  );
+};
+
 const ComposerAction: FC = () => {
   return (
     <div className="aui-composer-action-wrapper relative mx-2 mb-2 flex items-center justify-between">
-      <ComposerAddAttachment />
+      <div className="flex items-center gap-1">
+        <ComposerAddAttachment />
+        <ComposerAudioUpload />
+      </div>
       <div className="flex items-center gap-1">
         <ComposerPrimitive.If dictation={false}>
           <ComposerPrimitive.Dictate asChild={true}>
@@ -342,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
@@ -349,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">

studio/frontend/src/features/chat/api/chat-adapter.ts

@@ -1,6 +1,6 @@
 import type { ChatModelAdapter } from "@assistant-ui/react";
 import { toast } from "sonner";
-import { streamChatCompletions } from "./chat-api";
+import { generateAudio, streamChatCompletions } from "./chat-api";
 import { db } from "../db";
 import { useChatRuntimeStore } from "../stores/chat-runtime-store";
 import {
@@ -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")
@@ -92,6 +95,26 @@ function findLatestUserImageBase64(messages: RunMessages): string | undefined {
   return undefined;
 }
 
+function findLatestUserAudioBase64(messages: RunMessages): string | undefined {
+  // Check message content parts (from compare view's CompareMessagePart with type: "audio")
+  for (let i = messages.length - 1; i >= 0; i -= 1) {
+    const message = messages[i];
+    if (!message || message.role !== "user") continue;
+
+    for (const part of message.content ?? []) {
+      if (part.type === "audio" && "audio" in part) {
+        const audioPart = (part as unknown as { type: "audio"; audio: string | { data: string; format: string } }).audio;
+        const raw = typeof audioPart === "string" ? audioPart : audioPart?.data;
+        if (raw) return raw.startsWith("data:") ? raw.split(",")[1] : raw;
+      }
+    }
+  }
+
+  // Check the runtime store (from main composer's audio upload)
+  const pendingAudio = useChatRuntimeStore.getState().pendingAudioBase64;
+  return pendingAudio ?? undefined;
+}
+
 async function resolveUseAdapter(
   threadId: string | undefined,
 ): Promise<boolean | undefined> {
@@ -135,8 +158,69 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
         });
       }
       const imageBase64 = findLatestUserImageBase64(messages);
+      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);
 
+      // ── Audio model path (non-streaming) ─────────────────────
+      const activeModel = runtime.models.find(
+        (m) => m.id === params.checkpoint,
+      );
+      if (activeModel?.isAudio && !activeModel?.hasAudioInput) {
+        const threadKey = unstable_threadId || "__default";
+        runtime.setThreadRunning(threadKey, true);
+        try {
+          yield {
+            content: [{ type: "text" as const, text: "Generating audio..." }],
+          };
+
+          const result = await generateAudio(
+            {
+              model: params.checkpoint,
+              messages: outboundMessages,
+              stream: false,
+              temperature: params.temperature,
+              top_p: params.topP,
+              max_tokens: params.maxTokens,
+              top_k: params.topK,
+              min_p: params.minP,
+              repetition_penalty: params.repetitionPenalty,
+              ...(useAdapter === undefined ? {} : { use_adapter: useAdapter }),
+            },
+            abortSignal,
+          );
+
+          const audioUrl = `data:audio/wav;base64,${result.audio.data}`;
+          yield {
+            content: [
+              {
+                type: "text" as const,
+                text: `<audio-player src="${audioUrl}" />`,
+              },
+            ],
+          };
+        } catch (err) {
+          if (!abortSignal.aborted) {
+            toast.error("Audio generation failed", {
+              description:
+                err instanceof Error ? err.message : "Unknown error",
+            });
+          }
+          throw err;
+        } finally {
+          runtime.setThreadRunning(threadKey, false);
+        }
+        return;
+      }
+
       const threadKey = unstable_threadId || "__default";
       let waitingFirstChunk = true;
       let firstTokenSettled = false;
@@ -194,6 +278,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
             min_p: params.minP,
             repetition_penalty: params.repetitionPenalty,
             image_base64: imageBase64,
+            audio_base64: audioBase64,
             ...(useAdapter === undefined ? {} : { use_adapter: useAdapter }),
           },
           abortSignal,

studio/frontend/src/features/chat/api/chat-api.ts

@@ -1,5 +1,6 @@
 import { authFetch } from "@/features/auth";
 import type {
+  AudioGenerationResponse,
   GgufVariantsResponse,
   InferenceStatusResponse,
   ListLorasResponse,
@@ -155,3 +156,22 @@ export async function* streamChatCompletions(
     }
   }
 }
+
+export async function generateAudio(
+  payload: OpenAIChatCompletionsRequest,
+  signal: AbortSignal,
+): Promise<AudioGenerationResponse> {
+  const response = await authFetch("/api/inference/chat/completions", {
+    method: "POST",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify({ ...payload, stream: false }),
+    signal,
+  });
+
+  if (!response.ok) {
+    const body = await response.json().catch(() => null);
+    throw new Error(parseErrorText(response.status, body));
+  }
+
+  return (await response.json()) as AudioGenerationResponse;
+}

studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts

@@ -44,12 +44,17 @@ function describeModel(model: {
   is_lora?: boolean;
   is_vision?: boolean;
   is_gguf?: boolean;
+  is_audio?: boolean;
+  has_audio_input?: boolean;
 }): string | undefined {
   const tags: string[] = [];
   if (model.is_gguf) tags.push("GGUF");
   if (model.is_lora) tags.push("LoRA");
   if (model.is_vision) tags.push("Vision");
-  if (!model.is_lora && !model.is_vision && !model.is_gguf) tags.push("Base");
+  if (model.is_audio) tags.push("Audio");
+  if (model.has_audio_input) tags.push("Audio Input");
+  if (!model.is_lora && !model.is_vision && !model.is_gguf && !model.is_audio && !model.has_audio_input)
+    tags.push("Base");
   return tags.join(" · ");
 }
 
@@ -59,6 +64,9 @@ function toChatModelSummary(model: {
   is_lora?: boolean;
   is_vision?: boolean;
   is_gguf?: boolean;
+  is_audio?: boolean;
+  audio_type?: string | null;
+  has_audio_input?: boolean;
 }): ChatModelSummary {
   return {
     id: model.id,
@@ -67,6 +75,9 @@ function toChatModelSummary(model: {
     isLora: Boolean(model.is_lora),
     isVision: Boolean(model.is_vision),
     isGguf: Boolean(model.is_gguf),
+    isAudio: Boolean(model.is_audio),
+    audioType: model.audio_type ?? null,
+    hasAudioInput: Boolean(model.has_audio_input),
   };
 }
 

studio/frontend/src/features/chat/shared-composer.tsx

@@ -1,7 +1,9 @@
 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, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react";
+import { ArrowUpIcon, HeadphonesIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react";
+import { useChatRuntimeStore } from "./stores/chat-runtime-store";
 import {
   type KeyboardEvent,
   type MutableRefObject,
@@ -17,7 +19,8 @@ import {
 
 export type CompareMessagePart =
   | { type: "text"; text: string }
-  | { type: "image"; image: string };
+  | { type: "image"; image: string }
+  | { type: "audio"; audio: string };
 
 export interface CompareHandle {
   append: (content: CompareMessagePart[]) => void;
@@ -182,9 +185,18 @@ export function SharedComposer({
   const [text, setText] = useState("");
   const [running, setRunning] = useState(false);
   const [pendingImages, setPendingImages] = useState<PendingImage[]>([]);
+  const [pendingAudio, setPendingAudio] = useState<{ name: string; base64: string } | null>(null);
   const [dragging, setDragging] = useState(false);
   const textareaRef = useRef<HTMLTextAreaElement>(null);
   const fileInputRef = useRef<HTMLInputElement>(null);
+  const audioInputRef = useRef<HTMLInputElement>(null);
+
+  const activeModel = useChatRuntimeStore((s) => {
+    const checkpoint = s.params.checkpoint;
+    return s.models.find((m) => m.id === checkpoint);
+  });
+  const setPendingAudioStore = useChatRuntimeStore((s) => s.setPendingAudio);
+  const clearPendingAudioStore = useChatRuntimeStore((s) => s.clearPendingAudio);
 
   const { isDictating, start: startDictation, stop: stopDictation, supported: dictationSupported } = useDictation(
     setText,
@@ -204,12 +216,22 @@ export function SharedComposer({
     const next: PendingImage[] = [];
     for (let i = 0; i < files.length; i++) {
       const file = files[i];
-      if (!file?.type.match(/^image\/(jpeg|png|webp|gif)$/i)) continue;
+      if (!file) continue;
+      // Handle audio files
+      if (file.type.match(/^audio\//i) && file.size <= MAX_AUDIO_SIZE) {
+        fileToBase64(file).then((base64) => {
+          setPendingAudio({ name: file.name, base64 });
+          setPendingAudioStore(base64, file.name);
+        });
+        continue;
+      }
+      // Handle image files
+      if (!file.type.match(/^image\/(jpeg|png|webp|gif)$/i)) continue;
       if (file.size > MAX_IMAGE_SIZE) continue;
       next.push({ id: crypto.randomUUID(), file });
     }
     setPendingImages((prev) => [...prev, ...next]);
-  }, []);
+  }, [setPendingAudioStore]);
 
   const removePendingImage = useCallback((id: string) => {
     setPendingImages((prev) => prev.filter((p) => p.id !== id));
@@ -217,7 +239,7 @@ export function SharedComposer({
 
   async function send() {
     const msg = text.trim();
-    if (!msg && pendingImages.length === 0) return;
+    if (!msg && pendingImages.length === 0 && !pendingAudio) return;
 
     const content: CompareMessagePart[] = [];
     for (const { file } of pendingImages) {
@@ -228,6 +250,9 @@ export function SharedComposer({
         // skip failed image
       }
     }
+    if (pendingAudio) {
+      content.push({ type: "audio", audio: pendingAudio.base64 });
+    }
     if (msg) {
       content.push({ type: "text", text: msg });
     }
@@ -238,6 +263,8 @@ export function SharedComposer({
     }
     setText("");
     setPendingImages([]);
+    setPendingAudio(null);
+    clearPendingAudioStore();
     textareaRef.current?.focus();
   }
 
@@ -257,7 +284,7 @@ export function SharedComposer({
     }
   }
 
-  const canSend = (text.trim().length > 0 || pendingImages.length > 0) && !running;
+  const canSend = (text.trim().length > 0 || pendingImages.length > 0 || pendingAudio !== null) && !running;
 
   return (
     <div
@@ -273,7 +300,7 @@ export function SharedComposer({
         addFiles(e.dataTransfer.files);
       }}
     >
-      {pendingImages.length > 0 && (
+      {(pendingImages.length > 0 || pendingAudio) && (
         <div className="mb-2 flex w-full flex-row flex-wrap items-center gap-2 px-1.5 pt-0.5 pb-1">
           {pendingImages.map(({ id, file }) => (
             <PendingImageThumb
@@ -282,6 +309,20 @@ export function SharedComposer({
               onRemove={() => removePendingImage(id)}
             />
           ))}
+          {pendingAudio && (
+            <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">{pendingAudio.name}</span>
+              <button
+                type="button"
+                onClick={() => { setPendingAudio(null); clearPendingAudioStore(); }}
+                className="flex size-4 items-center justify-center rounded-full hover:bg-destructive hover:text-destructive-foreground"
+                aria-label="Remove audio"
+              >
+                <XIcon className="size-3" />
+              </button>
+            </div>
+          )}
         </div>
       )}
       <textarea
@@ -317,6 +358,31 @@ export function SharedComposer({
           >
             <PlusIcon className="size-5 stroke-[1.5px]" />
           </TooltipIconButton>
+          {activeModel?.hasAudioInput && (
+            <>
+              <input
+                ref={audioInputRef}
+                type="file"
+                accept={AUDIO_ACCEPT}
+                className="hidden"
+                onChange={(e) => {
+                  addFiles(e.target.files);
+                  e.target.value = "";
+                }}
+              />
+              <TooltipIconButton
+                tooltip="Upload audio"
+                side="bottom"
+                variant="ghost"
+                size="icon"
+                className="size-8 rounded-full text-muted-foreground hover:bg-muted-foreground/15"
+                onClick={() => audioInputRef.current?.click()}
+                aria-label="Upload audio"
+              >
+                <HeadphonesIcon className="size-4 stroke-[1.5px]" />
+              </TooltipIconButton>
+            </>
+          )}
         </div>
         <div className="flex items-center gap-1">
           {dictationSupported && (

studio/frontend/src/features/chat/stores/chat-runtime-store.ts

@@ -40,6 +40,8 @@ type ChatRuntimeStore = {
   autoTitle: boolean;
   modelsError: string | null;
   activeGgufVariant: string | null;
+  pendingAudioBase64: string | null;
+  pendingAudioName: string | null;
   setParams: (params: InferenceParams) => void;
   setModels: (models: ChatModelSummary[]) => void;
   setLoras: (loras: ChatLoraSummary[]) => void;
@@ -48,6 +50,8 @@ type ChatRuntimeStore = {
   setModelsError: (error: string | null) => void;
   setCheckpoint: (modelId: string, ggufVariant?: string | null) => void;
   clearCheckpoint: () => void;
+  setPendingAudio: (base64: string, name: string) => void;
+  clearPendingAudio: () => void;
 };
 
 export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
@@ -58,6 +62,8 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
   autoTitle: loadBool(AUTO_TITLE_KEY, false),
   modelsError: null,
   activeGgufVariant: null,
+  pendingAudioBase64: null,
+  pendingAudioName: null,
   setParams: (params) => set({ params }),
   setModels: (models) => set({ models }),
   setLoras: (loras) => set({ loras }),
@@ -93,4 +99,8 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set) => ({
       },
       activeGgufVariant: null,
     })),
+  setPendingAudio: (base64, name) =>
+    set({ pendingAudioBase64: base64, pendingAudioName: name }),
+  clearPendingAudio: () =>
+    set({ pendingAudioBase64: null, pendingAudioName: null }),
 }));

studio/frontend/src/features/chat/types/api.ts

@@ -4,6 +4,9 @@ export interface BackendModelDetails {
   is_vision?: boolean;
   is_lora?: boolean;
   is_gguf?: boolean;
+  is_audio?: boolean;
+  audio_type?: string | null;
+  has_audio_input?: boolean;
 }
 
 export interface ListModelsResponse {
@@ -53,6 +56,9 @@ export interface LoadModelResponse {
   is_vision: boolean;
   is_lora: boolean;
   is_gguf?: boolean;
+  is_audio?: boolean;
+  audio_type?: string | null;
+  has_audio_input?: boolean;
   inference?: {
     temperature?: number;
     top_p?: number;
@@ -70,10 +76,29 @@ export interface InferenceStatusResponse {
   is_vision: boolean;
   is_gguf?: boolean;
   gguf_variant?: string | null;
+  is_audio?: boolean;
+  audio_type?: string | null;
+  has_audio_input?: boolean;
   loading: string[];
   loaded: string[];
 }
 
+export interface AudioGenerationResponse {
+  id: string;
+  object: string;
+  model: string;
+  audio: {
+    data: string;
+    format: string;
+    sample_rate: number;
+  };
+  choices: Array<{
+    index: number;
+    message: { role: string; content: string };
+    finish_reason: string;
+  }>;
+}
+
 export interface OpenAIChatMessage {
   role: "system" | "user" | "assistant";
   content: string;
@@ -90,6 +115,7 @@ export interface OpenAIChatCompletionsRequest {
   min_p: number;
   repetition_penalty: number;
   image_base64?: string;
+  audio_base64?: string;
   use_adapter?: boolean | string | null;
 }
 

studio/frontend/src/features/chat/types/runtime.ts

@@ -27,6 +27,9 @@ export interface ChatModelSummary {
   isVision: boolean;
   isLora: boolean;
   isGguf?: boolean;
+  isAudio?: boolean;
+  audioType?: string | null;
+  hasAudioInput?: boolean;
 }
 
 export interface ChatLoraSummary {

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

@@ -16,6 +16,7 @@ const CHATML_ROLES = ["system", "user", "assistant"] as const;
 const ALPACA_ROLES = ["instruction", "input", "output"] as const;
 const SHAREGPT_ROLES = ["system", "human", "gpt"] as const;
 const VLM_ROLES = ["image", "text"] as const;
+const AUDIO_ROLES = ["audio", "text", "speaker_id"] as const;
 
 const ROLE_LABELS: Record<string, string> = {
   system: "System",
@@ -28,9 +29,12 @@ const ROLE_LABELS: Record<string, string> = {
   output: "Output",
   image: "Image",
   text: "Text",
+  audio: "Audio",
+  speaker_id: "Speaker ID",
 };
 
-export function getAvailableRoles(isVlm: boolean, format?: string): readonly string[] {
+export function getAvailableRoles(isVlm: boolean, format?: string, isAudio?: boolean): readonly string[] {
+  if (isAudio) return AUDIO_ROLES;
   if (isVlm) return VLM_ROLES;
   if (format === "alpaca") return ALPACA_ROLES;
   if (format === "sharegpt") return SHAREGPT_ROLES;
@@ -41,8 +45,10 @@ export function isMappingComplete(
   mapping: Record<string, string>,
   isVlm: boolean,
   format?: string,
+  isAudio?: boolean,
 ): boolean {
   const roles = new Set(Object.values(mapping));
+  if (isAudio) return roles.has("audio") && roles.has("text");
   if (isVlm) return roles.has("image") && roles.has("text");
   if (format === "alpaca") return roles.has("instruction") && roles.has("output");
   if (format === "sharegpt") return roles.has("human") && roles.has("gpt");
@@ -85,22 +91,26 @@ export function DatasetMappingCard({
   mappingOk,
   autoDetected = false,
   isVlm = false,
+  isAudio = false,
   format,
 }: {
   mapping: Record<string, string>;
   mappingOk: boolean;
   autoDetected?: boolean;
   isVlm?: boolean;
+  isAudio?: boolean;
   format?: string;
 }) {
   const entries = Object.entries(mapping);
-  const requiredLabel = isVlm
-    ? "image and text"
-    : format === "alpaca"
-      ? "instruction and output"
-      : format === "sharegpt"
-        ? "human and gpt"
-        : "user and assistant";
+  const requiredLabel = isAudio
+    ? "audio and text"
+    : isVlm
+      ? "image and text"
+      : format === "alpaca"
+        ? "instruction and output"
+        : format === "sharegpt"
+          ? "human and gpt"
+          : "user and assistant";
 
   return (
     <div
@@ -228,6 +238,7 @@ const TO_CANONICAL: Record<string, string> = {
   instruction: "user", input: "system", output: "assistant",
   human: "user", gpt: "assistant",
   image: "image", text: "text",
+  audio: "audio", speaker_id: "speaker_id",
 };
 
 /** Chatml → format-specific role names (only for formats that differ). */
@@ -257,10 +268,18 @@ export function deriveDefaultMapping(
   data: CheckFormatResponse,
   isVlm: boolean,
   format?: string,
+  isAudio?: boolean,
 ): Record<string, string> {
   if (data.suggested_mapping) {
     return remapRolesForFormat({ ...data.suggested_mapping }, format);
   }
+  if (isAudio) {
+    const result: Record<string, string> = {};
+    if (data.detected_audio_column) result[data.detected_audio_column] = "audio";
+    if (data.detected_text_column) result[data.detected_text_column] = "text";
+    if (data.detected_speaker_column) result[data.detected_speaker_column] = "speaker_id";
+    return result;
+  }
   if (isVlm) {
     const result: Record<string, string> = {};
     if (data.detected_image_column) result[data.detected_image_column] = "image";

studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx

@@ -62,15 +62,16 @@ export function DatasetPreviewDialog({
   );
   const { isStarting, startError, startTrainingRun } = useTrainingActions();
 
-  // If the backend reports multimodal data, treat as VLM even if the prop
-  // hasn't caught up yet (isDatasetMultimodal may still be null in the store).
-  const effectiveIsVlm = isVlm || !!data?.is_multimodal;
+  // If the backend reports image data, treat as VLM even if the prop
+  // hasn't caught up yet (isDatasetImage may still be null in the store).
+  const effectiveIsAudio = !!data?.is_audio;
+  const effectiveIsVlm = isVlm || !!data?.is_image;
 
   const hasHeuristicMapping = !data?.requires_manual_mapping && !!data?.suggested_mapping;
   const mappingEnabled = !!data?.requires_manual_mapping || hasHeuristicMapping;
   const showMappingFooter = mode === "mapping" && mappingEnabled;
-  const mappingOk = isMappingComplete(manualMapping, effectiveIsVlm, datasetFormat);
-  const availableRoles = getAvailableRoles(effectiveIsVlm, datasetFormat);
+  const mappingOk = isMappingComplete(manualMapping, effectiveIsVlm, datasetFormat, effectiveIsAudio);
+  const availableRoles = getAvailableRoles(effectiveIsVlm, datasetFormat, effectiveIsAudio);
   const isHfDataset = !!datasetName && datasetName.includes("/");
 
   // When format changes, remap existing mapping roles to the new format's role names
@@ -150,7 +151,7 @@ export function DatasetPreviewDialog({
     if (!data?.requires_manual_mapping && !data?.suggested_mapping) return;
     // Don't overwrite if mapping already has entries
     if (Object.keys(manualMapping).length > 0) return;
-    const derived = deriveDefaultMapping(data, effectiveIsVlm, datasetFormat);
+    const derived = deriveDefaultMapping(data, effectiveIsVlm, datasetFormat, effectiveIsAudio);
     if (Object.keys(derived).length === 0) return;
     setManualMapping(derived);
   }, [open, datasetName, data, effectiveIsVlm, datasetFormat, manualMapping, setManualMapping]);
@@ -353,6 +354,7 @@ export function DatasetPreviewDialog({
                   mappingOk={mappingOk}
                   autoDetected={hasHeuristicMapping}
                   isVlm={effectiveIsVlm}
+                  isAudio={effectiveIsAudio}
                   format={datasetFormat}
                 />
               )}

studio/frontend/src/features/studio/sections/params-section.tsx

@@ -114,7 +114,7 @@ function SliderRow({
 export function ParamsSection(): ReactElement {
   const store = useTrainingConfigStore();
   const isLora = store.trainingMethod !== "full";
-  const showVisionLora = store.isVisionModel && store.isDatasetMultimodal === true;
+  const showVisionLora = store.isVisionModel && store.isDatasetImage === true;
   const [loraOpen, setLoraOpen] = useState(false);
   const [hyperOpen, setHyperOpen] = useState(false);
   const maxStepsSliderMax = Math.max(500, store.maxSteps, 30);

studio/frontend/src/features/studio/sections/training-section.tsx

@@ -42,8 +42,8 @@ export function TrainingSection() {
   const store = useTrainingConfigStore();
   const { isStarting, startError, startTrainingRun } = useTrainingActions();
   const isIncompatible =
-    !store.isVisionModel && store.isDatasetMultimodal === true;
-  const fileInputRef = useRef<HTMLInputElement>(null);
+    !store.isVisionModel && store.isDatasetImage === true;
+    const fileInputRef = useRef<HTMLInputElement>(null);
 
   const handleFileUpload = (e: React.ChangeEvent<HTMLInputElement>) => {
     const file = e.target.files?.[0];

studio/frontend/src/features/studio/studio-page.tsx

@@ -93,7 +93,7 @@ export function StudioPage(): ReactElement {
           datasetSplit={config.datasetSplit}
           mode={dialogMode}
           initialData={dialogInitial}
-          isVlm={config.isVisionModel && config.isDatasetMultimodal === true}
+          isVlm={config.isVisionModel && config.isDatasetImage === true}
         />
 
         {canGoBack && (

studio/frontend/src/features/training/api/mappers.ts

@@ -68,7 +68,8 @@ export function buildTrainingStartPayload(
     finetune_language_layers: config.finetuneLanguageLayers,
     finetune_attention_modules: config.finetuneAttentionModules,
     finetune_mlp_modules: config.finetuneMLPModules,
-    is_dataset_multimodal: !!config.isDatasetMultimodal,
+    is_dataset_image: !!config.isDatasetImage,
+    is_dataset_audio: config.isDatasetAudio,
     enable_wandb: config.enableWandb,
     wandb_token: config.enableWandb ? config.wandbToken.trim() || null : null,
     wandb_project: config.enableWandb

studio/frontend/src/features/training/api/models-api.ts

@@ -46,6 +46,7 @@ interface BackendLoggingDefaults {
 }
 
 export interface BackendModelConfig {
+  audio_type?: string | null;
   training?: BackendTrainingDefaults;
   lora?: BackendLoraDefaults;
   logging?: BackendLoggingDefaults;
@@ -93,9 +94,11 @@ export async function checkVisionModel(modelName: string): Promise<boolean> {
 export async function getModelConfig(
   modelName: string,
   signal?: AbortSignal,
+  hfToken?: string,
 ): Promise<ModelConfigResponse> {
   const encoded = encodeURIComponent(modelName);
-  const response = await authFetch(`/api/models/config/${encoded}`, { signal });
+  const params = hfToken ? `?hf_token=${encodeURIComponent(hfToken)}` : "";
+  const response = await authFetch(`/api/models/config/${encoded}${params}`, { signal });
   if (!response.ok) {
     throw new Error(`Failed to fetch model config (${response.status})`);
   }

studio/frontend/src/features/training/hooks/use-training-actions.ts

@@ -49,7 +49,7 @@ export function useTrainingActions() {
 
     try {
       const datasetName = getDatasetName(config);
-      let isVlm = config.isVisionModel && config.isDatasetMultimodal === true;
+      let isVlm = config.isVisionModel && config.isDatasetImage === true;
 
       if (datasetName) {
         const check = await checkDatasetFormat({
@@ -60,12 +60,22 @@ export function useTrainingActions() {
           isVlm,
         });
 
-        // Backend auto-detects multimodal even if we didn't know yet
-        if (check.is_multimodal && config.isVisionModel) {
+        // Backend auto-detects image/audio from dataset content.
+        // Sync these flags into the store so buildTrainingStartPayload picks them up.
+        const isAudio = !!check.is_audio;
+        const isImage = !!check.is_image;
+
+        if (isImage && config.isVisionModel) {
           isVlm = true;
         }
+        if (isImage !== config.isDatasetImage || isAudio !== config.isDatasetAudio) {
+          useTrainingConfigStore.setState({
+            isDatasetImage: isImage,
+            isDatasetAudio: isAudio,
+          });
+        }
 
-        if (check.requires_manual_mapping && !hasManualMapping(config, isVlm)) {
+        if (check.requires_manual_mapping && !hasManualMapping(config, isVlm, isAudio)) {
           // Pre-fill from suggested_mapping or VLM detected columns
           const hint: Record<string, string> = {};
           if (check.suggested_mapping) {
@@ -73,6 +83,10 @@ export function useTrainingActions() {
             for (const [col, role] of Object.entries(check.suggested_mapping)) {
               hint[col] = table ? (table[role] ?? role) : role;
             }
+          } else if (isAudio) {
+            if (check.detected_audio_column) hint[check.detected_audio_column] = "audio";
+            if (check.detected_text_column) hint[check.detected_text_column] = "text";
+            if (check.detected_speaker_column) hint[check.detected_speaker_column] = "speaker_id";
           } else if (isVlm) {
             if (check.detected_image_column) hint[check.detected_image_column] = "image";
             if (check.detected_text_column) hint[check.detected_text_column] = "text";
@@ -88,7 +102,8 @@ export function useTrainingActions() {
         }
       }
 
-      const payload = buildTrainingStartPayload(config);
+      // Re-read config after potential store updates from dataset check
+      const payload = buildTrainingStartPayload(useTrainingConfigStore.getState());
       const response = await startTraining(payload);
 
       if (response.status === "error") {
@@ -159,12 +174,11 @@ function getDatasetName(config: TrainingConfigState): string | null {
     : config.uploadedFile;
 }
 
-function hasManualMapping(config: TrainingConfigState, isVlm = false): boolean {
+function hasManualMapping(config: TrainingConfigState, isVlm = false, isAudio = false): boolean {
   const mapping = config.datasetManualMapping;
   const roles = new Set(Object.values(mapping));
-  if (isVlm) {
-    return roles.has("image") && roles.has("text");
-  }
+  if (isAudio) return roles.has("audio") && roles.has("text");
+  if (isVlm) return roles.has("image") && roles.has("text");
   const fmt = config.datasetFormat;
   if (fmt === "alpaca") return roles.has("instruction") && roles.has("output");
   if (fmt === "sharegpt") return roles.has("human") && roles.has("gpt");

studio/frontend/src/features/training/stores/training-config-store.ts

@@ -37,7 +37,8 @@ const initialState: TrainingConfigState = {
   modelDefaultsError: null,
   modelDefaultsAppliedFor: null,
   isCheckingDataset: false,
-  isDatasetMultimodal: null,
+  isDatasetImage: null,
+  isDatasetAudio: false,
   ...DEFAULT_HYPERPARAMS,
 };
 
@@ -58,7 +59,8 @@ const NON_PERSISTED_STATE_KEYS: ReadonlySet<keyof TrainingConfigState> = new Set
   "modelDefaultsError",
   "modelDefaultsAppliedFor",
   "isCheckingDataset",
-  "isDatasetMultimodal",
+  "isDatasetImage",
+  "isDatasetAudio",
   "trainOnCompletions",
 ]);
 
@@ -108,7 +110,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
           modelDefaultsError: null,
         });
 
-        void getModelConfig(modelName, controller.signal)
+        void getModelConfig(modelName, controller.signal, get().hfToken || undefined)
           .then((modelDetails) => {
             if (controller.signal.aborted) return;
             if (get().selectedModel !== modelName) return;
@@ -116,9 +118,9 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             _trainOnCompletionsManuallySet = false;
             const patch = mapBackendModelConfigToTrainingPatch(modelDetails.config);
 
-            // If vision model + multimodal dataset already known, override
+            // If vision model + image dataset already known, override
             // trainOnCompletions to false regardless of backend default.
-            if (modelDetails.is_vision && get().isDatasetMultimodal === true) {
+            if (modelDetails.is_vision && get().isDatasetImage === true) {
               patch.trainOnCompletions = false;
             }
 
@@ -174,14 +176,16 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
         })
           .then((res) => {
             if (controller.signal.aborted) return;
-            const isMultimodal = !!res.is_multimodal;
+            const isImage = !!res.is_image;
+            const isAudio = !!res.is_audio;
             const updates: Record<string, unknown> = {
-              isDatasetMultimodal: isMultimodal,
+              isDatasetImage: isImage,
+              isDatasetAudio: isAudio,
               isCheckingDataset: false,
             };
             if (!_trainOnCompletionsManuallySet) {
               const { isVisionModel } = get();
-              if (isVisionModel && isMultimodal) {
+              if (isVisionModel && isImage) {
                 updates.trainOnCompletions = false;
               }
             }
@@ -189,7 +193,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
           })
           .catch(() => {
             if (controller.signal.aborted) return;
-            set({ isDatasetMultimodal: null, isCheckingDataset: false });
+            set({ isDatasetImage: null, isCheckingDataset: false });
           });
       };
 
@@ -207,6 +211,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             selectedModel: null,
             isCheckingVision: false,
             isVisionModel: false,
+            isDatasetAudio: false,
             isLoadingModelDefaults: false,
             modelDefaultsError: null,
             modelDefaultsAppliedFor: null,
@@ -222,6 +227,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             set({
               isCheckingVision: false,
               isVisionModel: false,
+              isDatasetAudio: false,
               isLoadingModelDefaults: false,
               modelDefaultsError: null,
               modelDefaultsAppliedFor: null,
@@ -259,7 +265,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             datasetManualMapping: emptyManualMapping(),
             datasetSliceStart: null,
             datasetSliceEnd: null,
-            isDatasetMultimodal: null,
+            isDatasetImage: null,
             isCheckingDataset: false,
           });
         },
@@ -272,7 +278,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             datasetSplit: null,
             datasetEvalSplit: null,
             datasetManualMapping: emptyManualMapping(),
-            isDatasetMultimodal: null,
+            isDatasetImage: null,
             isCheckingDataset: false,
           });
         },
@@ -280,7 +286,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
           set({
             datasetSplit,
             datasetManualMapping: emptyManualMapping(),
-            isDatasetMultimodal: null,
+            isDatasetImage: null,
             isCheckingDataset: false,
           });
 
@@ -296,7 +302,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
         ensureDatasetChecked: () => {
           const state = get();
           if (state.isCheckingDataset) return;
-          if (state.isDatasetMultimodal !== null) return;
+          if (state.isDatasetImage !== null) return;
 
           const datasetName =
             state.datasetSource === "huggingface"

studio/frontend/src/features/training/types/api.ts

@@ -40,7 +40,8 @@ export interface TrainingStartRequest {
   finetune_language_layers: boolean;
   finetune_attention_modules: boolean;
   finetune_mlp_modules: boolean;
-  is_dataset_multimodal: boolean;
+  is_dataset_image: boolean;
+  is_dataset_audio: boolean;
   enable_wandb: boolean;
   wandb_token: string | null;
   wandb_project: string | null;

studio/frontend/src/features/training/types/config.ts

@@ -61,7 +61,8 @@ export interface TrainingConfigState {
   modelDefaultsError: string | null;
   modelDefaultsAppliedFor: string | null;
   isCheckingDataset: boolean;
-  isDatasetMultimodal: boolean | null;
+  isDatasetImage: boolean | null;
+  isDatasetAudio: boolean;
   finetuneVisionLayers: boolean;
   finetuneLanguageLayers: boolean;
   finetuneAttentionModules: boolean;

studio/frontend/src/features/training/types/datasets.ts

@@ -4,10 +4,13 @@ export type CheckFormatResponse = {
   columns: string[];
   suggested_mapping?: Record<string, string> | null;
   detected_image_column?: string | null;
+  detected_audio_column?: string | null;
   detected_text_column?: string | null;
+  detected_speaker_column?: string | null;
   preview_samples?: Record<string, unknown>[] | null;
   total_rows?: number | null;
-  is_multimodal?: boolean;
+  is_image?: boolean;
+  is_audio?: boolean;
   multimodal_columns?: string[] | null;
   warning?: string | null;
 };

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<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);
+  });
+}