revamping up the code and adding inference

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/core/inference/OuteTTS

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

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,92 @@ 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
+                    from huggingface_hub import snapshot_download
+
+                    # Spark-TTS: download full repo, then load from /LLM subfolder
+                    hf_repo = config.path
+                    local_dir = hf_repo.split("/")[-1]
+                    repo_path = snapshot_download(hf_repo, local_dir=local_dir)
+                    abs_repo_path = os.path.abspath(repo_path)
+                    llm_path = os.path.join(abs_repo_path, "LLM")
+                    logger.info(f"Spark-TTS: downloaded repo to {repo_path}, loading LLM from {llm_path}")
+
+                    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
+                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 this audio type
+                model_repo_path = self.models[model_name].get("model_repo_path")
+                self._audio_codec_manager.load_codec(audio_type, self.device, model_repo_path=model_repo_path)
+
+                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}")
@@ -186,6 +268,10 @@ class InferenceBackend:
         """
         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]
@@ -801,6 +887,122 @@ class InferenceBackend:
             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
+        user_text = "Transcribe this audio."
+        if messages:
+            for msg in reversed(messages):
+                if msg["role"] == "user" and msg.get("content"):
+                    user_text = msg["content"]
+                    break
+
+        # Build messages in Gemma 3n format — audio goes INTO apply_chat_template
+        audio_messages = []
+        if system_prompt:
+            audio_messages.append({"role": "system", "content": [{"type": "text", "text": system_prompt}]})
+        audio_messages.append({
+            "role": "user",
+            "content": [
+                {"type": "audio", "audio": audio_array},
+                {"type": "text", "text": user_text},
+            ],
+        })
+
+        # 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",
+        ).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,
+            )
+
+            generation_kwargs = dict(
+                **inputs,
+                streamer=streamer,
+                max_new_tokens=max_new_tokens,
+                use_cache=True,
+                do_sample=temperature > 0,
+                temperature=temperature,
+                top_p=top_p,
+                top_k=top_k,
+                min_p=min_p,
+            )
+
+            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_stream(self,
                        prompt: str,
                        temperature: float = 0.7,
@@ -925,6 +1127,165 @@ class InferenceBackend:
     # ... 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:
             logger.error("No active model available")

studio/backend/core/training/inference/OuteTTS

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

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/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__)
 
@@ -184,6 +189,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,
         )
 
@@ -330,14 +338,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()),
         )
@@ -350,11 +367,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]"]:
@@ -444,6 +568,79 @@ async def openai_chat_completions(
             )
         model_name = backend.active_model_name or payload.model
 
+        # ── Audio TTS path: auto-route to audio generation ────
+        model_info = backend.models.get(backend.active_model_name, {})
+        if model_info.get("is_audio"):
+            return await generate_audio(payload, request)
+
+        # ── 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():
+                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

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

studio/backend/utils/models/model_config.py

@@ -220,6 +220,11 @@ MODEL_NAME_MAPPING = {
     ],
     "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",
@@ -402,6 +407,10 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
             logger.info(f"Model {model_name} detected as VLM: has img_processor")
             return True
 
+        # Check 4: Exclude audio models that have ForConditionalGeneration but aren't VLMs
+        # (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration)
+        # These are handled by is_audio_model() instead
+
         # Check 4: Has image_token_index (common in VLMs for image placeholder tokens)
         if hasattr(config, 'image_token_index'):
             logger.info(f"Model {model_name} detected as VLM: has image_token_index")
@@ -425,6 +434,38 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
 pass
 
 
+def is_audio_model(model_name: str) -> Optional[str]:
+    """
+    Check if a model is a TTS audio model by looking up its YAML config.
+
+    Returns the audio_type string ('snac', 'csm', 'bicodec', 'dac') or None.
+    """
+    try:
+        defaults = load_model_defaults(model_name)
+        audio_type = defaults.get('audio_type')
+        if audio_type and isinstance(audio_type, str) and audio_type in ('snac', 'csm', 'bicodec', 'dac'):
+            logger.info(f"Model {model_name} detected as audio model: audio_type={audio_type}")
+            return audio_type
+        return None
+    except Exception as e:
+        logger.debug(f"Could not determine if {model_name} is audio model: {e}")
+        return None
+
+
+def has_audio_input_model(model_name: str) -> bool:
+    """
+    Check if a model accepts audio input (ASR/speech understanding) by looking up its YAML config.
+
+    Returns True if the model has 'audio_input: true' in its defaults.
+    """
+    try:
+        defaults = load_model_defaults(model_name)
+        return bool(defaults.get('audio_input'))
+    except Exception as e:
+        logger.debug(f"Could not determine if {model_name} has audio input: {e}")
+        return False
+
+
 def detect_gguf_model(path: str) -> Optional[str]:
     """
     Check if the given local path is or contains a GGUF model file.
@@ -909,6 +950,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_hf_repo: Optional[str] = None  # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF")
     gguf_variant: Optional[str] = None  # Quantization variant (e.g. "Q4_K_M")
@@ -944,6 +988,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 = is_audio_model(base_model)
+
             display_name = lora_path_obj.name
             identifier = lora_path  # Use path as identifier for local LoRAs
 
@@ -955,6 +1002,8 @@ class ModelConfig:
                 is_cached=True,  # Local LoRAs are always "cached"
                 is_vision=is_vision,
                 is_lora=True,
+                is_audio=audio_type is not None,
+                audio_type=audio_type,
                 base_model=base_model,
             )
 
@@ -1105,11 +1154,15 @@ class ModelConfig:
                 logger.warning(f"Could not determine base model for LoRA '{path}'")
                 return None
             vision = is_vision_model(base_model, hf_token=hf_token)
+            audio_type_val = is_audio_model(base_model)
+            has_audio_in = has_audio_input_model(base_model)
         else:
             vision = is_vision_model(identifier, hf_token=hf_token)
-        
+            audio_type_val = is_audio_model(identifier)
+            has_audio_in = has_audio_input_model(identifier)
+
         display_name = Path(path).name if is_local else identifier.split("/")[-1]
-        
+
         return cls(
             identifier=identifier,
             display_name=display_name,
@@ -1118,6 +1171,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,
+            audio_type=audio_type_val,
+            has_audio_input=has_audio_in,
             base_model=base_model,
         )
 

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

@@ -33,13 +33,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 +165,50 @@ const ComposerAnimated: FC = () => {
   );
 };
 
+const AUDIO_ACCEPT = "audio/wav,audio/mpeg,audio/webm,audio/ogg,audio/flac,audio/mp4";
+const MAX_AUDIO_SIZE = 50 * 1024 * 1024;
+
+function fileToBase64(file: File): Promise<string> {
+  return new Promise((resolve, reject) => {
+    const reader = new FileReader();
+    reader.onload = () => {
+      const result = reader.result as string;
+      const commaIndex = result.indexOf(",");
+      resolve(commaIndex >= 0 ? result.slice(commaIndex + 1) : result);
+    };
+    reader.onerror = () => reject(new Error("Failed to read file"));
+    reader.readAsDataURL(file);
+  });
+}
+
+const PendingAudioChip: FC = () => {
+  const audioName = useChatRuntimeStore((s) => s.pendingAudioName);
+  const clearPendingAudio = useChatRuntimeStore((s) => s.clearPendingAudio);
+  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 +222,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}>

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 {
@@ -92,6 +92,25 @@ 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 raw = (part as { type: "audio"; audio: string }).audio;
+        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 +154,64 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
         });
       }
       const imageBase64 = findLatestUserImageBase64(messages);
+      const audioBase64 = findLatestUserAudioBase64(messages);
+      // Clear pending audio from store after extracting (consumed on send)
+      if (audioBase64) {
+        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) {
+        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 +269,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,8 @@
 import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
 import { Button } from "@/components/ui/button";
 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 +18,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;
@@ -27,6 +29,8 @@ export interface CompareHandle {
 
 const IMAGE_ACCEPT = "image/jpeg,image/png,image/webp,image/gif";
 const MAX_IMAGE_SIZE = 20 * 1024 * 1024;
+const AUDIO_ACCEPT = "audio/wav,audio/mpeg,audio/webm,audio/ogg,audio/flac,audio/mp4";
+const MAX_AUDIO_SIZE = 50 * 1024 * 1024;
 
 function fileToBase64DataURL(file: File): Promise<string> {
   return new Promise((resolve, reject) => {
@@ -182,9 +186,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 +217,27 @@ 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) {
+        const reader = new FileReader();
+        reader.onload = () => {
+          const result = reader.result as string;
+          const commaIndex = result.indexOf(",");
+          const base64 = commaIndex >= 0 ? result.slice(commaIndex + 1) : result;
+          setPendingAudio({ name: file.name, base64 });
+          setPendingAudioStore(base64, file.name);
+        };
+        reader.readAsDataURL(file);
+        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 +245,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 +256,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 +269,8 @@ export function SharedComposer({
     }
     setText("");
     setPendingImages([]);
+    setPendingAudio(null);
+    clearPendingAudioStore();
     textareaRef.current?.focus();
   }
 
@@ -257,7 +290,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 +306,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 +315,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 +364,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 {