merge nightly

.gitignore

@@ -11,6 +11,8 @@ __pycache__/
 
 # Virtual environments
 .venv/
+.venv_overlay/
+.venv_t5/
 venv/
 env/
 environment.yaml

install_python_stack.py

@@ -174,6 +174,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",
@@ -252,8 +259,11 @@ def install_python_stack() -> int:
         [sys.executable, str(SINGLE_ENV / "patch_metadata.py")],
     )
 
-    # 13. Final check
-    run("Running pip check", [sys.executable, "-m", "pip", "check"], quiet=False)
+    # 13. 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

setup.ps1

@@ -685,12 +685,43 @@ $CuTag = Get-PytorchCudaTag
 Write-Host "   Installing PyTorch with CUDA support ($CuTag)..." -ForegroundColor Cyan
 pip install torch torchvision torchaudio --index-url "https://download.pytorch.org/whl/$CuTag" 2>&1 | Out-Null
 
+# Install Triton for Windows (enables torch.compile — without it training can hang)
+Write-Host "   Installing Triton for Windows..." -ForegroundColor Cyan
+pip install "triton-windows<3.7" 2>&1 | Out-Null
+Write-Host "[OK] Triton for Windows installed (enables torch.compile)" -ForegroundColor Green
+
 # Ordered heavy dependency installation — shared cross-platform script
 Write-Host "   Running ordered dependency installation..." -ForegroundColor Cyan
 python "$PSScriptRoot\install_python_stack.py"
 # Restore ErrorActionPreference after pip/python work
 $ErrorActionPreference = $prevEAP
 
+# ── Pre-install transformers 5.x into .venv_t5/ ──
+# Models like GLM-4.7-Flash need transformers>=5.2.0. Instead of pip-installing
+# at runtime (slow, ~10-15s), we pre-install into a separate directory.
+# The training subprocess just prepends .venv_t5/ to sys.path — instant switch.
+Write-Host ""
+Write-Host "   Pre-installing transformers 5.x for newer model support..." -ForegroundColor Cyan
+$VenvT5Dir = Join-Path $PSScriptRoot ".venv_t5"
+if (Test-Path $VenvT5Dir) { Remove-Item -Recurse -Force $VenvT5Dir }
+New-Item -ItemType Directory -Path $VenvT5Dir -Force | Out-Null
+$prevEAP_t5 = $ErrorActionPreference
+$ErrorActionPreference = "Continue"
+pip install --target $VenvT5Dir --no-deps "transformers==5.2.0" 2>&1 | Out-Null
+if ($LASTEXITCODE -ne 0) {
+    Write-Host "[FAIL] Could not install transformers 5.2.0 into .venv_t5/" -ForegroundColor Red
+    $ErrorActionPreference = $prevEAP_t5
+    exit 1
+}
+pip install --target $VenvT5Dir --no-deps "huggingface_hub==1.3.0" 2>&1 | Out-Null
+if ($LASTEXITCODE -ne 0) {
+    Write-Host "[FAIL] Could not install huggingface_hub 1.3.0 into .venv_t5/" -ForegroundColor Red
+    $ErrorActionPreference = $prevEAP_t5
+    exit 1
+}
+$ErrorActionPreference = $prevEAP_t5
+Write-Host "[OK] Transformers 5.x pre-installed to .venv_t5/" -ForegroundColor Green
+
 # ==========================================================================
 #  PHASE 3.5: Install OpenSSL dev (for HTTPS support in llama-server)
 # ==========================================================================

setup.sh

@@ -181,10 +181,24 @@ if [ "$IS_COLAB" = true ]; then
 else
     # Local: create venv (always start fresh to preserve correct install order)
     rm -rf .venv
+    rm -rf .venv_overlay  # Remove legacy overlay (no longer used)
+    rm -rf .venv_t5       # Will be rebuilt below
     "$BEST_PY" -m venv .venv
     source .venv/bin/activate
     install_python_stack
-    
+
+    # ── 6b. Pre-install transformers 5.x into .venv_t5/ ──
+    # Models like GLM-4.7-Flash need transformers>=5.2.0. Instead of pip-installing
+    # at runtime (slow, ~10-15s), we pre-install into a separate directory.
+    # The training subprocess just prepends .venv_t5/ to sys.path — instant switch.
+    echo ""
+    echo "   Pre-installing transformers 5.x for newer model support..."
+    VENV_T5_DIR="$SCRIPT_DIR/.venv_t5"
+    mkdir -p "$VENV_T5_DIR"
+    run_quiet "pip install transformers 5.x" pip install --target "$VENV_T5_DIR" --no-deps "transformers==5.2.0"
+    run_quiet "pip install huggingface_hub for t5" pip install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==1.3.0"
+    echo "✅ Transformers 5.x pre-installed to .venv_t5/"
+
     # ── 7. WSL: pre-install GGUF build dependencies ──
     # On WSL, sudo requires a password and can't be entered during GGUF export
     # (runs in a non-interactive subprocess). Install build deps here instead.

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/__init__.py

@@ -1,30 +1,18 @@
 """
 Unified core module for Unsloth backend
+
+Imports are LAZY (via __getattr__) so that training subprocesses can
+import core.training.worker without pulling in heavy ML dependencies
+like unsloth, transformers, or torch before the version activation
+code has a chance to run.
 """
 
-# Inference
-from .inference import InferenceBackend, get_inference_backend
-
-# Training
-from .training import UnslothTrainer, get_trainer, TrainingBackend, get_training_backend, TrainingProgress
-
-# Configuration (from utils)
-from utils.models import is_vision_model, ModelConfig, scan_trained_loras, load_model_defaults, get_base_model_from_lora
-
-# Utilities (from utils)
-from utils.paths import normalize_path, is_local_path, is_model_cached
-from utils.utils import without_hf_auth, format_error_message
-from utils.hardware import get_device, is_apple_silicon, clear_gpu_cache, get_gpu_memory_info, log_gpu_memory, DeviceType
-from utils.datasets import format_and_template_dataset
-
 __all__ = [
     # Inference
     'InferenceBackend',
     'get_inference_backend',
 
     # Training
-    'UnslothTrainer',
-    'get_trainer',
     'get_training_backend',
     'TrainingBackend',
     'TrainingProgress',
@@ -50,3 +38,72 @@ __all__ = [
     'clear_gpu_cache',
     'DeviceType',
 ]
+
+
+def __getattr__(name):
+    # Inference
+    if name in ('InferenceBackend', 'get_inference_backend'):
+        from .inference import InferenceBackend, get_inference_backend
+        globals()['InferenceBackend'] = InferenceBackend
+        globals()['get_inference_backend'] = get_inference_backend
+        return globals()[name]
+
+    # Training
+    if name in ('TrainingBackend', 'get_training_backend', 'TrainingProgress'):
+        from .training import TrainingBackend, get_training_backend, TrainingProgress
+        globals()['TrainingBackend'] = TrainingBackend
+        globals()['get_training_backend'] = get_training_backend
+        globals()['TrainingProgress'] = TrainingProgress
+        return globals()[name]
+
+    # Config (from utils.models)
+    if name in ('is_vision_model', 'ModelConfig', 'scan_trained_loras',
+                'load_model_defaults', 'get_base_model_from_lora'):
+        from utils.models import (
+            is_vision_model, ModelConfig, scan_trained_loras,
+            load_model_defaults, get_base_model_from_lora,
+        )
+        globals()['is_vision_model'] = is_vision_model
+        globals()['ModelConfig'] = ModelConfig
+        globals()['scan_trained_loras'] = scan_trained_loras
+        globals()['load_model_defaults'] = load_model_defaults
+        globals()['get_base_model_from_lora'] = get_base_model_from_lora
+        return globals()[name]
+
+    # Paths
+    if name in ('normalize_path', 'is_local_path', 'is_model_cached'):
+        from utils.paths import normalize_path, is_local_path, is_model_cached
+        globals()['normalize_path'] = normalize_path
+        globals()['is_local_path'] = is_local_path
+        globals()['is_model_cached'] = is_model_cached
+        return globals()[name]
+
+    # Utils
+    if name in ('without_hf_auth', 'format_error_message'):
+        from utils.utils import without_hf_auth, format_error_message
+        globals()['without_hf_auth'] = without_hf_auth
+        globals()['format_error_message'] = format_error_message
+        return globals()[name]
+
+    # Hardware
+    if name in ('get_device', 'is_apple_silicon', 'clear_gpu_cache',
+                'get_gpu_memory_info', 'log_gpu_memory', 'DeviceType'):
+        from utils.hardware import (
+            get_device, is_apple_silicon, clear_gpu_cache,
+            get_gpu_memory_info, log_gpu_memory, DeviceType,
+        )
+        globals()['get_device'] = get_device
+        globals()['is_apple_silicon'] = is_apple_silicon
+        globals()['clear_gpu_cache'] = clear_gpu_cache
+        globals()['get_gpu_memory_info'] = get_gpu_memory_info
+        globals()['log_gpu_memory'] = log_gpu_memory
+        globals()['DeviceType'] = DeviceType
+        return globals()[name]
+
+    # Datasets
+    if name == 'format_and_template_dataset':
+        from utils.datasets import format_and_template_dataset
+        globals()['format_and_template_dataset'] = format_and_template_dataset
+        return format_and_template_dataset
+
+    raise AttributeError(f"module 'core' has no attribute {name!r}")

studio/backend/core/export/__init__.py

@@ -1,9 +1,17 @@
 """
 Export submodule - Model export operations
+
+The default get_export_backend() returns an ExportOrchestrator that
+delegates to a subprocess. The original ExportBackend runs inside
+the subprocess and can be imported directly from .export when needed.
 """
-from .export import ExportBackend, get_export_backend
+from .orchestrator import ExportOrchestrator, get_export_backend
+
+# Expose ExportOrchestrator as ExportBackend for backward compat
+ExportBackend = ExportOrchestrator
 
 __all__ = [
     'ExportBackend',
+    'ExportOrchestrator',
     'get_export_backend',
 ]

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
                 )
@@ -453,7 +518,14 @@ class ExportBackend:
                 # Write export metadata so the Chat page can identify the base model
                 self._write_export_metadata(abs_save_dir)
 
-                logger.info(f"GGUF model saved successfully in {abs_save_dir}")
+                # Log final file locations (after relocation) so it's clear
+                # where the GGUF files actually ended up.
+                final_ggufs = sorted(glob.glob(os.path.join(abs_save_dir, "*.gguf")))
+                logger.info(
+                    "GGUF export complete. Final files in %s:\n  %s",
+                    abs_save_dir,
+                    "\n  ".join(os.path.basename(f) for f in final_ggufs) or "(none)",
+                )
 
             # Push to hub if requested
             if push_to_hub:

studio/backend/core/export/orchestrator.py

@@ -0,0 +1,392 @@
+"""
+Export orchestrator — subprocess-based.
+
+Provides the same API as ExportBackend, but delegates all ML work
+to a persistent subprocess. The subprocess is spawned on first checkpoint
+load and stays alive for subsequent export operations.
+
+When switching between checkpoints that need different transformers versions,
+the old subprocess is killed and a new one is spawned with the correct version.
+
+Pattern follows core/inference/orchestrator.py.
+"""
+import atexit
+import logging
+import multiprocessing as mp
+import queue
+import threading
+import time
+from pathlib import Path
+from typing import Any, List, Optional, Tuple
+
+logger = logging.getLogger(__name__)
+
+_CTX = mp.get_context("spawn")
+
+
+class ExportOrchestrator:
+    """
+    Export backend orchestrator — subprocess-based.
+
+    Exposes the same API surface as ExportBackend so routes/export.py
+    needs minimal changes. Internally, all heavy ML operations happen in
+    a persistent subprocess.
+    """
+
+    def __init__(self):
+        # Subprocess state
+        self._proc: Optional[mp.Process] = None
+        self._cmd_queue: Any = None
+        self._resp_queue: Any = None
+        self._lock = threading.Lock()
+
+        # Local state mirrors (updated from subprocess responses)
+        self.current_checkpoint: Optional[str] = None
+        self.is_vision: bool = False
+        self.is_peft: bool = False
+
+        atexit.register(self._cleanup)
+        logger.info("ExportOrchestrator initialized (subprocess mode)")
+
+    # ------------------------------------------------------------------
+    # Subprocess lifecycle
+    # ------------------------------------------------------------------
+
+    def _spawn_subprocess(self, config: dict) -> None:
+        """Spawn a new export subprocess."""
+        from .worker import run_export_process
+
+        self._cmd_queue = _CTX.Queue()
+        self._resp_queue = _CTX.Queue()
+
+        self._proc = _CTX.Process(
+            target=run_export_process,
+            kwargs={
+                "cmd_queue": self._cmd_queue,
+                "resp_queue": self._resp_queue,
+                "config": config,
+            },
+            daemon=True,
+        )
+        self._proc.start()
+        logger.info("Export subprocess started (pid=%s)", self._proc.pid)
+
+    def _shutdown_subprocess(self, timeout: float = 10.0) -> None:
+        """Gracefully shut down the export subprocess."""
+        if self._proc is None or not self._proc.is_alive():
+            self._proc = None
+            return
+
+        # 1. Drain stale responses
+        self._drain_queue()
+
+        # 2. Send shutdown command
+        try:
+            self._cmd_queue.put({"type": "shutdown"})
+        except (OSError, ValueError):
+            pass
+
+        # 3. Wait for graceful shutdown
+        try:
+            self._proc.join(timeout=timeout)
+        except Exception:
+            pass
+
+        # 4. Force kill if still alive
+        if self._proc is not None and self._proc.is_alive():
+            logger.warning("Export subprocess did not exit gracefully, terminating")
+            try:
+                self._proc.terminate()
+                self._proc.join(timeout=5)
+            except Exception:
+                pass
+            if self._proc is not None and self._proc.is_alive():
+                logger.warning("Subprocess still alive after terminate, killing")
+                try:
+                    self._proc.kill()
+                    self._proc.join(timeout=3)
+                except Exception:
+                    pass
+
+        self._proc = None
+        self._cmd_queue = None
+        self._resp_queue = None
+        logger.info("Export subprocess shut down")
+
+    def _cleanup(self):
+        """atexit handler."""
+        self._shutdown_subprocess(timeout=5.0)
+
+    def _ensure_subprocess_alive(self) -> bool:
+        """Check if subprocess is alive."""
+        return self._proc is not None and self._proc.is_alive()
+
+    # ------------------------------------------------------------------
+    # Queue helpers
+    # ------------------------------------------------------------------
+
+    def _send_cmd(self, cmd: dict) -> None:
+        """Send a command to the subprocess."""
+        if self._cmd_queue is None:
+            raise RuntimeError("No export subprocess running")
+        try:
+            self._cmd_queue.put(cmd)
+        except (OSError, ValueError) as exc:
+            raise RuntimeError(f"Failed to send command to subprocess: {exc}")
+
+    def _read_resp(self, timeout: float = 1.0) -> Optional[dict]:
+        """Read a response from the subprocess (non-blocking with timeout)."""
+        if self._resp_queue is None:
+            return None
+        try:
+            return self._resp_queue.get(timeout=timeout)
+        except queue.Empty:
+            return None
+        except (EOFError, OSError, ValueError):
+            return None
+
+    def _wait_response(
+        self, expected_type: str, timeout: float = 3600.0
+    ) -> dict:
+        """Block until a response of the expected type arrives.
+
+        Export operations can take a very long time — GGUF conversion for
+        large models (30B+) easily takes 20-30 minutes. Default timeout
+        is 1 hour.
+        """
+        deadline = time.monotonic() + timeout
+
+        while time.monotonic() < deadline:
+            remaining = max(0.1, deadline - time.monotonic())
+            resp = self._read_resp(timeout=min(remaining, 2.0))
+
+            if resp is None:
+                # Check subprocess health
+                if not self._ensure_subprocess_alive():
+                    raise RuntimeError("Export subprocess crashed during wait")
+                continue
+
+            rtype = resp.get("type", "")
+
+            if rtype == expected_type:
+                return resp
+
+            if rtype == "error":
+                error_msg = resp.get("error", "Unknown error")
+                raise RuntimeError(f"Subprocess error: {error_msg}")
+
+            if rtype == "status":
+                logger.info("Export subprocess status: %s", resp.get("message", ""))
+                continue
+
+            # Other response types during wait — skip
+            logger.debug(
+                "Skipping response type '%s' while waiting for '%s'",
+                rtype, expected_type,
+            )
+
+        raise RuntimeError(
+            f"Timeout waiting for '{expected_type}' response after {timeout}s"
+        )
+
+    def _drain_queue(self) -> list:
+        """Drain all pending responses."""
+        events = []
+        if self._resp_queue is None:
+            return events
+        while True:
+            try:
+                events.append(self._resp_queue.get_nowait())
+            except queue.Empty:
+                return events
+            except (EOFError, OSError, ValueError):
+                return events
+
+    # ------------------------------------------------------------------
+    # Public API — same interface as ExportBackend
+    # ------------------------------------------------------------------
+
+    def load_checkpoint(
+        self,
+        checkpoint_path: str,
+        max_seq_length: int = 2048,
+        load_in_4bit: bool = True,
+    ) -> Tuple[bool, str]:
+        """Load a checkpoint for export.
+
+        Always spawns a fresh subprocess to ensure a clean Python interpreter.
+        """
+        project_root = str(
+            Path(__file__).resolve().parent.parent.parent.parent.parent
+        )
+
+        sub_config = {
+            "project_root": project_root,
+            "checkpoint_path": checkpoint_path,
+            "max_seq_length": max_seq_length,
+            "load_in_4bit": load_in_4bit,
+        }
+
+        # Always kill existing subprocess and spawn fresh.
+        if self._ensure_subprocess_alive():
+            self._shutdown_subprocess()
+        elif self._proc is not None:
+            self._shutdown_subprocess(timeout=2)
+
+        logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path)
+        self._spawn_subprocess(sub_config)
+
+        try:
+            resp = self._wait_response("loaded", timeout=300)
+        except RuntimeError as exc:
+            self._shutdown_subprocess(timeout=5)
+            self.current_checkpoint = None
+            self.is_vision = False
+            self.is_peft = False
+            return False, str(exc)
+
+        if resp.get("success"):
+            self.current_checkpoint = resp.get("checkpoint")
+            self.is_vision = resp.get("is_vision", False)
+            self.is_peft = resp.get("is_peft", False)
+            logger.info("Checkpoint '%s' loaded in subprocess", checkpoint_path)
+            return True, resp.get("message", "Loaded successfully")
+        else:
+            error = resp.get("message", "Failed to load checkpoint")
+            logger.error("Failed to load checkpoint: %s", error)
+            self.current_checkpoint = None
+            self.is_vision = False
+            self.is_peft = False
+            return False, error
+
+    def export_merged_model(
+        self,
+        save_directory: str,
+        format_type: str = "16-bit (FP16)",
+        push_to_hub: bool = False,
+        repo_id: Optional[str] = None,
+        hf_token: Optional[str] = None,
+        private: bool = False,
+    ) -> Tuple[bool, str]:
+        """Export merged PEFT model."""
+        return self._run_export("merged", {
+            "save_directory": save_directory,
+            "format_type": format_type,
+            "push_to_hub": push_to_hub,
+            "repo_id": repo_id,
+            "hf_token": hf_token,
+            "private": private,
+        })
+
+    def export_base_model(
+        self,
+        save_directory: str,
+        push_to_hub: bool = False,
+        repo_id: Optional[str] = None,
+        hf_token: Optional[str] = None,
+        private: bool = False,
+        base_model_id: Optional[str] = None,
+    ) -> Tuple[bool, str]:
+        """Export base model (non-PEFT)."""
+        return self._run_export("base", {
+            "save_directory": save_directory,
+            "push_to_hub": push_to_hub,
+            "repo_id": repo_id,
+            "hf_token": hf_token,
+            "private": private,
+            "base_model_id": base_model_id,
+        })
+
+    def export_gguf(
+        self,
+        save_directory: str,
+        quantization_method: str = "Q4_K_M",
+        push_to_hub: bool = False,
+        repo_id: Optional[str] = None,
+        hf_token: Optional[str] = None,
+    ) -> Tuple[bool, str]:
+        """Export model in GGUF format."""
+        return self._run_export("gguf", {
+            "save_directory": save_directory,
+            "quantization_method": quantization_method,
+            "push_to_hub": push_to_hub,
+            "repo_id": repo_id,
+            "hf_token": hf_token,
+        })
+
+    def export_lora_adapter(
+        self,
+        save_directory: str,
+        push_to_hub: bool = False,
+        repo_id: Optional[str] = None,
+        hf_token: Optional[str] = None,
+        private: bool = False,
+    ) -> Tuple[bool, str]:
+        """Export LoRA adapter only."""
+        return self._run_export("lora", {
+            "save_directory": save_directory,
+            "push_to_hub": push_to_hub,
+            "repo_id": repo_id,
+            "hf_token": hf_token,
+            "private": private,
+        })
+
+    def _run_export(self, export_type: str, params: dict) -> Tuple[bool, str]:
+        """Send an export command to the subprocess and wait for result."""
+        if not self._ensure_subprocess_alive():
+            return False, "No export subprocess running. Load a checkpoint first."
+
+        cmd = {"type": "export", "export_type": export_type, **params}
+
+        try:
+            self._send_cmd(cmd)
+            resp = self._wait_response(
+                f"export_{export_type}_done",
+                timeout=3600,  # GGUF for 30B+ models can take 30+ min
+            )
+            return resp.get("success", False), resp.get("message", "")
+        except RuntimeError as exc:
+            return False, str(exc)
+
+    def cleanup_memory(self) -> bool:
+        """Cleanup export-related models from memory."""
+        if not self._ensure_subprocess_alive():
+            # No subprocess — just clear local state
+            self.current_checkpoint = None
+            self.is_vision = False
+            self.is_peft = False
+            return True
+
+        try:
+            self._send_cmd({"type": "cleanup"})
+            resp = self._wait_response("cleanup_done", timeout=30)
+            success = resp.get("success", False)
+        except RuntimeError:
+            success = False
+
+        # Shut down subprocess after cleanup — no model loaded
+        self._shutdown_subprocess()
+
+        self.current_checkpoint = None
+        self.is_vision = False
+        self.is_peft = False
+        return success
+
+    def scan_checkpoints(
+        self, outputs_dir: str = "./outputs"
+    ) -> List[Tuple[str, list]]:
+        """Scan for checkpoints — no ML imports needed, runs locally."""
+        from utils.models.checkpoints import scan_checkpoints
+        return scan_checkpoints(outputs_dir=outputs_dir)
+
+
+# ========== GLOBAL INSTANCE ==========
+_export_backend = None
+
+
+def get_export_backend() -> ExportOrchestrator:
+    """Get global export backend instance (orchestrator)."""
+    global _export_backend
+    if _export_backend is None:
+        _export_backend = ExportOrchestrator()
+    return _export_backend

studio/backend/core/export/worker.py

@@ -0,0 +1,350 @@
+"""
+Export subprocess entry point.
+
+Each export session runs in a persistent subprocess (mp.get_context("spawn")).
+This gives us a clean Python interpreter with no stale module state —
+solving the transformers version-switching problem completely.
+
+The subprocess stays alive while a model is loaded, accepting commands
+(load, export_merged, export_base, export_gguf, export_lora, cleanup,
+shutdown) via mp.Queue.
+
+Pattern follows core/inference/worker.py and core/training/worker.py.
+"""
+from __future__ import annotations
+
+import logging
+import os
+import sys
+import time
+import traceback
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def _activate_transformers_version(model_name: str, project_root: str) -> None:
+    """Activate the correct transformers version BEFORE any ML imports.
+
+    If the model needs transformers 5.x, prepend the pre-installed .venv_t5/
+    directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/).
+    """
+    # Ensure backend is on path for utils imports
+    backend_path = os.path.join(project_root, "studio", "backend")
+    if backend_path not in sys.path:
+        sys.path.insert(0, backend_path)
+
+    from utils.transformers_version import needs_transformers_5, _resolve_base_model
+
+    resolved = _resolve_base_model(model_name)
+    if needs_transformers_5(resolved):
+        venv_t5 = os.path.join(project_root, ".venv_t5")
+        if os.path.isdir(venv_t5):
+            sys.path.insert(0, venv_t5)
+            logger.info("Activated transformers 5.x from %s", venv_t5)
+        else:
+
+            # Fallback: pip install at runtime (slower, ~10-15s)
+            logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5)
+            import subprocess as sp
+            os.makedirs(venv_t5, exist_ok=True)
+            r1 = sp.run(
+                [sys.executable, "-m", "pip", "install", "--target", venv_t5,
+                 "--no-deps", "transformers==5.2.0"],
+                stdout=sp.PIPE, stderr=sp.STDOUT,
+            )
+            r2 = sp.run(
+                [sys.executable, "-m", "pip", "install", "--target", venv_t5,
+                 "--no-deps", "huggingface_hub==1.3.0"],
+                stdout=sp.PIPE, stderr=sp.STDOUT,
+            )
+            if r1.returncode != 0 or r2.returncode != 0:
+                raise RuntimeError(
+                    f"Failed to install transformers 5.x into {venv_t5}. "
+                    f"pip returncode: transformers={r1.returncode}, huggingface_hub={r2.returncode}"
+                )
+            sys.path.insert(0, venv_t5)
+        # Propagate to child subprocesses (e.g. GGUF converter)
+        _pp = os.environ.get("PYTHONPATH", "")
+        os.environ["PYTHONPATH"] = venv_t5 + (os.pathsep + _pp if _pp else "")
+    else:
+        logger.info("Using default transformers (4.57.x) for %s", model_name)
+
+
+def _send_response(resp_queue: Any, response: dict) -> None:
+    """Send a response to the parent process."""
+    try:
+        resp_queue.put(response)
+    except (OSError, ValueError) as exc:
+        logger.error("Failed to send response: %s", exc)
+
+
+def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
+    """Handle a load_checkpoint command."""
+    checkpoint_path = cmd["checkpoint_path"]
+    max_seq_length = cmd.get("max_seq_length", 2048)
+    load_in_4bit = cmd.get("load_in_4bit", True)
+
+    try:
+        _send_response(resp_queue, {
+            "type": "status",
+            "message": f"Loading checkpoint: {checkpoint_path}",
+            "ts": time.time(),
+        })
+
+        success, message = backend.load_checkpoint(
+            checkpoint_path=checkpoint_path,
+            max_seq_length=max_seq_length,
+            load_in_4bit=load_in_4bit,
+        )
+
+        _send_response(resp_queue, {
+            "type": "loaded",
+            "success": success,
+            "message": message,
+            "checkpoint": checkpoint_path if success else None,
+            "is_vision": backend.is_vision if success else False,
+            "is_peft": backend.is_peft if success else False,
+            "ts": time.time(),
+        })
+
+    except Exception as exc:
+        _send_response(resp_queue, {
+            "type": "loaded",
+            "success": False,
+            "message": str(exc),
+            "stack": traceback.format_exc(limit=20),
+            "ts": time.time(),
+        })
+
+
+def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
+    """Handle any export command (merged, base, gguf, lora)."""
+    export_type = cmd["export_type"]  # "merged", "base", "gguf", "lora"
+    response_type = f"export_{export_type}_done"
+
+    try:
+        if export_type == "merged":
+            success, message = backend.export_merged_model(
+                save_directory=cmd.get("save_directory", ""),
+                format_type=cmd.get("format_type", "16-bit (FP16)"),
+                push_to_hub=cmd.get("push_to_hub", False),
+                repo_id=cmd.get("repo_id"),
+                hf_token=cmd.get("hf_token"),
+                private=cmd.get("private", False),
+            )
+        elif export_type == "base":
+            success, message = backend.export_base_model(
+                save_directory=cmd.get("save_directory", ""),
+                push_to_hub=cmd.get("push_to_hub", False),
+                repo_id=cmd.get("repo_id"),
+                hf_token=cmd.get("hf_token"),
+                private=cmd.get("private", False),
+                base_model_id=cmd.get("base_model_id"),
+            )
+        elif export_type == "gguf":
+            success, message = backend.export_gguf(
+                save_directory=cmd.get("save_directory", ""),
+                quantization_method=cmd.get("quantization_method", "Q4_K_M"),
+                push_to_hub=cmd.get("push_to_hub", False),
+                repo_id=cmd.get("repo_id"),
+                hf_token=cmd.get("hf_token"),
+            )
+        elif export_type == "lora":
+            success, message = backend.export_lora_adapter(
+                save_directory=cmd.get("save_directory", ""),
+                push_to_hub=cmd.get("push_to_hub", False),
+                repo_id=cmd.get("repo_id"),
+                hf_token=cmd.get("hf_token"),
+                private=cmd.get("private", False),
+            )
+        else:
+            success, message = False, f"Unknown export type: {export_type}"
+
+        _send_response(resp_queue, {
+            "type": response_type,
+            "success": success,
+            "message": message,
+            "ts": time.time(),
+        })
+
+    except Exception as exc:
+        _send_response(resp_queue, {
+            "type": response_type,
+            "success": False,
+            "message": str(exc),
+            "stack": traceback.format_exc(limit=20),
+            "ts": time.time(),
+        })
+
+
+def _handle_cleanup(backend, resp_queue: Any) -> None:
+    """Handle a cleanup command."""
+    try:
+        success = backend.cleanup_memory()
+        _send_response(resp_queue, {
+            "type": "cleanup_done",
+            "success": success,
+            "ts": time.time(),
+        })
+    except Exception as exc:
+        _send_response(resp_queue, {
+            "type": "cleanup_done",
+            "success": False,
+            "message": str(exc),
+            "ts": time.time(),
+        })
+
+
+def run_export_process(
+    *,
+    cmd_queue: Any,
+    resp_queue: Any,
+    config: dict,
+) -> None:
+    """Subprocess entrypoint. Persistent — runs command loop until shutdown.
+
+    Args:
+        cmd_queue: mp.Queue for receiving commands from parent.
+        resp_queue: mp.Queue for sending responses to parent.
+        config: Initial configuration dict with checkpoint_path and project_root.
+    """
+    import queue as _queue
+
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+    project_root = config["project_root"]
+    checkpoint_path = config["checkpoint_path"]
+
+    # ── 1. Activate correct transformers version BEFORE any ML imports ──
+    try:
+        _activate_transformers_version(checkpoint_path, project_root)
+    except Exception as exc:
+        _send_response(resp_queue, {
+            "type": "error",
+            "error": f"Failed to activate transformers version: {exc}",
+            "stack": traceback.format_exc(limit=20),
+            "ts": time.time(),
+        })
+        return
+
+    # ── 1b. On Windows, check Triton availability (must be before import torch) ──
+    if sys.platform == "win32":
+        try:
+            import triton  # noqa: F401
+            logger.info("Triton available — torch.compile enabled")
+        except ImportError:
+            os.environ["TORCHDYNAMO_DISABLE"] = "1"
+            logger.warning(
+                "Triton not found on Windows — torch.compile disabled. "
+                'Install for better performance: pip install "triton-windows<3.7"'
+            )
+
+    # ── 2. Import ML libraries (fresh in this clean process) ──
+    try:
+        _send_response(resp_queue, {
+            "type": "status",
+            "message": "Importing ML libraries...",
+            "ts": time.time(),
+        })
+
+        backend_path = os.path.join(project_root, "studio", "backend")
+        if backend_path not in sys.path:
+            sys.path.insert(0, backend_path)
+
+        from core.export.export import ExportBackend
+
+        import transformers
+        logger.info("Export subprocess loaded transformers %s", transformers.__version__)
+
+    except Exception as exc:
+        _send_response(resp_queue, {
+            "type": "error",
+            "error": f"Failed to import ML libraries: {exc}",
+            "stack": traceback.format_exc(limit=20),
+            "ts": time.time(),
+        })
+        return
+
+    # ── 3. Create export backend and load initial checkpoint ──
+    try:
+        backend = ExportBackend()
+
+        _handle_load(backend, config, resp_queue)
+
+    except Exception as exc:
+        _send_response(resp_queue, {
+            "type": "error",
+            "error": f"Failed to initialize export backend: {exc}",
+            "stack": traceback.format_exc(limit=20),
+            "ts": time.time(),
+        })
+        return
+
+    # ── 4. Command loop — process commands until shutdown ──
+    logger.info("Export subprocess ready, entering command loop")
+
+    while True:
+        try:
+            cmd = cmd_queue.get(timeout=1.0)
+        except _queue.Empty:
+            continue
+        except (EOFError, OSError):
+            logger.info("Command queue closed, shutting down")
+            return
+
+        if cmd is None:
+            continue
+
+        cmd_type = cmd.get("type", "")
+        logger.info("Received command: %s", cmd_type)
+
+        try:
+            if cmd_type == "load":
+                # Load a new checkpoint (reusing this subprocess)
+                backend.cleanup_memory()
+                _handle_load(backend, cmd, resp_queue)
+
+            elif cmd_type == "export":
+                _handle_export(backend, cmd, resp_queue)
+
+            elif cmd_type == "cleanup":
+                _handle_cleanup(backend, resp_queue)
+
+            elif cmd_type == "status":
+                _send_response(resp_queue, {
+                    "type": "status_response",
+                    "checkpoint": backend.current_checkpoint,
+                    "is_vision": backend.is_vision,
+                    "is_peft": backend.is_peft,
+                    "ts": time.time(),
+                })
+
+            elif cmd_type == "shutdown":
+                logger.info("Shutdown command received, cleaning up and exiting")
+                try:
+                    backend.cleanup_memory()
+                except Exception:
+                    pass
+                _send_response(resp_queue, {
+                    "type": "shutdown_ack",
+                    "ts": time.time(),
+                })
+                return
+
+            else:
+                logger.warning("Unknown command type: %s", cmd_type)
+                _send_response(resp_queue, {
+                    "type": "error",
+                    "error": f"Unknown command type: {cmd_type}",
+                    "ts": time.time(),
+                })
+
+        except Exception as exc:
+            logger.error("Error handling command '%s': %s", cmd_type, exc, exc_info=True)
+            _send_response(resp_queue, {
+                "type": "error",
+                "error": f"Command '{cmd_type}' failed: {exc}",
+                "stack": traceback.format_exc(limit=20),
+                "ts": time.time(),
+            })

studio/backend/core/inference/__init__.py

@@ -1,11 +1,19 @@
 """
 Inference submodule - Inference backend for model loading and generation
+
+The default get_inference_backend() returns an InferenceOrchestrator that
+delegates to a subprocess. The original InferenceBackend runs inside
+the subprocess and can be imported directly from .inference when needed.
 """
-from .inference import InferenceBackend, get_inference_backend
+from .orchestrator import InferenceOrchestrator, get_inference_backend
 from .llama_cpp import LlamaCppBackend
 
+# Expose InferenceOrchestrator as InferenceBackend for backward compat
+InferenceBackend = InferenceOrchestrator
+
 __all__ = [
     'InferenceBackend',
+    'InferenceOrchestrator',
     'get_inference_backend',
     'LlamaCppBackend',
 ]

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

@@ -0,0 +1,800 @@
+"""
+Inference orchestrator — subprocess-based.
+
+Provides the same API as InferenceBackend, but delegates all ML work
+to a persistent subprocess. The subprocess is spawned on first model load
+and stays alive for subsequent requests.
+
+When switching between models that need different transformers versions
+(e.g. GLM-4.7-Flash needs 5.x, Qwen needs 4.57.x), the old subprocess
+is killed and a new one is spawned with the correct version.
+
+Pattern follows core/training/training.py.
+"""
+import atexit
+import base64
+import logging
+import multiprocessing as mp
+import queue
+import threading
+import time
+import uuid
+from io import BytesIO
+from pathlib import Path
+from typing import Any, Generator, Optional, Tuple, Union
+
+logger = logging.getLogger(__name__)
+
+_CTX = mp.get_context("spawn")
+
+
+class InferenceOrchestrator:
+    """
+    Inference backend orchestrator — subprocess-based.
+
+    Exposes the same API surface as InferenceBackend so routes/inference.py
+    needs minimal changes. Internally, all heavy ML operations happen in
+    a persistent subprocess.
+    """
+
+    def __init__(self):
+        # Subprocess state
+        self._proc: Optional[mp.Process] = None
+        self._cmd_queue: Any = None
+        self._resp_queue: Any = None
+        self._cancel_event: Any = None  # mp.Event — set to cancel generation instantly
+        self._lock = threading.Lock()
+        self._gen_lock = threading.Lock()  # Serializes generation — one request at a time
+
+        # Local state mirrors (updated from subprocess responses)
+        self.active_model_name: Optional[str] = None
+        self.models: dict = {}
+        self.loading_models: set = set()
+        self.loaded_local_models: list = []
+        self.default_models = [
+            "unsloth/Qwen3-4B-Instruct-2507",
+            "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit",
+            "unsloth/Mistral-Nemo-Instruct-2407-bnb-4bit",
+            "unsloth/Phi-3.5-mini-instruct",
+            "unsloth/Gemma-3-4B-it",
+            "unsloth/Qwen2-VL-2B-Instruct-bnb-4bit",
+        ]
+
+        # Version tracking for subprocess reuse
+        self._current_transformers_major: Optional[str] = None  # "4" or "5"
+
+        atexit.register(self._cleanup)
+        logger.info("InferenceOrchestrator initialized (subprocess mode)")
+
+    # ------------------------------------------------------------------
+    # Subprocess lifecycle
+    # ------------------------------------------------------------------
+
+    def _spawn_subprocess(self, config: dict) -> None:
+        """Spawn a new inference subprocess."""
+        from .worker import run_inference_process
+
+        self._cmd_queue = _CTX.Queue()
+        self._resp_queue = _CTX.Queue()
+        self._cancel_event = _CTX.Event()
+
+        self._proc = _CTX.Process(
+            target=run_inference_process,
+            kwargs={
+                "cmd_queue": self._cmd_queue,
+                "resp_queue": self._resp_queue,
+                "cancel_event": self._cancel_event,
+                "config": config,
+            },
+            daemon=True,
+        )
+        self._proc.start()
+        logger.info("Inference subprocess started (pid=%s)", self._proc.pid)
+
+    def _cancel_generation(self) -> None:
+        """Cancel any ongoing generation in the subprocess (instant)."""
+        if self._cancel_event is not None:
+            self._cancel_event.set()
+
+    def _shutdown_subprocess(self, timeout: float = 10.0) -> None:
+        """Gracefully shut down the inference subprocess."""
+        if self._proc is None or not self._proc.is_alive():
+            self._proc = None
+            return
+
+        # 1. Cancel any ongoing generation first (instant via mp.Event)
+        self._cancel_generation()
+        time.sleep(0.5)  # Brief wait for generation to stop
+
+        # 2. Drain stale responses from queue
+        self._drain_queue()
+
+        # 3. Send shutdown command
+        try:
+            self._cmd_queue.put({"type": "shutdown"})
+        except (OSError, ValueError):
+            pass
+
+        # 4. Wait for graceful shutdown
+        try:
+            self._proc.join(timeout=timeout)
+        except Exception:
+            pass
+
+        # 5. Force kill if still alive
+        if self._proc is not None and self._proc.is_alive():
+            logger.warning("Inference subprocess did not exit gracefully, terminating")
+            try:
+                self._proc.terminate()
+                self._proc.join(timeout=5)
+            except Exception:
+                pass
+            if self._proc is not None and self._proc.is_alive():
+                logger.warning("Subprocess still alive after terminate, killing")
+                try:
+                    self._proc.kill()
+                    self._proc.join(timeout=3)
+                except Exception:
+                    pass
+
+        self._proc = None
+        self._cmd_queue = None
+        self._resp_queue = None
+        self._cancel_event = None
+        logger.info("Inference subprocess shut down")
+
+    def _cleanup(self):
+        """atexit handler."""
+        self._shutdown_subprocess(timeout=5.0)
+
+    def _ensure_subprocess_alive(self) -> bool:
+        """Check if subprocess is alive."""
+        return self._proc is not None and self._proc.is_alive()
+
+    # ------------------------------------------------------------------
+    # Queue helpers
+    # ------------------------------------------------------------------
+
+    def _send_cmd(self, cmd: dict) -> None:
+        """Send a command to the subprocess."""
+        if self._cmd_queue is None:
+            raise RuntimeError("No inference subprocess running")
+        try:
+            self._cmd_queue.put(cmd)
+        except (OSError, ValueError) as exc:
+            raise RuntimeError(f"Failed to send command to subprocess: {exc}")
+
+    def _read_resp(self, timeout: float = 1.0) -> Optional[dict]:
+        """Read a response from the subprocess (non-blocking with timeout)."""
+        if self._resp_queue is None:
+            return None
+        try:
+            return self._resp_queue.get(timeout=timeout)
+        except queue.Empty:
+            return None
+        except (EOFError, OSError, ValueError):
+            return None
+
+    def _wait_response(
+        self, expected_type: str, timeout: float = 120.0
+    ) -> dict:
+        """Block until a response of the expected type arrives.
+
+        Also handles 'status' and 'error' events during the wait.
+        Returns the matching response dict.
+        Raises RuntimeError on timeout or subprocess crash.
+        """
+        deadline = time.monotonic() + timeout
+
+        while time.monotonic() < deadline:
+            remaining = max(0.1, deadline - time.monotonic())
+            resp = self._read_resp(timeout=min(remaining, 1.0))
+
+            if resp is None:
+                # Check subprocess health
+                if not self._ensure_subprocess_alive():
+                    raise RuntimeError("Inference subprocess crashed during wait")
+                continue
+
+            rtype = resp.get("type", "")
+
+            if rtype == expected_type:
+                return resp
+
+            if rtype == "error":
+                error_msg = resp.get("error", "Unknown error")
+                raise RuntimeError(f"Subprocess error: {error_msg}")
+
+            if rtype == "status":
+                logger.info("Subprocess status: %s", resp.get("message", ""))
+                continue
+
+            # Other response types during wait — skip
+            logger.debug("Skipping response type '%s' while waiting for '%s'", rtype, expected_type)
+
+        raise RuntimeError(
+            f"Timeout waiting for '{expected_type}' response after {timeout}s"
+        )
+
+    def _drain_queue(self) -> list:
+        """Drain all pending responses."""
+        events = []
+        if self._resp_queue is None:
+            return events
+        while True:
+            try:
+                events.append(self._resp_queue.get_nowait())
+            except queue.Empty:
+                return events
+            except (EOFError, OSError, ValueError):
+                return events
+
+    def _drain_until_gen_done(self, timeout: float = 5.0) -> None:
+        """Consume resp_queue events until gen_done/gen_error, discarding them.
+
+        Called after cancel to ensure stale tokens from the cancelled
+        generation don't leak into the next request.
+        """
+        deadline = time.monotonic() + timeout
+        while time.monotonic() < deadline:
+            resp = self._read_resp(timeout=min(0.5, deadline - time.monotonic()))
+            if resp is None:
+                if not self._ensure_subprocess_alive():
+                    return
+                continue
+            rtype = resp.get("type", "")
+            if rtype in ("gen_done", "gen_error"):
+                return
+        logger.warning("Timed out waiting for gen_done after cancel")
+
+    # ------------------------------------------------------------------
+    # Public API — same interface as InferenceBackend
+    # ------------------------------------------------------------------
+
+    def load_model(
+        self,
+        config,  # ModelConfig
+        max_seq_length: int = 2048,
+        dtype=None,
+        load_in_4bit: bool = True,
+        hf_token: Optional[str] = None,
+    ) -> bool:
+        """Load a model for inference.
+
+        Always spawns a fresh subprocess for each model load. This ensures
+        a clean Python interpreter — no stale unsloth patches, torch.compile
+        caches, or inspect.getsource() failures from a previous model.
+        """
+        from utils.transformers_version import needs_transformers_5
+
+        model_name = config.identifier
+        self.loading_models.add(model_name)
+
+        try:
+            needed_major = "5" if needs_transformers_5(model_name) else "4"
+            project_root = str(Path(__file__).resolve().parent.parent.parent.parent.parent)
+
+            # Build config dict for subprocess
+            sub_config = {
+                "project_root": project_root,
+                "model_name": model_name,
+                "max_seq_length": max_seq_length,
+                "load_in_4bit": load_in_4bit,
+                "hf_token": hf_token or "",
+                "gguf_variant": getattr(config, "gguf_variant", None),
+            }
+
+            # Always kill existing subprocess and spawn fresh.
+            # Reusing a subprocess after unsloth patches torch internals
+            # causes inspect.getsource() failures on the next model load.
+            if self._ensure_subprocess_alive():
+                self._cancel_generation()
+                time.sleep(0.3)
+                self._shutdown_subprocess()
+
+            elif self._proc is not None:
+                # Dead subprocess — clean up
+                self._shutdown_subprocess(timeout=2)
+
+            logger.info(
+                "Spawning fresh inference subprocess for '%s' (transformers %s.x)",
+                model_name, needed_major,
+            )
+            self._spawn_subprocess(sub_config)
+            resp = self._wait_response("loaded", timeout=180)
+
+            # Update local state from response
+            if resp.get("success"):
+                self._current_transformers_major = needed_major
+                model_info = resp.get("model_info", {})
+                self.active_model_name = model_info.get("identifier", model_name)
+                self.models[self.active_model_name] = {
+                    "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)
+                return True
+            else:
+                error = resp.get("error", "Failed to load model")
+                self.loading_models.discard(model_name)
+                self.active_model_name = None
+                self.models.clear()
+                raise Exception(error)
+
+        except Exception:
+            self.loading_models.discard(model_name)
+            self.active_model_name = None
+            self.models.clear()
+            raise
+
+    def unload_model(self, model_name: str) -> bool:
+        """Unload a model from the subprocess."""
+        if not self._ensure_subprocess_alive():
+            # No subprocess — just clear local state
+            self.models.pop(model_name, None)
+            if self.active_model_name == model_name:
+                self.active_model_name = None
+            return True
+
+        try:
+            self._send_cmd({
+                "type": "unload",
+                "model_name": model_name,
+            })
+            resp = self._wait_response("unloaded", timeout=30)
+
+            # Update local state
+            self.models.pop(model_name, None)
+            if self.active_model_name == model_name:
+                self.active_model_name = None
+
+            logger.info("Model '%s' unloaded from subprocess", model_name)
+            return True
+
+        except Exception as exc:
+            logger.error("Error unloading model '%s': %s", model_name, exc)
+            # Clear local state anyway
+            self.models.pop(model_name, None)
+            if self.active_model_name == model_name:
+                self.active_model_name = None
+            return False
+
+    def generate_chat_response(
+        self,
+        messages: list,
+        system_prompt: str = "",
+        image=None,
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        top_k: int = 40,
+        min_p: float = 0.0,
+        max_new_tokens: int = 256,
+        repetition_penalty: float = 1.1,
+        cancel_event=None,
+    ) -> Generator[str, None, None]:
+        """Generate response, streaming tokens from subprocess."""
+        yield from self._generate_inner(
+            messages=messages,
+            system_prompt=system_prompt,
+            image=image,
+            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,
+            use_adapter=None,
+        )
+
+    def generate_with_adapter_control(
+        self,
+        use_adapter: Optional[Union[bool, str]] = None,
+        cancel_event=None,
+        **gen_kwargs,
+    ) -> Generator[str, None, None]:
+        """Generate with adapter control, streaming tokens from subprocess."""
+        yield from self._generate_inner(
+            use_adapter=use_adapter,
+            cancel_event=cancel_event,
+            **gen_kwargs,
+        )
+
+    def _generate_inner(
+        self,
+        messages: list = None,
+        system_prompt: str = "",
+        image=None,
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        top_k: int = 40,
+        min_p: float = 0.0,
+        max_new_tokens: int = 256,
+        repetition_penalty: float = 1.1,
+        cancel_event=None,
+        use_adapter=None,
+    ) -> Generator[str, None, None]:
+        """Inner generation logic — sends command to subprocess, yields tokens.
+
+        Serialized by _gen_lock: only one generation runs at a time.
+        This prevents concurrent readers from consuming each other's
+        tokens off the shared resp_queue.
+        """
+        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
+
+        # Serialize generation — single GPU, one generation at a time.
+        # Without this lock, two concurrent readers on the same resp_queue
+        # can consume and drop each other's token events.
+        with self._gen_lock:
+            yield from self._generate_locked(
+                messages=messages,
+                system_prompt=system_prompt,
+                image=image,
+                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,
+                use_adapter=use_adapter,
+            )
+
+    def _generate_locked(
+        self,
+        messages: list = None,
+        system_prompt: str = "",
+        image=None,
+        temperature: float = 0.7,
+        top_p: float = 0.9,
+        top_k: int = 40,
+        min_p: float = 0.0,
+        max_new_tokens: int = 256,
+        repetition_penalty: float = 1.1,
+        cancel_event=None,
+        use_adapter=None,
+    ) -> Generator[str, None, None]:
+        """Actual generation logic — must be called under _gen_lock."""
+        request_id = str(uuid.uuid4())
+
+        # Convert PIL Image to base64 if needed
+        image_b64 = None
+        if image is not None:
+            image_b64 = self._pil_to_base64(image)
+
+        cmd = {
+            "type": "generate",
+            "request_id": request_id,
+            "messages": messages or [],
+            "system_prompt": system_prompt,
+            "image_base64": image_b64,
+            "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
+
+        try:
+            self._send_cmd(cmd)
+        except RuntimeError as exc:
+            yield f"Error: {exc}"
+            return
+
+        # Yield tokens from response queue — we are the only reader
+        # because _gen_lock is held.
+        while True:
+            resp = self._read_resp(timeout=30.0)
+
+            if resp is None:
+                # Check subprocess health
+                if not self._ensure_subprocess_alive():
+                    yield "Error: Inference subprocess crashed during generation"
+                    return
+                continue
+
+            rtype = resp.get("type", "")
+
+            # Status messages — skip
+            if rtype == "status":
+                continue
+
+            # Error without request_id = subprocess-level error
+            resp_rid = resp.get("request_id")
+            if rtype == "error" and not resp_rid:
+                yield f"Error: {resp.get('error', 'Unknown error')}"
+                return
+
+            if rtype == "token":
+                # Check cancel from route (e.g. SSE connection closed)
+                if cancel_event is not None and cancel_event.is_set():
+                    self._cancel_generation()
+                    # Wait for the subprocess to acknowledge cancellation
+                    # (gen_done/gen_error) so stale events don't leak into
+                    # the next generation request.
+                    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
+
+    def reset_generation_state(self):
+        """Cancel any ongoing generation and reset state."""
+        self._cancel_generation()
+        if not self._ensure_subprocess_alive():
+            return
+        try:
+            self._send_cmd({"type": "reset"})
+        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)
+    # ------------------------------------------------------------------
+
+    def resize_image(self, img, max_size: int = 800):
+        """Resize image while maintaining aspect ratio.
+        No ML imports needed — runs locally in parent process.
+        """
+        if img is None:
+            return None
+        if img.size[0] > max_size or img.size[1] > max_size:
+            from PIL import Image
+            ratio = min(max_size / img.size[0], max_size / img.size[1])
+            new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio))
+            return img.resize(new_size, Image.Resampling.LANCZOS)
+        return img
+
+    @staticmethod
+    def _pil_to_base64(img) -> str:
+        """Convert a PIL Image to base64 string for IPC."""
+        buf = BytesIO()
+        img.save(buf, format="PNG")
+        return base64.b64encode(buf.getvalue()).decode("ascii")
+
+    def get_current_model(self) -> Optional[str]:
+        """Get currently active model name."""
+        return self.active_model_name
+
+    def is_model_loading(self) -> bool:
+        """Check if any model is currently loading."""
+        return len(self.loading_models) > 0
+
+    def get_loading_model(self) -> Optional[str]:
+        """Get name of currently loading model."""
+        return next(iter(self.loading_models)) if self.loading_models else None
+
+    def check_vision_model_compatibility(self) -> bool:
+        """Check if current model supports vision."""
+        if self.active_model_name and self.active_model_name in self.models:
+            return self.models[self.active_model_name].get("is_vision", False)
+        return False
+
+
+# ========== GLOBAL INSTANCE ==========
+_inference_backend = None
+
+
+def get_inference_backend() -> InferenceOrchestrator:
+    """Get global inference backend instance (orchestrator)."""
+    global _inference_backend
+    if _inference_backend is None:
+        _inference_backend = InferenceOrchestrator()
+    return _inference_backend

studio/backend/core/inference/worker.py

@@ -0,0 +1,593 @@
+"""
+Inference subprocess entry point.
+
+Each inference session runs in a persistent subprocess (mp.get_context("spawn")).
+This gives us a clean Python interpreter with no stale module state —
+solving the transformers version-switching problem completely.
+
+The subprocess stays alive while a model is loaded, accepting commands
+(generate, load, unload) via mp.Queue. It exits on shutdown or unload.
+
+Pattern follows core/training/worker.py.
+"""
+from __future__ import annotations
+
+import base64
+import logging
+import os
+import queue as _queue
+import sys
+import time
+import traceback
+from io import BytesIO
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def _activate_transformers_version(model_name: str, project_root: str) -> None:
+    """Activate the correct transformers version BEFORE any ML imports.
+
+    If the model needs transformers 5.x, prepend the pre-installed .venv_t5/
+    directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/).
+    """
+    # Ensure backend is on path for utils imports
+    backend_path = os.path.join(project_root, "studio", "backend")
+    if backend_path not in sys.path:
+        sys.path.insert(0, backend_path)
+
+    from utils.transformers_version import needs_transformers_5, _resolve_base_model
+
+    resolved = _resolve_base_model(model_name)
+    if needs_transformers_5(resolved):
+        venv_t5 = os.path.join(project_root, ".venv_t5")
+        if os.path.isdir(venv_t5):
+            sys.path.insert(0, venv_t5)
+            logger.info("Activated transformers 5.x from %s", venv_t5)
+        else:
+            # Fallback: pip install at runtime (slower, ~10-15s)
+            logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5)
+            import subprocess as sp
+            os.makedirs(venv_t5, exist_ok=True)
+            r1 = sp.run(
+                [sys.executable, "-m", "pip", "install", "--target", venv_t5,
+                 "--no-deps", "transformers==5.2.0"],
+                stdout=sp.PIPE, stderr=sp.STDOUT,
+            )
+            r2 = sp.run(
+                [sys.executable, "-m", "pip", "install", "--target", venv_t5,
+                 "--no-deps", "huggingface_hub==1.3.0"],
+                stdout=sp.PIPE, stderr=sp.STDOUT,
+            )
+            if r1.returncode != 0 or r2.returncode != 0:
+                raise RuntimeError(
+                    f"Failed to install transformers 5.x into {venv_t5}. "
+                    f"pip returncode: transformers={r1.returncode}, huggingface_hub={r2.returncode}"
+                )
+            sys.path.insert(0, venv_t5)
+        # Propagate to child subprocesses (e.g. GGUF converter)
+        _pp = os.environ.get("PYTHONPATH", "")
+        os.environ["PYTHONPATH"] = venv_t5 + (os.pathsep + _pp if _pp else "")
+    else:
+        logger.info("Using default transformers (4.57.x) for %s", model_name)
+
+
+def _decode_image(image_base64: str):
+    """Decode base64 string to PIL.Image."""
+    from PIL import Image
+    image_data = base64.b64decode(image_base64)
+    return Image.open(BytesIO(image_data))
+
+
+def _resize_image(img, max_size: int = 800):
+    """Resize image while maintaining aspect ratio."""
+    if img is None:
+        return None
+    if img.size[0] > max_size or img.size[1] > max_size:
+        from PIL import Image
+        ratio = min(max_size / img.size[0], max_size / img.size[1])
+        new_size = (int(img.size[0] * ratio), int(img.size[1] * ratio))
+        return img.resize(new_size, Image.Resampling.LANCZOS)
+    return img
+
+
+def _send_response(resp_queue: Any, response: dict) -> None:
+    """Send a response to the parent process."""
+    try:
+        resp_queue.put(response)
+    except (OSError, ValueError) as exc:
+        logger.error("Failed to send response: %s", exc)
+
+
+def _build_model_config(config: dict):
+    """Build a ModelConfig from the config dict."""
+    from utils.models import ModelConfig
+
+    model_name = config["model_name"]
+    hf_token = config.get("hf_token")
+    hf_token = hf_token if hf_token and hf_token.strip() else None
+    gguf_variant = config.get("gguf_variant")
+
+    mc = ModelConfig.from_identifier(
+        model_id=model_name,
+        hf_token=hf_token,
+        gguf_variant=gguf_variant,
+    )
+    if not mc:
+        raise ValueError(f"Invalid model identifier: {model_name}")
+    return mc
+
+
+def _handle_load(backend, config: dict, resp_queue: Any) -> None:
+    """Handle a load command: load a model into the backend."""
+    try:
+        mc = _build_model_config(config)
+
+        hf_token = config.get("hf_token")
+        hf_token = hf_token if hf_token and hf_token.strip() else None
+
+        # Auto-detect quantization for LoRA adapters
+        load_in_4bit = config.get("load_in_4bit", True)
+        if mc.is_lora and mc.path:
+            import json
+            from pathlib import Path
+            adapter_cfg_path = Path(mc.path) / "adapter_config.json"
+            if adapter_cfg_path.exists():
+                try:
+                    with open(adapter_cfg_path) as f:
+                        adapter_cfg = json.load(f)
+                    training_method = adapter_cfg.get("unsloth_training_method")
+                    if training_method == "lora" and load_in_4bit:
+                        logger.info("adapter_config.json says lora — setting load_in_4bit=False")
+                        load_in_4bit = False
+                    elif training_method == "qlora" and not load_in_4bit:
+                        logger.info("adapter_config.json says qlora — setting load_in_4bit=True")
+                        load_in_4bit = True
+                    elif not training_method:
+                        if mc.base_model and "-bnb-4bit" not in mc.base_model.lower() and load_in_4bit:
+                            logger.info("No training method, base model has no -bnb-4bit — setting load_in_4bit=False")
+                            load_in_4bit = False
+                except Exception as e:
+                    logger.warning("Could not read adapter_config.json: %s", e)
+
+        success = backend.load_model(
+            config=mc,
+            max_seq_length=config.get("max_seq_length", 2048),
+            load_in_4bit=load_in_4bit,
+            hf_token=hf_token,
+        )
+
+        if success:
+            # Build model_info for the parent to mirror
+            model_info = {
+                "identifier": mc.identifier,
+                "display_name": mc.display_name,
+                "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",
+                "success": True,
+                "model_info": model_info,
+                "ts": time.time(),
+            })
+        else:
+            _send_response(resp_queue, {
+                "type": "loaded",
+                "success": False,
+                "error": "Failed to load model",
+                "ts": time.time(),
+            })
+
+    except Exception as exc:
+        _send_response(resp_queue, {
+            "type": "loaded",
+            "success": False,
+            "error": str(exc),
+            "stack": traceback.format_exc(limit=20),
+            "ts": time.time(),
+        })
+
+
+def _handle_generate(
+    backend,
+    cmd: dict,
+    resp_queue: Any,
+    cancel_event,
+) -> None:
+    """Handle a generate command: stream tokens back via resp_queue.
+
+    cancel_event is an mp.Event shared with the parent process.
+    The parent can set it at any time (e.g. user stops generation,
+    or user loads a new model while generating) and generation
+    stops within 1-2 tokens.
+    """
+    request_id = cmd.get("request_id", "")
+
+    try:
+        # Decode image if provided
+        image = None
+        image_b64 = cmd.get("image_base64")
+        if image_b64:
+            image = _decode_image(image_b64)
+            image = _resize_image(image)
+
+        # Build generation kwargs
+        gen_kwargs = {
+            "messages": cmd["messages"],
+            "system_prompt": cmd.get("system_prompt", ""),
+            "image": image,
+            "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", 256),
+            "repetition_penalty": cmd.get("repetition_penalty", 1.1),
+            "cancel_event": cancel_event,
+        }
+
+        # Choose generation path
+        use_adapter = cmd.get("use_adapter")
+        if use_adapter is not None:
+            generator = backend.generate_with_adapter_control(
+                use_adapter=use_adapter,
+                **gen_kwargs,
+            )
+        else:
+            generator = backend.generate_chat_response(**gen_kwargs)
+
+        for cumulative_text in generator:
+            # cancel_event is an mp.Event — checked instantly, no queue polling
+            if cancel_event.is_set():
+                logger.info("Generation cancelled for request %s", request_id)
+                break
+
+            _send_response(resp_queue, {
+                "type": "token",
+                "request_id": request_id,
+                "text": cumulative_text,
+                "ts": time.time(),
+            })
+
+        _send_response(resp_queue, {
+            "type": "gen_done",
+            "request_id": request_id,
+            "ts": time.time(),
+        })
+
+    except Exception as exc:
+        logger.error("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_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", "")
+    try:
+        if model_name and model_name in backend.models:
+            backend.unload_model(model_name)
+        elif backend.active_model_name:
+            backend.unload_model(backend.active_model_name)
+
+        _send_response(resp_queue, {
+            "type": "unloaded",
+            "model_name": model_name,
+            "ts": time.time(),
+        })
+    except Exception as exc:
+        logger.error("Unload error: %s", exc)
+        _send_response(resp_queue, {
+            "type": "unloaded",
+            "model_name": model_name,
+            "error": str(exc),
+            "ts": time.time(),
+        })
+
+
+def run_inference_process(
+    *,
+    cmd_queue: Any,
+    resp_queue: Any,
+    cancel_event,
+    config: dict,
+) -> None:
+    """Subprocess entrypoint. Persistent — runs command loop until shutdown.
+
+    Args:
+        cmd_queue: mp.Queue for receiving commands from parent.
+        resp_queue: mp.Queue for sending responses to parent.
+        cancel_event: mp.Event shared with parent — set by parent to cancel generation.
+        config: Initial configuration dict with model info and project_root.
+    """
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+    project_root = config["project_root"]
+    model_name = config["model_name"]
+
+    # ── 1. Activate correct transformers version BEFORE any ML imports ──
+    try:
+        _activate_transformers_version(model_name, project_root)
+    except Exception as exc:
+        _send_response(resp_queue, {
+            "type": "error",
+            "error": f"Failed to activate transformers version: {exc}",
+            "stack": traceback.format_exc(limit=20),
+            "ts": time.time(),
+        })
+        return
+
+    # ── 1b. On Windows, check Triton availability (must be before import torch) ──
+    if sys.platform == "win32":
+        try:
+            import triton  # noqa: F401
+            logger.info("Triton available — torch.compile enabled")
+        except ImportError:
+            os.environ["TORCHDYNAMO_DISABLE"] = "1"
+            logger.warning(
+                "Triton not found on Windows — torch.compile disabled. "
+                'Install for better performance: pip install "triton-windows<3.7"'
+            )
+
+    # ── 2. Import ML libraries (fresh in this clean process) ──
+    try:
+        _send_response(resp_queue, {
+            "type": "status",
+            "message": "Importing ML libraries...",
+            "ts": time.time(),
+        })
+
+        backend_path = os.path.join(project_root, "studio", "backend")
+        if backend_path not in sys.path:
+            sys.path.insert(0, backend_path)
+
+        from core.inference.inference import InferenceBackend
+
+        import transformers
+        logger.info("Subprocess loaded transformers %s", transformers.__version__)
+
+    except Exception as exc:
+        _send_response(resp_queue, {
+            "type": "error",
+            "error": f"Failed to import ML libraries: {exc}",
+            "stack": traceback.format_exc(limit=20),
+            "ts": time.time(),
+        })
+        return
+
+    # ── 3. Create inference backend and load initial model ──
+    try:
+        backend = InferenceBackend()
+
+        _send_response(resp_queue, {
+            "type": "status",
+            "message": "Loading model...",
+            "ts": time.time(),
+        })
+
+        _handle_load(backend, config, resp_queue)
+
+    except Exception as exc:
+        _send_response(resp_queue, {
+            "type": "error",
+            "error": f"Failed to initialize inference backend: {exc}",
+            "stack": traceback.format_exc(limit=20),
+            "ts": time.time(),
+        })
+        return
+
+    # ── 4. Command loop — process commands until shutdown ──
+    # cancel_event is an mp.Event shared with parent — parent can set it
+    # at any time to cancel generation instantly (no queue polling needed).
+    logger.info("Inference subprocess ready, entering command loop")
+
+    while True:
+        try:
+            cmd = cmd_queue.get(timeout=1.0)
+        except _queue.Empty:
+            continue
+        except (EOFError, OSError):
+            logger.info("Command queue closed, shutting down")
+            return
+
+        if cmd is None:
+            continue
+
+        cmd_type = cmd.get("type", "")
+        logger.info("Received command: %s", cmd_type)
+
+        try:
+            if cmd_type == "generate":
+                cancel_event.clear()
+                _handle_generate(backend, cmd, resp_queue, cancel_event)
+
+            elif cmd_type == "load":
+                # Load a new model (reusing this subprocess)
+                # First unload current model
+                if backend.active_model_name:
+                    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)
+
+            elif cmd_type == "cancel":
+                # Redundant with mp.Event but handle gracefully
+                cancel_event.set()
+                logger.info("Cancel command received")
+
+            elif cmd_type == "reset":
+                cancel_event.set()
+                backend.reset_generation_state()
+                _send_response(resp_queue, {
+                    "type": "reset_ack",
+                    "ts": time.time(),
+                })
+
+            elif cmd_type == "status":
+                # Return current status
+                _send_response(resp_queue, {
+                    "type": "status_response",
+                    "active_model": backend.active_model_name,
+                    "models": {
+                        name: {
+                            "is_vision": info.get("is_vision", False),
+                            "is_lora": info.get("is_lora", False),
+                        }
+                        for name, info in backend.models.items()
+                    },
+                    "loading": list(backend.loading_models),
+                    "ts": time.time(),
+                })
+
+            elif cmd_type == "shutdown":
+                logger.info("Shutdown command received, exiting")
+                # Unload all models
+                for model_name in list(backend.models.keys()):
+                    try:
+                        backend.unload_model(model_name)
+                    except Exception:
+                        pass
+                _send_response(resp_queue, {
+                    "type": "shutdown_ack",
+                    "ts": time.time(),
+                })
+                return
+
+            else:
+                logger.warning("Unknown command type: %s", cmd_type)
+                _send_response(resp_queue, {
+                    "type": "error",
+                    "error": f"Unknown command type: {cmd_type}",
+                    "ts": time.time(),
+                })
+
+        except Exception as exc:
+            logger.error("Error handling command '%s': %s", cmd_type, exc, exc_info=True)
+            _send_response(resp_queue, {
+                "type": "error",
+                "error": f"Command '{cmd_type}' failed: {exc}",
+                "stack": traceback.format_exc(limit=20),
+                "ts": time.time(),
+            })

studio/backend/core/training/__init__.py

@@ -1,12 +1,9 @@
 """
 Training submodule - Training backends and trainer classes
 """
-from .trainer import UnslothTrainer, get_trainer, TrainingProgress
-from .training import TrainingBackend, get_training_backend
+from .training import TrainingBackend, TrainingProgress, get_training_backend
 
 __all__ = [
-    'UnslothTrainer',
-    'get_trainer',
     'TrainingProgress',
     'TrainingBackend',
     'get_training_backend',

studio/backend/core/training/worker.py

@@ -0,0 +1,359 @@
+"""
+Training subprocess entry point.
+
+Each training job runs in a fresh subprocess (mp.get_context("spawn")).
+This gives us a clean Python interpreter with no stale module state —
+solving the transformers version-switching problem completely.
+
+Pattern follows core/data_recipe/jobs/worker.py.
+"""
+from __future__ import annotations
+
+import logging
+import os
+import sys
+import time
+import traceback
+from pathlib import Path
+from typing import Any
+
+logger = logging.getLogger(__name__)
+
+
+def _activate_transformers_version(model_name: str, project_root: str) -> None:
+    """Activate the correct transformers version BEFORE any ML imports.
+
+    If the model needs transformers 5.x, prepend the pre-installed .venv_t5/
+    directory to sys.path. Otherwise do nothing (default 4.57.x in .venv/).
+    """
+    # Ensure backend is on path for utils imports
+    backend_path = os.path.join(project_root, "studio", "backend")
+    if backend_path not in sys.path:
+        sys.path.insert(0, backend_path)
+
+    from utils.transformers_version import needs_transformers_5, _resolve_base_model
+
+    resolved = _resolve_base_model(model_name)
+    if needs_transformers_5(resolved):
+        venv_t5 = os.path.join(project_root, ".venv_t5")
+        if os.path.isdir(venv_t5):
+            sys.path.insert(0, venv_t5)
+            logger.info("Activated transformers 5.x from %s", venv_t5)
+        else:
+            # Fallback: pip install at runtime (slower, ~10-15s)
+            logger.warning(".venv_t5 not found at %s — installing at runtime", venv_t5)
+            import subprocess as sp
+            os.makedirs(venv_t5, exist_ok=True)
+            r1 = sp.run(
+                [sys.executable, "-m", "pip", "install", "--target", venv_t5,
+                 "--no-deps", "transformers==5.2.0"],
+                stdout=sp.PIPE, stderr=sp.STDOUT,
+            )
+            r2 = sp.run(
+                [sys.executable, "-m", "pip", "install", "--target", venv_t5,
+                 "--no-deps", "huggingface_hub==1.3.0"],
+                stdout=sp.PIPE, stderr=sp.STDOUT,
+            )
+            if r1.returncode != 0 or r2.returncode != 0:
+                raise RuntimeError(
+                    f"Failed to install transformers 5.x into {venv_t5}. "
+                    f"pip returncode: transformers={r1.returncode}, huggingface_hub={r2.returncode}"
+                )
+            sys.path.insert(0, venv_t5)
+        # Propagate to child subprocesses (e.g. GGUF converter)
+        _pp = os.environ.get("PYTHONPATH", "")
+        os.environ["PYTHONPATH"] = venv_t5 + (os.pathsep + _pp if _pp else "")
+    else:
+        logger.info("Using default transformers (4.57.x) for %s", model_name)
+
+
+def run_training_process(
+    *,
+    event_queue: Any,
+    stop_queue: Any,
+    config: dict,
+) -> None:
+    """Subprocess entrypoint. Fresh Python — no stale module state.
+
+    Args:
+        event_queue: mp.Queue for sending progress/status/error events to parent.
+        stop_queue: mp.Queue for receiving stop commands from parent.
+        config: Training configuration dict with all parameters.
+    """
+    os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+    project_root = config["project_root"]
+    model_name = config["model_name"]
+
+    # ── 1. Activate correct transformers version BEFORE any ML imports ──
+    try:
+        _activate_transformers_version(model_name, project_root)
+    except Exception as exc:
+        event_queue.put({
+            "type": "error",
+            "error": f"Failed to activate transformers version: {exc}",
+            "stack": traceback.format_exc(limit=20),
+            "ts": time.time(),
+        })
+        return
+
+    # ── 1b. On Windows, check Triton availability (must be before import torch) ──
+    if sys.platform == "win32":
+        try:
+            import triton  # noqa: F401
+            logger.info("Triton available — torch.compile enabled")
+        except ImportError:
+            os.environ["TORCHDYNAMO_DISABLE"] = "1"
+            logger.warning(
+                "Triton not found on Windows — torch.compile disabled. "
+                'Install for better performance: pip install "triton-windows<3.7"'
+            )
+
+    # ── 2. Now import ML libraries (fresh in this clean process) ──
+    try:
+        _send_status(event_queue, "Importing ML libraries...")
+
+        backend_path = os.path.join(project_root, "studio", "backend")
+        if backend_path not in sys.path:
+            sys.path.insert(0, backend_path)
+
+        from core.training.trainer import UnslothTrainer, TrainingProgress
+
+        import transformers
+        logger.info("Subprocess loaded transformers %s", transformers.__version__)
+    except Exception as exc:
+        event_queue.put({
+            "type": "error",
+            "error": f"Failed to import ML libraries: {exc}",
+            "stack": traceback.format_exc(limit=20),
+            "ts": time.time(),
+        })
+        return
+
+    # ── 3. Create a fresh trainer instance ──
+    trainer = UnslothTrainer()
+
+    # Wire up progress callback → event_queue
+    def _on_progress(progress: TrainingProgress):
+        if progress.step >= 0 and progress.loss > 0:
+            event_queue.put({
+                "type": "progress",
+                "step": progress.step,
+                "epoch": progress.epoch,
+                "loss": progress.loss,
+                "learning_rate": progress.learning_rate,
+                "total_steps": progress.total_steps,
+                "elapsed_seconds": progress.elapsed_seconds,
+                "eta_seconds": progress.eta_seconds,
+                "grad_norm": progress.grad_norm,
+                "num_tokens": progress.num_tokens,
+                "eval_loss": progress.eval_loss,
+                "status_message": progress.status_message,
+                "ts": time.time(),
+            })
+        if progress.status_message:
+            _send_status(event_queue, progress.status_message)
+
+    trainer.add_progress_callback(_on_progress)
+
+    # Wire up stop_queue polling to trainer.should_stop
+    import threading
+    import queue as _queue
+
+    def _poll_stop():
+        while True:
+            try:
+                msg = stop_queue.get(timeout=1.0)
+                if msg and msg.get("type") == "stop":
+                    save = msg.get("save", True)
+                    trainer.should_stop = True
+                    trainer.save_on_stop = save
+                    logger.info("Stop signal received (save=%s)", save)
+                    return
+            except _queue.Empty:
+                continue
+            except (EOFError, OSError):
+                return
+
+    stop_thread = threading.Thread(target=_poll_stop, daemon=True)
+    stop_thread.start()
+
+    # ── 4. Execute the training pipeline ──
+    try:
+        hf_token = config.get("hf_token", "")
+        hf_token = hf_token if hf_token and hf_token.strip() else None
+
+        # Load model
+        _send_status(event_queue, "Loading model...")
+        success = trainer.load_model(
+            model_name=model_name,
+            max_seq_length=config["max_seq_length"],
+            load_in_4bit=config["load_in_4bit"],
+            hf_token=hf_token,
+            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:
+                event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
+            else:
+                event_queue.put({
+                    "type": "error",
+                    "error": trainer.training_progress.error or "Failed to load model",
+                    "stack": "", "ts": time.time(),
+                })
+            return
+
+        # Prepare model (LoRA or full finetuning)
+        training_type = config.get("training_type", "LoRA/QLoRA")
+        use_lora = (training_type == "LoRA/QLoRA")
+        if use_lora:
+            _send_status(event_queue, "Configuring LoRA adapters...")
+            success = trainer.prepare_model_for_training(
+                use_lora=True,
+                finetune_vision_layers=config.get("finetune_vision_layers", True),
+                finetune_language_layers=config.get("finetune_language_layers", True),
+                finetune_attention_modules=config.get("finetune_attention_modules", True),
+                finetune_mlp_modules=config.get("finetune_mlp_modules", True),
+                target_modules=config.get("target_modules"),
+                lora_r=config.get("lora_r", 16),
+                lora_alpha=config.get("lora_alpha", 16),
+                lora_dropout=config.get("lora_dropout", 0.0),
+                use_gradient_checkpointing=config.get("gradient_checkpointing", "unsloth"),
+                use_rslora=config.get("use_rslora", False),
+                use_loftq=config.get("use_loftq", False),
+            )
+        else:
+            _send_status(event_queue, "Preparing model for full finetuning...")
+            success = trainer.prepare_model_for_training(use_lora=False)
+
+        if not success or trainer.should_stop:
+            if trainer.should_stop:
+                event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
+            else:
+                event_queue.put({
+                    "type": "error",
+                    "error": trainer.training_progress.error or "Failed to prepare model",
+                    "stack": "", "ts": time.time(),
+                })
+            return
+
+        # Load dataset
+        _send_status(event_queue, "Loading and formatting dataset...")
+        hf_dataset = config.get("hf_dataset", "")
+        dataset_result = trainer.load_and_format_dataset(
+            dataset_source=hf_dataset if hf_dataset and hf_dataset.strip() else None,
+            format_type=config.get("format_type", ""),
+            local_datasets=config.get("local_datasets") or None,
+            custom_format_mapping=config.get("custom_format_mapping"),
+            subset=config.get("subset"),
+            train_split=config.get("train_split", "train"),
+            eval_split=config.get("eval_split"),
+            eval_steps=config.get("eval_steps", 0.00),
+            dataset_slice_start=config.get("dataset_slice_start"),
+            dataset_slice_end=config.get("dataset_slice_end"),
+        )
+
+        if isinstance(dataset_result, tuple):
+            dataset, eval_dataset = dataset_result
+        else:
+            dataset = dataset_result
+            eval_dataset = None
+
+        # Disable eval if eval_steps <= 0
+        eval_steps = config.get("eval_steps", 0.00)
+        if eval_steps is not None and float(eval_steps) <= 0:
+            eval_dataset = None
+
+        if dataset is None or trainer.should_stop:
+            if trainer.should_stop:
+                event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
+            else:
+                event_queue.put({
+                    "type": "error",
+                    "error": trainer.training_progress.error or "Failed to load dataset",
+                    "stack": "", "ts": time.time(),
+                })
+            return
+
+        # Convert learning rate
+        try:
+            lr_value = float(config.get("learning_rate", "2e-4"))
+        except ValueError:
+            event_queue.put({
+                "type": "error",
+                "error": f"Invalid learning rate: {config.get('learning_rate')}",
+                "stack": "", "ts": time.time(),
+            })
+            return
+
+        # Generate output dir
+        output_dir = config.get("output_dir")
+        if not output_dir:
+            output_dir = f"./outputs/{model_name.replace('/', '_')}_{int(time.time())}"
+
+        # Start training (directly — no inner thread, we ARE the subprocess)
+        _send_status(event_queue, "Starting training...")
+        max_steps = config.get("max_steps", 0)
+        save_steps = config.get("save_steps", 0)
+
+        trainer._train_worker(
+            dataset,
+            output_dir=output_dir,
+            num_epochs=config.get("num_epochs", 3),
+            learning_rate=lr_value,
+            batch_size=config.get("batch_size", 2),
+            gradient_accumulation_steps=config.get("gradient_accumulation_steps", 4),
+            warmup_steps=config.get("warmup_steps"),
+            warmup_ratio=config.get("warmup_ratio"),
+            max_steps=max_steps if max_steps and max_steps > 0 else 0,
+            save_steps=save_steps if save_steps and save_steps > 0 else 0,
+            weight_decay=config.get("weight_decay", 0.01),
+            random_seed=config.get("random_seed", 3407),
+            packing=config.get("packing", False),
+            train_on_completions=config.get("train_on_completions", False),
+            enable_wandb=config.get("enable_wandb", False),
+            wandb_project=config.get("wandb_project", "unsloth-training"),
+            wandb_token=config.get("wandb_token"),
+            enable_tensorboard=config.get("enable_tensorboard", False),
+            tensorboard_dir=config.get("tensorboard_dir", "runs"),
+            eval_dataset=eval_dataset,
+            eval_steps=eval_steps,
+            max_seq_length=config.get("max_seq_length", 2048),
+            optim=config.get("optim", "adamw_8bit"),
+            lr_scheduler_type=config.get("lr_scheduler_type", "linear"),
+        )
+
+        # Check final state
+        progress = trainer.get_training_progress()
+        if progress.error:
+            event_queue.put({
+                "type": "error",
+                "error": progress.error,
+                "stack": "",
+                "ts": time.time(),
+            })
+        else:
+            event_queue.put({
+                "type": "complete",
+                "output_dir": output_dir,
+                "status_message": progress.status_message or "Training completed",
+                "ts": time.time(),
+            })
+
+    except Exception as exc:
+        event_queue.put({
+            "type": "error",
+            "error": str(exc),
+            "stack": traceback.format_exc(limit=20),
+            "ts": time.time(),
+        })
+
+
+def _send_status(event_queue: Any, message: str) -> None:
+    """Send a status update to the parent process."""
+    event_queue.put({
+        "type": "status",
+        "message": message,
+        "ts": time.time(),
+    })

studio/backend/main.py

@@ -3,6 +3,7 @@ Main FastAPI application for Unsloth UI Backend
 """
 import os
 import secrets
+import shutil
 from contextlib import asynccontextmanager
 
 from fastapi import FastAPI
@@ -35,6 +36,12 @@ async def lifespan(app: FastAPI):
     # Clean up any stale compiled cache from previous runs
     clear_unsloth_compiled_cache()
 
+    # Remove stale .venv_overlay from previous versions — no longer used.
+    # Version switching now uses .venv_t5/ (pre-installed by setup.sh).
+    overlay_dir = Path(__file__).resolve().parent.parent.parent / ".venv_overlay"
+    if overlay_dir.is_dir():
+        shutil.rmtree(overlay_dir, ignore_errors=True)
+
     # Detect hardware first — sets DEVICE global used everywhere
     detect_hardware()
 
@@ -89,6 +96,11 @@ app.include_router(auth_router, prefix="/api/auth", tags=["auth"])
 app.include_router(training_router, prefix="/api/train", tags=["training"])
 app.include_router(models_router, prefix="/api/models", tags=["models"])
 app.include_router(inference_router, prefix="/api/inference", tags=["inference"])
+
+# OpenAI-compatible endpoints: mount the same inference router at /v1
+# so external tools (Open WebUI, SillyTavern, etc.) can use the
+# standard /v1/chat/completions path.
+app.include_router(inference_router, prefix="/v1", tags=["openai-compat"])
 app.include_router(datasets_router, prefix="/api/datasets", tags=["datasets"])
 app.include_router(data_recipe_router, prefix="/api/data-recipe", tags=["data-recipe"])
 app.include_router(export_router, prefix="/api/export", tags=["export"])

studio/backend/models/datasets.py

@@ -28,11 +28,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

@@ -353,7 +353,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
@@ -396,11 +396,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/export.py

@@ -61,6 +61,42 @@ async def load_checkpoint(
     Wraps ExportBackend.load_checkpoint.
     """
     try:
+        # Version switching is handled automatically by the subprocess-based
+        # export backend — no need for ensure_transformers_version() here.
+
+        # Free GPU memory: shut down any running inference/training subprocesses
+        # before loading the export checkpoint (they'd compete for VRAM).
+        try:
+            from core.inference import get_inference_backend
+            inf = get_inference_backend()
+            if inf.active_model_name:
+                logger.info(
+                    "Unloading inference model '%s' to free GPU memory for export",
+                    inf.active_model_name,
+                )
+                inf._shutdown_subprocess()
+                inf.active_model_name = None
+                inf.models.clear()
+        except Exception as e:
+            logger.warning("Could not unload inference model: %s", e)
+
+        try:
+            from core.training import get_training_backend
+            trn = get_training_backend()
+            if trn.is_training_active():
+                logger.info("Stopping active training to free GPU memory for export")
+                trn.stop_training()
+                # Wait for training subprocess to actually exit before proceeding,
+                # otherwise it may still hold GPU memory when export tries to load.
+                for _ in range(60):  # up to 30s
+                    if not trn.is_training_active():
+                        break
+                    import time; time.sleep(0.5)
+                else:
+                    logger.warning("Training subprocess did not exit within 30s, proceeding anyway")
+        except Exception as e:
+            logger.warning("Could not stop training: %s", e)
+
         backend = get_export_backend()
         success, message = backend.load_checkpoint(
             checkpoint_path=request.checkpoint_path,

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__)
 
@@ -86,6 +91,9 @@ async def load_model(
     GGUF models are loaded via llama-server (llama.cpp) instead of Unsloth.
     """
     try:
+        # Version switching is handled automatically by the subprocess-based
+        # inference backend — no need for ensure_transformers_version() here.
+
         # Create config using clean factory method
         # is_lora is auto-detected from adapter_config.json on disk/HF
         config = ModelConfig.from_identifier(
@@ -160,10 +168,63 @@ async def load_model(
             logger.info("Unloading GGUF model before loading Unsloth model")
             llama_backend.unload_model()
 
+        # Shut down any export subprocess to free VRAM
+        try:
+            from core.export import get_export_backend
+            exp_backend = get_export_backend()
+            if exp_backend.current_checkpoint:
+                logger.info("Shutting down export subprocess to free GPU memory for inference")
+                exp_backend._shutdown_subprocess()
+                exp_backend.current_checkpoint = None
+                exp_backend.is_vision = False
+                exp_backend.is_peft = False
+        except Exception as e:
+            logger.warning("Could not shut down export subprocess: %s", e)
+
+        # Auto-detect quantization for LoRA adapters from adapter_config.json
+        # The training pipeline patches this file with "unsloth_training_method"
+        # which is 'qlora' or 'lora'. Only LoRA (16-bit) needs load_in_4bit=False.
+        load_in_4bit = request.load_in_4bit
+        if config.is_lora and config.path:
+            import json
+            from pathlib import Path
+            adapter_cfg_path = Path(config.path) / "adapter_config.json"
+            if adapter_cfg_path.exists():
+                try:
+                    with open(adapter_cfg_path) as f:
+                        adapter_cfg = json.load(f)
+                    training_method = adapter_cfg.get("unsloth_training_method")
+                    if training_method == "lora" and load_in_4bit:
+                        logger.info(
+                            f"adapter_config.json says unsloth_training_method='lora' — "
+                            f"setting load_in_4bit=False to match 16-bit training"
+                        )
+                        load_in_4bit = False
+                    elif training_method == "qlora" and not load_in_4bit:
+                        logger.info(
+                            f"adapter_config.json says unsloth_training_method='qlora' — "
+                            f"setting load_in_4bit=True to match QLoRA training"
+                        )
+                        load_in_4bit = True
+                    elif training_method:
+                        logger.info(f"Training method: {training_method}, load_in_4bit={load_in_4bit}")
+                    else:
+                        # No unsloth_training_method — fallback to base model name
+                        if config.base_model and "-bnb-4bit" not in config.base_model.lower() and load_in_4bit:
+                            logger.info(
+                                f"No unsloth_training_method in adapter_config.json. "
+                                f"Base model '{config.base_model}' has no -bnb-4bit suffix — "
+                                f"setting load_in_4bit=False"
+                            )
+                            load_in_4bit = False
+                except Exception as e:
+                    logger.warning(f"Could not read adapter_config.json: {e}")
+
+        # Load the model
         success = backend.load_model(
             config=config,
             max_seq_length=request.max_seq_length,
-            load_in_4bit=request.load_in_4bit,
+            load_in_4bit=load_in_4bit,
             hf_token=request.hf_token,
         )
 
@@ -185,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,
         )
 
@@ -331,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()),
         )
@@ -351,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]"]:
@@ -445,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
@@ -729,3 +995,40 @@ async def openai_chat_completions(
             backend.reset_generation_state()
             logger.error(f"Error during OpenAI completion: {e}", exc_info=True)
             raise HTTPException(status_code=500, detail=str(e))
+
+
+# =====================================================================
+# OpenAI-Compatible Models Listing  (/models → /v1/models)
+# =====================================================================
+
+@router.get("/models")
+async def openai_list_models(
+    current_subject: str = Depends(get_current_subject),
+):
+    """
+    OpenAI-compatible model listing endpoint.
+
+    Returns the currently loaded model in the format expected by
+    OpenAI-compatible clients (``GET /v1/models``).
+    """
+    models = []
+
+    # Check GGUF backend
+    llama_backend = get_llama_cpp_backend()
+    if llama_backend.is_loaded:
+        models.append({
+            "id": llama_backend.model_identifier,
+            "object": "model",
+            "owned_by": "local",
+        })
+
+    # Check Unsloth backend
+    backend = get_inference_backend()
+    if backend.active_model_name:
+        models.append({
+            "id": backend.active_model_name,
+            "object": "model",
+            "owned_by": "local",
+        })
+
+    return {"object": "list", "data": models}

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,41 +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)
-        
-        # Check if it's a vision model
+
+        # 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

@@ -9,7 +9,6 @@ from typing import Dict, Optional, Any
 import logging
 import asyncio
 from datetime import datetime
-import threading
 
 # Add backend directory to path
 # The backend code should be in the same directory structure
@@ -85,6 +84,11 @@ async def start_training(
     """
     try:
         logger.info(f"Starting training job with model: {request.model_name}")
+
+        # NOTE: No in-process ensure_transformers_version() call here.
+        # The subprocess (worker.py) activates the correct version in a
+        # fresh Python interpreter before importing any ML libraries.
+
         backend = get_training_backend()
 
         # Generate job ID and attach to backend for later status/progress calls
@@ -180,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 "",
@@ -188,84 +193,50 @@ async def start_training(
             "tensorboard_dir": request.tensorboard_dir or "",
         }
 
-        # Set initial "preparing" state
+        # Free GPU memory: shut down any running inference/export subprocesses
+        # before training starts (they'd compete for VRAM otherwise)
         try:
-            backend.trainer._update_progress(
-                status_message="Initializing training...",
-                is_training=False,
-            )
-        except Exception:
-            pass
-
-        def run_training():
-            try:
+            from core.inference import get_inference_backend
+            inf_backend = get_inference_backend()
+            if inf_backend.active_model_name:
                 logger.info(
-                    f"Starting training job {job_id} with model {request.model_name}"
+                    "Unloading inference model '%s' to free GPU memory for training",
+                    inf_backend.active_model_name,
                 )
+                inf_backend._shutdown_subprocess()
+                inf_backend.active_model_name = None
+                inf_backend.models.clear()
+        except Exception as e:
+            logger.warning("Could not unload inference model: %s", e)
 
-                # Update status to show we're loading model
-                try:
-                    backend.trainer._update_progress(status_message="Loading model...")
-                except Exception as e:
-                    logger.error(f"Error updating progress: {e}")
+        try:
+            from core.export import get_export_backend
+            exp_backend = get_export_backend()
+            if exp_backend.current_checkpoint:
+                logger.info("Shutting down export subprocess to free GPU memory for training")
+                exp_backend._shutdown_subprocess()
+                exp_backend.current_checkpoint = None
+                exp_backend.is_vision = False
+                exp_backend.is_peft = False
+        except Exception as e:
+            logger.warning("Could not shut down export subprocess: %s", e)
 
-                # start_training returns bool (not generator)
-                run_result = backend.start_training(**training_kwargs)
-                logger.info(
-                    "Training job %s backend.start_training returned type=%s value=%r",
-                    job_id,
-                    type(run_result).__name__,
-                    run_result,
-                )
-                if not run_result:
-                    progress_error = backend.trainer.training_progress.error
-                    raise RuntimeError(progress_error or "Training failed to start")
+        # start_training now spawns a subprocess (non-blocking)
+        success = backend.start_training(**training_kwargs)
 
-                logger.info(f"Training job {job_id} started successfully")
-
-            except Exception as e:
-                logger.error(f"Training error in job {job_id}: {e}", exc_info=True)
-                try:
-                    backend.trainer._update_progress(
-                        error=str(e),
-                        is_training=False,
-                    )
-                except Exception as update_error:
-                    logger.error(f"Failed to update progress: {update_error}")
-
-        # Start training in a daemon thread
-        training_thread = threading.Thread(
-            target=run_training,
-            daemon=True,
-            name=f"Training-{job_id}",
-        )
-        training_thread.start()
-
-        # Store thread reference for status checking
-        backend._training_thread = training_thread
-
-        # Give it a moment to start
-        import time
-
-        time.sleep(0.5)
-
-        # Verify training thread is alive
-        if not training_thread.is_alive():
-            logger.warning(f"Training thread died immediately for job {job_id}")
+        if not success:
+            progress_error = backend.trainer.training_progress.error
             return TrainingJobResponse(
                 job_id=job_id,
                 status="error",
-                message=(
-                    "Training thread failed to start. "
-                    "Check server logs for details."
-                ),
-                error="Thread not alive",
+                message=progress_error or "Failed to start training subprocess",
+                error=progress_error or "subprocess_start_failed",
             )
 
         return TrainingJobResponse(
             job_id=job_id,
             status="queued",
-            message="Training job queued and starting in background",
+            message="Training job queued and starting in subprocess",
             error=None,
         )
 
@@ -290,18 +261,10 @@ async def stop_training(
     """
     try:
         backend = get_training_backend()
-        trainer_thread = getattr(getattr(backend, "trainer", None), "training_thread", None)
-        thread_alive = bool(trainer_thread and trainer_thread.is_alive())
         is_active = backend.is_training_active()
-        logger.info(
-            "Stop requested: save=%s is_active=%s thread_alive=%s should_stop=%s",
-            body.save,
-            is_active,
-            thread_alive,
-            getattr(getattr(backend, "trainer", None), "should_stop", None),
-        )
+        logger.info("Stop requested: save=%s is_active=%s", body.save, is_active)
 
-        if not is_active and not thread_alive:
+        if not is_active:
             return TrainingStopResponse(
                 status="idle",
                 message="No training job is currently running"
@@ -332,25 +295,22 @@ async def reset_training(
     """
     try:
         backend = get_training_backend()
-        trainer_thread = getattr(getattr(backend, "trainer", None), "training_thread", None)
-        thread_alive = bool(trainer_thread and trainer_thread.is_alive())
         is_active = backend.is_training_active()
 
-        if is_active or thread_alive:
-            logger.warning(
-                "Rejected reset while training active: is_active=%s thread_alive=%s should_stop=%s",
-                is_active,
-                thread_alive,
-                getattr(getattr(backend, "trainer", None), "should_stop", None),
-            )
+        if is_active:
+            logger.warning("Rejected reset while training active: is_active=%s", is_active)
             raise HTTPException(
                 status_code=409,
                 detail="Training is still running. Stop training and wait for it to finish before resetting.",
             )
 
         logger.info("Reset training state: clearing runtime + metric history")
-        backend.trainer.should_stop = False
-        backend.trainer.training_progress = backend.trainer.training_progress.__class__()
+        backend._should_stop = False  # Clear stop flag so status returns to idle
+        backend.trainer._update_progress(
+            is_training=False, is_completed=False, error=None,
+            status_message="Ready to train", step=0, loss=0.0, epoch=0,
+            total_steps=0,
+        )
         backend.loss_history = []
         backend.lr_history = []
         backend.step_history = []
@@ -376,18 +336,11 @@ async def get_training_status(
     """
     try:
         backend = get_training_backend()
-        job_id: str = getattr(backend, "current_job_id", "")
+        job_id: str = getattr(backend, "current_job_id", "") or ""
 
         # Check if training is active
         is_active = backend.is_training_active()
 
-        # Check if there's a training thread running (preparation phase)
-        has_thread = (
-            hasattr(backend, "_training_thread")
-            and backend._training_thread
-            and backend._training_thread.is_alive()
-        )
-
         # Get progress info from trainer
         try:
             progress = backend.trainer.get_training_progress()
@@ -400,14 +353,14 @@ async def get_training_status(
         error_message = getattr(progress, "error", None) if progress else None
 
         # Check if training was stopped by user
-        trainer_stopped = getattr(backend.trainer, "should_stop", False)
+        trainer_stopped = getattr(backend, "_should_stop", False)
 
         # Derive high-level phase
         if error_message:
             phase = "error"
         elif is_active:
             msg_lower = status_message.lower()
-            if "loading" in msg_lower:
+            if "loading" in msg_lower or "importing" in msg_lower:
                 phase = "loading_model"
             elif any(
                 k in msg_lower for k in ["preparing", "initializing", "configuring"]
@@ -419,8 +372,6 @@ async def get_training_status(
             phase = "stopped"
         elif progress and getattr(progress, "is_completed", False):
             phase = "completed"
-        elif has_thread:
-            phase = "loading_model"
         else:
             phase = "idle"
 
@@ -534,7 +485,7 @@ async def stream_training_progress(
 
     async def event_generator():
         backend = get_training_backend()
-        job_id: str = getattr(backend, "current_job_id", "")
+        job_id: str = getattr(backend, "current_job_id", "") or ""
 
         # ── Helpers ──────────────────────────────────────────────
         def build_progress(

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

@@ -7,6 +7,9 @@ from typing import Optional, Dict, Any
 from utils.paths import normalize_path, is_local_path, is_model_cached
 from utils.utils import without_hf_auth
 import logging
+import os
+import subprocess
+import sys
 from pathlib import Path
 from typing import List, Tuple
 import json
@@ -213,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",
@@ -317,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",
@@ -364,7 +375,117 @@ def load_model_config(model_name: str, use_auth: bool = False, token: Optional[s
         model_name,
         trust_remote_code=True
     )
-pass
+
+
+# VLM architecture suffixes and known VLM model_type values.
+_VLM_ARCH_SUFFIXES = ("ForConditionalGeneration", "ForVisionText2Text")
+_VLM_MODEL_TYPES = {
+    'phi3_v', 'llava', 'llava_next', 'llava_onevision',
+    'internvl_chat', 'cogvlm2', 'minicpmv',
+}
+
+# Pre-computed project root and .venv_t5 path for subprocess version switching.
+_PROJECT_ROOT = str(Path(__file__).resolve().parent.parent.parent.parent.parent)
+_VENV_T5_DIR = os.path.join(_PROJECT_ROOT, ".venv_t5")
+_BACKEND_DIR = os.path.join(_PROJECT_ROOT, "studio", "backend")
+
+# Inline script executed in a subprocess with transformers 5.x activated.
+# Receives model_name and token via argv, prints JSON result to stdout.
+_VISION_CHECK_SCRIPT = r'''
+import sys, os, json
+os.environ["TOKENIZERS_PARALLELISM"] = "false"
+
+# Activate transformers 5.x
+venv_t5 = sys.argv[1]
+backend_dir = sys.argv[2]
+model_name = sys.argv[3]
+token = sys.argv[4] if len(sys.argv) > 4 and sys.argv[4] != "" else None
+
+sys.path.insert(0, venv_t5)
+if backend_dir not in sys.path:
+    sys.path.insert(0, backend_dir)
+
+try:
+    from transformers import AutoConfig
+    kwargs = {"trust_remote_code": True}
+    if token:
+        kwargs["token"] = token
+    config = AutoConfig.from_pretrained(model_name, **kwargs)
+
+    is_vlm = False
+    if hasattr(config, "architectures"):
+        is_vlm = any(
+            x.endswith(("ForConditionalGeneration", "ForVisionText2Text"))
+            for x in config.architectures
+        )
+    if not is_vlm and hasattr(config, "vision_config"):
+        is_vlm = True
+    if not is_vlm and hasattr(config, "img_processor"):
+        is_vlm = True
+    if not is_vlm and hasattr(config, "image_token_index"):
+        is_vlm = True
+    if not is_vlm and hasattr(config, "model_type"):
+        vlm_types = {"phi3_v","llava","llava_next","llava_onevision",
+                      "internvl_chat","cogvlm2","minicpmv"}
+        if config.model_type in vlm_types:
+            is_vlm = True
+
+    model_type = getattr(config, "model_type", "unknown")
+    archs = getattr(config, "architectures", [])
+    print(json.dumps({"is_vision": is_vlm, "model_type": model_type,
+                       "architectures": archs}))
+except Exception as exc:
+    print(json.dumps({"error": str(exc)}))
+    sys.exit(1)
+'''
+
+
+def _is_vision_model_subprocess(model_name: str, hf_token: Optional[str] = None) -> bool:
+    """Run is_vision_model check in a subprocess with transformers 5.x.
+
+    Same pattern as training/inference workers: spawn a clean subprocess
+    with .venv_t5/ prepended to sys.path so AutoConfig recognizes newer
+    architectures (glm4_moe_lite, etc.).
+    """
+    token_arg = hf_token or ""
+
+    try:
+        result = subprocess.run(
+            [sys.executable, "-c", _VISION_CHECK_SCRIPT,
+             _VENV_T5_DIR, _BACKEND_DIR, model_name, token_arg],
+            capture_output=True, text=True, timeout=60,
+        )
+
+        if result.returncode != 0:
+            stderr = result.stderr.strip()
+            logger.warning(
+                "Vision check subprocess failed for '%s': %s",
+                model_name, stderr or result.stdout.strip(),
+            )
+            return False
+
+        data = json.loads(result.stdout.strip())
+        if "error" in data:
+            logger.warning(
+                "Vision check subprocess error for '%s': %s",
+                model_name, data["error"],
+            )
+            return False
+
+        is_vlm = data["is_vision"]
+        logger.info(
+            "Vision check (subprocess, transformers 5.x) for '%s': "
+            "model_type=%s, architectures=%s, is_vision=%s",
+            model_name, data.get("model_type"), data.get("architectures"), is_vlm,
+        )
+        return is_vlm
+
+    except subprocess.TimeoutExpired:
+        logger.warning("Vision check subprocess timed out for '%s'", model_name)
+        return False
+    except Exception as exc:
+        logger.warning("Vision check subprocess failed for '%s': %s", model_name, exc)
+        return False
 
 
 def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
@@ -372,17 +493,39 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
     Detect vision-language models (VLMs) by checking architecture in config.
     Works for fine-tuned models since they inherit the base architecture.
 
+    For models that require transformers 5.x (e.g. GLM-4.7-Flash), the check
+    runs in a subprocess with .venv_t5/ activated — same pattern as the
+    training and inference workers.
+
     Args:
         model_name: Model identifier (HF repo or local path)
         hf_token: Optional HF token for accessing gated/private models
     """
+    # Models that need transformers 5.x must be checked in a subprocess
+    # because AutoConfig in the main process (transformers 4.57.x) doesn't
+    # recognize their architectures.
+    from utils.transformers_version import needs_transformers_5
+    if needs_transformers_5(model_name):
+        logger.info(
+            "Model '%s' needs transformers 5.x — checking vision via subprocess",
+            model_name,
+        )
+        return _is_vision_model_subprocess(model_name, hf_token=hf_token)
+
     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(
-                x.endswith(("ForConditionalGeneration", "ForVisionText2Text"))
+                x.endswith(_VLM_ARCH_SUFFIXES)
                 for x in config.architectures
             )
             if is_vlm:
@@ -406,11 +549,7 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
 
         # Check 5: Known VLM model_type values that may not match above checks
         if hasattr(config, 'model_type'):
-            vlm_model_types = {
-                'phi3_v', 'llava', 'llava_next', 'llava_onevision',
-                'internvl_chat', 'cogvlm2', 'minicpmv',
-            }
-            if config.model_type in vlm_model_types:
+            if config.model_type in _VLM_MODEL_TYPES:
                 logger.info(f"Model {model_name} detected as VLM: model_type={config.model_type}")
                 return True
 
@@ -419,7 +558,115 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool:
     except Exception as e:
         logger.warning(f"Could not determine if {model_name} is vision model: {e}")
         return False
-pass
+
+
+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:
@@ -905,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
@@ -941,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")
@@ -977,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
 
@@ -988,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,
             )
 
@@ -1165,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,
@@ -1179,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,
         )
 
@@ -1258,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/backend/utils/transformers_version.py

@@ -0,0 +1,266 @@
+"""
+Automatic transformers version switching.
+
+Some newer model architectures (Ministral-3, GLM-4.7-Flash, Qwen3-30B-A3B MoE,
+tiny_qwen3_moe) require transformers>=5.2.0, while everything else needs the
+default 4.57.x that ships with Unsloth.
+
+When loading a LoRA adapter with a custom name, we resolve the base model from
+``adapter_config.json`` and check *that* against the model list.
+
+Strategy:
+  Training and inference run in subprocesses that activate the correct version
+  via sys.path (prepending .venv_t5/ for 5.x models). See:
+    - core/training/worker.py
+    - core/inference/worker.py
+
+  For export (still in-process), ensure_transformers_version() does a lightweight
+  sys.path swap using the same .venv_t5/ directory pre-installed by setup.sh.
+"""
+
+import importlib
+import json
+import logging
+import os
+import subprocess
+import sys
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+# Ensure our logger is visible even if root logger isn't configured for INFO.
+if not logger.handlers:
+    _handler = logging.StreamHandler()
+    _handler.setLevel(logging.INFO)
+    _handler.setFormatter(
+        logging.Formatter("[%(name)s|%(levelname)s]%(message)s")
+    )
+    logger.addHandler(_handler)
+    logger.setLevel(logging.INFO)
+
+# ---------------------------------------------------------------------------
+# Detection
+# ---------------------------------------------------------------------------
+
+# Lowercase substrings — if ANY appears anywhere in the lowered model name,
+# we need transformers 5.x.
+TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = (
+    "ministral-3-",       # Ministral-3-{3,8,14}B-{Instruct,Reasoning,Base}-2512
+    "glm-4.7-flash",      # GLM-4.7-Flash
+    "qwen3-30b-a3b",      # Qwen3-30B-A3B-Instruct-2507 and variants
+    "qwen3.5",            # Qwen3.5 family (35B-A3B, etc.)
+    "qwen3-next",         # Qwen3-Next and variants
+    "tiny_qwen3_moe",     # imdatta0/tiny_qwen3_moe_2.8B_0.7B
+)
+
+# Versions
+TRANSFORMERS_5_VERSION = "5.2.0"
+TRANSFORMERS_DEFAULT_VERSION = "4.57.1"
+
+# Pre-installed directory for transformers 5.x — created by setup.sh / setup.ps1
+_PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent.parent  # studio/backend/utils/ → project root
+_VENV_T5_DIR = str(_PROJECT_ROOT / ".venv_t5")
+
+
+def _resolve_base_model(model_name: str) -> str:
+    """If *model_name* points to a LoRA adapter, return its base model.
+
+    Checks for ``adapter_config.json`` locally first.  Only calls the heavier
+    ``get_base_model_from_lora`` for paths that are actual local directories
+    (avoids noisy warnings for plain HF model IDs).
+
+    Returns the original *model_name* unchanged if it is not a LoRA adapter.
+    """
+    # --- Fast local check ---------------------------------------------------
+    local_path = Path(model_name)
+    adapter_cfg_path = local_path / "adapter_config.json"
+    if adapter_cfg_path.is_file():
+        try:
+            with open(adapter_cfg_path) as f:
+                cfg = json.load(f)
+            base = cfg.get("base_model_name_or_path")
+            if base:
+                logger.info(
+                    "Resolved LoRA adapter '%s' → base model '%s'",
+                    model_name, base,
+                )
+                return base
+        except Exception as exc:
+            logger.debug("Could not read %s: %s", adapter_cfg_path, exc)
+
+    # --- Only try the heavier fallback for local directories ----------------
+    if local_path.is_dir():
+        try:
+            from utils.models import get_base_model_from_lora
+            base = get_base_model_from_lora(model_name)
+            if base:
+                logger.info(
+                    "Resolved LoRA adapter '%s' → base model '%s' "
+                    "(via get_base_model_from_lora)",
+                    model_name, base,
+                )
+                return base
+        except Exception as exc:
+            logger.debug(
+                "get_base_model_from_lora failed for '%s': %s",
+                model_name, exc,
+            )
+
+    return model_name
+
+
+def needs_transformers_5(model_name: str) -> bool:
+    """Return True if *model_name* belongs to an architecture that requires
+    ``transformers>=5.2.0``."""
+    lowered = model_name.lower()
+    return any(sub in lowered for sub in TRANSFORMERS_5_MODEL_SUBSTRINGS)
+
+
+# ---------------------------------------------------------------------------
+# Version switching (in-process — used only by export)
+# ---------------------------------------------------------------------------
+
+def _get_in_memory_version() -> str | None:
+    """Return the transformers version currently loaded in this process."""
+    tf = sys.modules.get("transformers")
+    if tf is not None:
+        return getattr(tf, "__version__", None)
+    return None
+
+
+# All top-level prefixes that hold references to transformers internals.
+_PURGE_PREFIXES = (
+    "transformers",
+    "huggingface_hub",
+    "unsloth",
+    "unsloth_zoo",
+    "peft",
+    "trl",
+    "accelerate",
+    "auto_gptq",
+    # NOTE: bitsandbytes is intentionally EXCLUDED — it registers torch custom
+    # operators at import time via torch.library.define(). Those registrations
+    # live in torch's global operator registry which survives module purge.
+    # Re-importing bitsandbytes after purge → duplicate registration → crash.
+    # Our own modules that import from transformers at module level
+    # (e.g. model_config.py: `from transformers import AutoConfig`)
+    "utils.models",
+    "core.training",
+    "core.inference",
+    "core.export",
+)
+
+
+def _purge_modules() -> int:
+    """Remove all cached modules for transformers and its dependents.
+
+    Returns the number of modules purged.
+    """
+    importlib.invalidate_caches()
+    to_remove = [
+        k for k in list(sys.modules.keys())
+        if any(k == p or k.startswith(p + ".") for p in _PURGE_PREFIXES)
+    ]
+    for key in to_remove:
+        del sys.modules[key]
+    return len(to_remove)
+
+
+def _ensure_venv_t5_exists() -> bool:
+    """Ensure .venv_t5/ exists. Install at runtime if missing."""
+    if os.path.isdir(_VENV_T5_DIR) and os.listdir(_VENV_T5_DIR):
+        return True
+
+    logger.warning(".venv_t5 not found at %s — installing at runtime", _VENV_T5_DIR)
+    os.makedirs(_VENV_T5_DIR, exist_ok=True)
+    for pkg in (f"transformers=={TRANSFORMERS_5_VERSION}", "huggingface_hub==1.3.0"):
+        cmd = [
+            sys.executable, "-m", "pip", "install",
+            "--target", _VENV_T5_DIR,
+            "--no-deps",
+            pkg,
+        ]
+        result = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)
+        if result.returncode != 0:
+            logger.error("pip install failed:\n%s", result.stdout)
+            return False
+    logger.info("Installed transformers 5.x to %s", _VENV_T5_DIR)
+    return True
+
+
+def _activate_5x() -> None:
+    """Prepend .venv_t5/ to sys.path, purge stale modules, reimport."""
+    if not _ensure_venv_t5_exists():
+        raise RuntimeError(f"Cannot activate transformers 5.x: .venv_t5 missing at {_VENV_T5_DIR}")
+
+    if _VENV_T5_DIR not in sys.path:
+        sys.path.insert(0, _VENV_T5_DIR)
+        logger.info("Prepended %s to sys.path", _VENV_T5_DIR)
+
+    count = _purge_modules()
+    logger.info("Purged %d cached modules", count)
+
+    import transformers
+    logger.info("Loaded transformers %s", transformers.__version__)
+
+
+def _deactivate_5x() -> None:
+    """Remove .venv_t5/ from sys.path, purge stale modules, reimport."""
+    while _VENV_T5_DIR in sys.path:
+        sys.path.remove(_VENV_T5_DIR)
+    logger.info("Removed %s from sys.path", _VENV_T5_DIR)
+
+    count = _purge_modules()
+    logger.info("Purged %d cached modules", count)
+
+    import transformers
+    logger.info("Reverted to transformers %s", transformers.__version__)
+
+
+def ensure_transformers_version(model_name: str) -> None:
+    """Ensure the correct ``transformers`` version is active for *model_name*.
+
+    Uses sys.path with .venv_t5/ (pre-installed by setup.sh):
+      • Need 5.x → prepend .venv_t5/ to sys.path, purge modules.
+      • Need 4.x → remove .venv_t5/ from sys.path, purge modules.
+
+    For LoRA adapters with custom names, the base model is resolved from
+    ``adapter_config.json`` before checking.
+
+    NOTE: Training and inference use subprocess isolation instead of this
+    function. This is only used by the export path (routes/export.py).
+    """
+    # Resolve LoRA adapters to their base model for accurate detection
+    resolved = _resolve_base_model(model_name)
+    want_5 = needs_transformers_5(resolved)
+    target_version = TRANSFORMERS_5_VERSION if want_5 else TRANSFORMERS_DEFAULT_VERSION
+    target_major = int(target_version.split(".")[0])
+
+    # Check what's actually loaded in memory
+    in_memory = _get_in_memory_version()
+
+    logger.info(
+        "Version check for '%s' (resolved: '%s'): need=%s, in_memory=%s",
+        model_name, resolved, target_version, in_memory,
+    )
+
+    # --- Already correct? ---------------------------------------------------
+    if in_memory is not None:
+        in_memory_major = int(in_memory.split(".")[0])
+        if in_memory_major == target_major:
+            logger.info(
+                "transformers %s already loaded — correct for '%s'",
+                in_memory, model_name,
+            )
+            return
+
+    # --- Switch version -----------------------------------------------------
+    if want_5:
+        logger.info("Activating transformers %s via .venv_t5…", TRANSFORMERS_5_VERSION)
+        _activate_5x()
+    else:
+        logger.info("Reverting to default transformers %s…", TRANSFORMERS_DEFAULT_VERSION)
+        _deactivate_5x()
+
+    final = _get_in_memory_version()
+    logger.info("✓ transformers version is now %s", final)

studio/frontend/package.json

@@ -31,10 +31,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",
@@ -65,7 +65,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/components/ui/chart.tsx

@@ -46,22 +46,94 @@ function ChartContainer({
 }) {
   const uniqueId = React.useId();
   const chartId = `chart-${id || uniqueId.replace(/:/g, "")}`;
+  const containerRef = React.useRef<HTMLDivElement | null>(null);
+  const [containerSize, setContainerSize] = React.useState<{
+    width: number;
+    height: number;
+  } | null>(null);
+
+  React.useEffect(() => {
+    const element = containerRef.current;
+    if (!element) return;
+
+    const updateSizeState = () => {
+      const { width, height } = element.getBoundingClientRect();
+      const nextSize =
+        width > 0 && height > 0
+          ? {
+              width: Math.round(width),
+              height: Math.round(height),
+            }
+          : null;
+
+      setContainerSize((currentSize) => {
+        if (!nextSize) {
+          // Keep the last valid size once mounted to avoid unmount/remount thrash.
+          return currentSize;
+        }
+        if (
+          currentSize &&
+          currentSize.width === nextSize.width &&
+          currentSize.height === nextSize.height
+        ) {
+          return currentSize;
+        }
+        return nextSize;
+      });
+    };
+
+    updateSizeState();
+
+    if (typeof ResizeObserver === "undefined") {
+      const recheckSize = () => {
+        if (document.visibilityState === "visible") {
+          updateSizeState();
+        }
+      };
+
+      window.addEventListener("resize", recheckSize);
+      window.addEventListener("orientationchange", recheckSize);
+      document.addEventListener("visibilitychange", recheckSize);
+
+      return () => {
+        window.removeEventListener("resize", recheckSize);
+        window.removeEventListener("orientationchange", recheckSize);
+        document.removeEventListener("visibilitychange", recheckSize);
+      };
+    }
+
+    const observer = new ResizeObserver(() => {
+      updateSizeState();
+    });
+    observer.observe(element);
+
+    return () => observer.disconnect();
+  }, []);
 
   return (
     <ChartContext.Provider value={{ config }}>
       <div
+        ref={containerRef}
         data-slot="chart"
         data-chart={chartId}
         className={cn(
-          "[&_.recharts-cartesian-axis-tick_text]:fill-muted-foreground [&_.recharts-cartesian-grid_line[stroke='#ccc']]:stroke-border/50 [&_.recharts-curve.recharts-tooltip-cursor]:stroke-border [&_.recharts-polar-grid_[stroke='#ccc']]:stroke-border [&_.recharts-radial-bar-background-sector]:fill-muted [&_.recharts-rectangle.recharts-tooltip-cursor]:fill-muted [&_.recharts-reference-line_[stroke='#ccc']]:stroke-border flex aspect-video justify-center text-xs [&_.recharts-dot[stroke='#fff']]:stroke-transparent [&_.recharts-layer]:outline-hidden [&_.recharts-sector]:outline-hidden [&_.recharts-sector[stroke='#fff']]:stroke-transparent [&_.recharts-surface]:outline-hidden",
+          "[&_.recharts-cartesian-axis-tick_text]:fill-muted-foreground [&_.recharts-cartesian-grid_line[stroke='#ccc']]:stroke-border/50 [&_.recharts-curve.recharts-tooltip-cursor]:stroke-border [&_.recharts-polar-grid_[stroke='#ccc']]:stroke-border [&_.recharts-radial-bar-background-sector]:fill-muted [&_.recharts-rectangle.recharts-tooltip-cursor]:fill-muted [&_.recharts-reference-line_[stroke='#ccc']]:stroke-border flex min-w-0 aspect-video justify-center text-xs [&_.recharts-dot[stroke='#fff']]:stroke-transparent [&_.recharts-layer]:outline-hidden [&_.recharts-sector]:outline-hidden [&_.recharts-sector[stroke='#fff']]:stroke-transparent [&_.recharts-surface]:outline-hidden",
           className,
         )}
         {...props}
       >
         <ChartStyle id={chartId} config={config} />
-        <RechartsPrimitive.ResponsiveContainer>
-          {children}
-        </RechartsPrimitive.ResponsiveContainer>
+        {containerSize ? (
+          <RechartsPrimitive.ResponsiveContainer
+            width="100%"
+            height="100%"
+            minWidth={0}
+            minHeight={1}
+            initialDimension={containerSize}
+          >
+            {children}
+          </RechartsPrimitive.ResponsiveContainer>
+        ) : null}
       </div>
     </ChartContext.Provider>
   );
@@ -100,30 +172,30 @@ ${colorConfig
   );
 };
 
-const ChartTooltip = RechartsPrimitive.Tooltip;
-
-function ChartTooltipContent({
-  active,
-  payload,
-  className,
+const ChartTooltip = RechartsPrimitive.Tooltip;
+
+function ChartTooltipContent({
+  active,
+  payload,
+  className,
   indicator = "dot",
   hideLabel = false,
   hideIndicator = false,
   label,
   labelFormatter,
   labelClassName,
-  formatter,
-  color,
-  nameKey,
-  labelKey,
-}: Partial<RechartsPrimitive.TooltipContentProps<any, any>> &
-  React.ComponentProps<"div"> & {
-    hideLabel?: boolean;
-    hideIndicator?: boolean;
-    indicator?: "line" | "dot" | "dashed";
-    nameKey?: string;
-    labelKey?: string;
-  }) {
+  formatter,
+  color,
+  nameKey,
+  labelKey,
+}: Partial<RechartsPrimitive.TooltipContentProps<any, any>> &
+  React.ComponentProps<"div"> & {
+    hideLabel?: boolean;
+    hideIndicator?: boolean;
+    indicator?: "line" | "dot" | "dashed";
+    nameKey?: string;
+    labelKey?: string;
+  }) {
   const { config } = useChart();
 
   const tooltipLabel = React.useMemo(() => {
@@ -248,20 +320,20 @@ function ChartTooltipContent({
   );
 }
 
-const ChartLegend = RechartsPrimitive.Legend;
-
-function ChartLegendContent({
-  className,
-  hideIcon = false,
-  payload,
-  verticalAlign = "bottom",
-  nameKey,
-}: React.ComponentProps<"div"> &
-  Pick<RechartsPrimitive.DefaultLegendContentProps, "payload" | "verticalAlign"> & {
-    hideIcon?: boolean;
-    nameKey?: string;
-  }) {
-  const { config } = useChart();
+const ChartLegend = RechartsPrimitive.Legend;
+
+function ChartLegendContent({
+  className,
+  hideIcon = false,
+  payload,
+  verticalAlign = "bottom",
+  nameKey,
+}: React.ComponentProps<"div"> &
+  Pick<RechartsPrimitive.DefaultLegendContentProps, "payload" | "verticalAlign"> & {
+    hideIcon?: boolean;
+    nameKey?: string;
+  }) {
+  const { config } = useChart();
 
   if (!payload?.length) {
     return null;

studio/frontend/src/components/ui/slider.tsx

@@ -1,57 +1,94 @@
-import { Slider as SliderPrimitive } from "radix-ui";
-import * as React from "react";
-
-import { cn } from "@/lib/utils";
-
-function Slider({
-  className,
-  defaultValue,
-  value,
-  min = 0,
-  max = 100,
-  ...props
-}: React.ComponentProps<typeof SliderPrimitive.Root>) {
-  const _values = React.useMemo(
-    () =>
-      Array.isArray(value)
-        ? value
-        : Array.isArray(defaultValue)
-          ? defaultValue
-          : [min, max],
-    [value, defaultValue, min, max],
-  );
-
-  return (
-    <SliderPrimitive.Root
-      data-slot="slider"
-      defaultValue={defaultValue}
-      value={value}
-      min={min}
-      max={max}
-      className={cn(
-        "data-vertical:min-h-40 relative flex w-full touch-none items-center select-none data-disabled:opacity-50 data-vertical:h-full data-vertical:w-auto data-vertical:flex-col",
-        className,
-      )}
-      {...props}
-    >
-      <SliderPrimitive.Track
-        data-slot="slider-track"
-        className="bg-muted rounded-4xl data-horizontal:h-3 data-horizontal:w-full data-vertical:h-full data-vertical:w-3 bg-muted relative grow overflow-hidden data-horizontal:w-full data-vertical:h-full cursor-pointer"
-      >
-        <SliderPrimitive.Range
-          data-slot="slider-range"
-          className="bg-primary absolute select-none data-horizontal:h-full data-vertical:w-full"
-        />
-      </SliderPrimitive.Track>
-      {Array.from({ length: _values.length }, (_, index) => (
-        <SliderPrimitive.Thumb
-          data-slot="slider-thumb"
-          key={index}
-          className="border-primary ring-ring/50 size-4 rounded-4xl border bg-white shadow-sm block shrink-0 select-none cursor-pointer disabled:pointer-events-none disabled:opacity-50 transition-transform duration-100 ease-out hover:scale-110 hover:ring-4 active:scale-95 focus-visible:ring-4 focus-visible:outline-hidden"
-        />
-      ))}
-    </SliderPrimitive.Root>
-  );
-}
-
-export { Slider };
+import { Slider as SliderPrimitive } from "radix-ui";
+import * as React from "react";
+
+import { cn } from "@/lib/utils";
+
+function Slider({
+  className,
+  defaultValue,
+  value,
+  min = 0,
+  max = 100,
+  orientation = "horizontal",
+  onValueChange,
+  ...props
+}: React.ComponentProps<typeof SliderPrimitive.Root>) {
+  const isControlled = Array.isArray(value);
+  const [uncontrolledValues, setUncontrolledValues] =
+    React.useState<number[]>(() =>
+      Array.isArray(defaultValue) ? defaultValue : [min, max],
+    );
+
+  const values = isControlled ? value : uncontrolledValues;
+  const handleValueChange = React.useCallback(
+    (nextValues: number[]) => {
+      if (!isControlled) {
+        setUncontrolledValues(nextValues);
+      }
+      onValueChange?.(nextValues);
+    },
+    [isControlled, onValueChange],
+  );
+
+  // For single-thumb horizontal sliders, render the fill bar as a sibling of
+  // the track (outside its overflow-hidden container) so it can align flush
+  // with the thumb center without being clipped. The Range inside the track
+  // is hidden in this case to avoid double-painting.
+  const isSingleThumbHorizontal =
+    values.length === 1 && orientation === "horizontal";
+  const fillPercent = isSingleThumbHorizontal
+    ? Math.min(
+        100,
+        Math.max(
+          0,
+          max === min ? 0 : (((values[0] ?? min) - min) / (max - min)) * 100,
+        ),
+      )
+    : null;
+
+  return (
+    <SliderPrimitive.Root
+      data-slot="slider"
+      defaultValue={defaultValue}
+      value={value}
+      min={min}
+      max={max}
+      orientation={orientation}
+      onValueChange={handleValueChange}
+      className={cn(
+        "data-vertical:min-h-40 relative flex w-full touch-none items-center select-none data-disabled:opacity-50 data-vertical:h-full data-vertical:w-auto data-vertical:flex-col",
+        className,
+      )}
+      {...props}
+    >
+      <SliderPrimitive.Track
+        data-slot="slider-track"
+        className="bg-muted rounded-4xl data-horizontal:h-3 data-horizontal:w-full data-vertical:h-full data-vertical:w-3 bg-muted relative grow overflow-hidden data-horizontal:w-full data-vertical:h-full cursor-pointer"
+      >
+        <SliderPrimitive.Range
+          data-slot="slider-range"
+          className={cn(
+            "bg-primary absolute select-none data-horizontal:h-full data-vertical:w-full",
+            isSingleThumbHorizontal && "opacity-0",
+          )}
+        />
+      </SliderPrimitive.Track>
+      {isSingleThumbHorizontal && (
+        <div
+          aria-hidden={true}
+          className="absolute inset-y-0 left-0 my-auto h-3 rounded-4xl bg-primary pointer-events-none"
+          style={{ width: `${fillPercent}%` }}
+        />
+      )}
+      {Array.from({ length: values.length }, (_, index) => (
+        <SliderPrimitive.Thumb
+          data-slot="slider-thumb"
+          key={index}
+          className="border-primary ring-ring/50 size-4 rounded-4xl border bg-white shadow-sm block shrink-0 select-none cursor-pointer disabled:pointer-events-none disabled:opacity-50 transition-transform duration-100 ease-out hover:scale-110 hover:ring-4 active:scale-95 focus-visible:ring-4 focus-visible:outline-hidden"
+        />
+      ))}
+    </SliderPrimitive.Root>
+  );
+}
+
+export { Slider };

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,
@@ -99,7 +100,7 @@ export async function* streamChatCompletions(
   payload: OpenAIChatCompletionsRequest,
   signal: AbortSignal,
 ): AsyncGenerator<OpenAIChatChunk> {
-  const response = await authFetch("/api/inference/chat/completions", {
+  const response = await authFetch("/v1/chat/completions", {
     method: "POST",
     headers: { "Content-Type": "application/json" },
     body: JSON.stringify(payload),
@@ -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/runtime-provider.tsx

@@ -202,7 +202,7 @@ async function generateTitleWithModel(payload: {
     return joined.length > 60 ? joined.slice(0, 60).trimEnd() : joined;
   }
 
-  const response = await authFetch("/api/inference/chat/completions", {
+  const response = await authFetch("/v1/chat/completions", {
     method: "POST",
     headers: { "Content-Type": "application/json" },
     body: JSON.stringify({

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/export/components/export-dialog.tsx

@@ -181,6 +181,8 @@ export function ExportDialog({
                         </InputGroupAddon>
                         <InputGroupInput
                           type="password"
+                          autoComplete="new-password"
+                          name="hf-token"
                           placeholder="hf_..."
                           value={hfToken}
                           onChange={(e) => onHfTokenChange(e.target.value)}

studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx

@@ -186,6 +186,8 @@ export function DatasetStep() {
               </InputGroupAddon>
               <InputGroupInput
                 type="password"
+                autoComplete="new-password"
+                name="hf-token"
                 placeholder="hf_..."
                 value={hfToken}
                 onChange={(e) => setHfToken(e.target.value)}

studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx

@@ -152,6 +152,8 @@ export function ModelSelectionStep() {
           </InputGroupAddon>
           <InputGroupInput
             type="password"
+            autoComplete="new-password"
+            name="hf-token"
             placeholder="hf_..."
             value={hfToken}
             onChange={(e) => setHfToken(e.target.value)}

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

@@ -64,15 +64,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 = datasetSource === "huggingface";
 
   // When format changes, remap existing mapping roles to the new format's role names
@@ -152,10 +153,10 @@ 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]);
+  }, [open, datasetName, data, effectiveIsVlm, datasetFormat, effectiveIsAudio, manualMapping, setManualMapping]);
 
   const rows = data?.preview_samples ?? [];
   const columns = data?.columns ?? [];
@@ -355,6 +356,7 @@ export function DatasetPreviewDialog({
                   mappingOk={mappingOk}
                   autoDetected={hasHeuristicMapping}
                   isVlm={effectiveIsVlm}
+                  isAudio={effectiveIsAudio}
                   format={datasetFormat}
                 />
               )}

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

@@ -93,6 +93,8 @@ export function DatasetSection() {
     setUploadedFile,
     hfToken,
     modelType,
+    isVisionModel,
+    isCheckingVision,
     datasetSliceStart,
     setDatasetSliceStart,
     datasetSliceEnd,
@@ -115,6 +117,8 @@ export function DatasetSection() {
       setUploadedFile: s.setUploadedFile,
       hfToken: s.hfToken,
       modelType: s.modelType,
+      isVisionModel: s.isVisionModel,
+      isCheckingVision: s.isCheckingVision,
       datasetSliceStart: s.datasetSliceStart,
       setDatasetSliceStart: s.setDatasetSliceStart,
       datasetSliceEnd: s.datasetSliceEnd,
@@ -185,6 +189,9 @@ export function DatasetSection() {
     }
     setSearchQuery(val);
   }
+
+  const effectiveModelType = !isCheckingVision && isVisionModel ? "vision" : modelType;
+
   const {
     results: hfResults,
     isLoading,
@@ -192,7 +199,7 @@ export function DatasetSection() {
     fetchMore,
     error: hfSearchError,
   } = useHfDatasetSearch(pickerTab === "huggingface" ? debouncedQuery : "", {
-    modelType,
+    modelType: effectiveModelType,
     accessToken: hfToken || undefined,
     enabled: pickerTab === "huggingface",
   });

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

@@ -581,6 +581,8 @@ export function ModelSection() {
             </InputGroupAddon>
             <InputGroupInput
               type="password"
+              autoComplete="new-password"
+              name="hf-token"
               placeholder="hf_..."
               value={hfToken}
               onChange={(e) => setHfToken(e.target.value)}

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];
@@ -105,7 +105,7 @@ export function TrainingSection() {
         <div className="relative  ">
           <ChartContainer
             config={chartConfig}
-            className="min-h-[180px] w-full relative right-8 w-full blur "
+            className="h-[180px] w-full relative right-8 blur"
           >
             <LineChart data={placeholderData} accessibilityLayer={true}>
               <CartesianGrid vertical={false} strokeDasharray="3 3" />

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

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

@@ -72,7 +72,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/components/hf-dataset-subset-split-selectors.tsx

@@ -48,33 +48,34 @@ export function HfDatasetSubsetSplitSelectors({
   const {
     subsets: hfSubsets,
     splits: hfSplits,
-    hasMultipleSubsets,
     isLoading,
     error,
   } = useHfDatasetSplits(enabled ? datasetName : null, datasetSubset, {
     accessToken,
   });
 
+  // Auto-select subset and split in one pass to avoid racing effects
   useEffect(() => {
-    if (hfSubsets.length === 1 && datasetSubset !== hfSubsets[0]) {
-      setDatasetSubset(hfSubsets[0]);
-    }
-  }, [hfSubsets, datasetSubset, setDatasetSubset]);
+    if (hfSubsets.length === 0) return;
 
-  useEffect(() => {
+    // --- subset ---
+    if (!datasetSubset || !hfSubsets.includes(datasetSubset)) {
+      const pick = hfSubsets.includes("default") ? "default" : hfSubsets[0];
+      setDatasetSubset(pick);
+      return;
+    }
+
+    // --- split (only once subset is settled) ---
     if (hfSplits.length === 0) return;
-    if (hasMultipleSubsets && !datasetSubset) return;
-    if (hfSplits.length === 1 && datasetSplit !== hfSplits[0]) {
-      setDatasetSplit(hfSplits[0]);
-    } else if (!datasetSplit && hfSplits.includes("train")) {
-      setDatasetSplit("train");
-    } else if (!datasetSplit) {
-      setDatasetSplit(hfSplits[0]);
+    if (!datasetSplit || !hfSplits.includes(datasetSplit)) {
+      const pick = hfSplits.includes("train") ? "train" : hfSplits[0];
+      setDatasetSplit(pick);
     }
   }, [
+    hfSubsets,
     hfSplits,
-    hasMultipleSubsets,
     datasetSubset,
+    setDatasetSubset,
     datasetSplit,
     setDatasetSplit,
   ]);

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,8 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             datasetManualMapping: emptyManualMapping(),
             datasetSliceStart: null,
             datasetSliceEnd: null,
-            isDatasetMultimodal: null,
+            isDatasetImage: null,
+            isDatasetAudio: false,
             isCheckingDataset: false,
           });
         },
@@ -272,7 +279,8 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             datasetSplit: null,
             datasetEvalSplit: null,
             datasetManualMapping: emptyManualMapping(),
-            isDatasetMultimodal: null,
+            isDatasetImage: null,
+            isDatasetAudio: false,
             isCheckingDataset: false,
           });
         },
@@ -280,7 +288,8 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
           set({
             datasetSplit,
             datasetManualMapping: emptyManualMapping(),
-            isDatasetMultimodal: null,
+            isDatasetImage: null,
+            isDatasetAudio: false,
             isCheckingDataset: false,
           });
 
@@ -296,7 +305,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"
@@ -329,7 +338,8 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
             datasetManualMapping: emptyManualMapping(),
             datasetSliceStart: null,
             datasetSliceEnd: null,
-            isDatasetMultimodal: null,
+            isDatasetImage: null,
+            isDatasetAudio: false,
             isCheckingDataset: false,
           });
         },

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/hooks/use-hf-dataset-splits.ts

@@ -86,6 +86,14 @@ export function useHfDatasetSplits(
   const [isLoading, setIsLoading] = useState(false);
   const [error, setError] = useState<string | null>(null);
 
+  
+  const [prevDatasetName, setPrevDatasetName] = useState(datasetName);
+  if (datasetName !== prevDatasetName) {
+    setPrevDatasetName(datasetName);
+    setEntries([]);
+    setError(null);
+  }
+
   const accessToken = options?.accessToken;
 
   const fetchSplits = useCallback(

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

studio/frontend/vite.config.ts

@@ -17,6 +17,10 @@ export default defineConfig({
         target: "http://127.0.0.1:8000",
         changeOrigin: true,
       },
+      "/v1": {
+        target: "http://127.0.0.1:8000",
+        changeOrigin: true,
+      },
       "/seed/inspect": {
         target: "http://127.0.0.1:8004",
         changeOrigin: true,