diff --git a/.gitignore b/.gitignore index f7f914fb96..d89fdb9693 100755 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,8 @@ __pycache__/ # Virtual environments .venv/ +.venv_overlay/ +.venv_t5/ venv/ env/ environment.yaml diff --git a/install_python_stack.py b/install_python_stack.py index 706d8ac4fd..c6b4e556f4 100644 --- a/install_python_stack.py +++ b/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 diff --git a/setup.ps1 b/setup.ps1 index 8a2ec56033..e46a7522d2 100644 --- a/setup.ps1 +++ b/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) # ========================================================================== diff --git a/setup.sh b/setup.sh index 899677444b..6edb470d8d 100755 --- a/setup.sh +++ b/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. diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml index ddd58a1225..dd5ae51ab0 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B-it.yaml +++ b/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 diff --git a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml b/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml index 1ca686aea7..e53e163a04 100644 --- a/studio/backend/assets/configs/model_defaults/gemma/unsloth_gemma-3n-E4B.yaml +++ b/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 diff --git a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml index 9c65107699..72b5b018e1 100644 --- a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml +++ b/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 diff --git a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml index 84cf750262..d20751b0c7 100644 --- a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml +++ b/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 diff --git a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml index 294da47e10..f5f49fe1e6 100644 --- a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml +++ b/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 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml index 1bbbcdf66c..5a3c4abb48 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml +++ b/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 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml index d41c1c65fb..1906ecda51 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml @@ -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 diff --git a/studio/backend/core/__init__.py b/studio/backend/core/__init__.py index a4a700d45b..2b1e3246c0 100644 --- a/studio/backend/core/__init__.py +++ b/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}") diff --git a/studio/backend/core/export/__init__.py b/studio/backend/core/export/__init__.py index 66154f48eb..0a883f5f3f 100644 --- a/studio/backend/core/export/__init__.py +++ b/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', ] diff --git a/studio/backend/core/export/export.py b/studio/backend/core/export/export.py index 3d3db2d560..9e1c5cff84 100644 --- a/studio/backend/core/export/export.py +++ b/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: diff --git a/studio/backend/core/export/orchestrator.py b/studio/backend/core/export/orchestrator.py new file mode 100644 index 0000000000..ec18d3533e --- /dev/null +++ b/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 diff --git a/studio/backend/core/export/worker.py b/studio/backend/core/export/worker.py new file mode 100644 index 0000000000..9e7e72e9dd --- /dev/null +++ b/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(), + }) diff --git a/studio/backend/core/inference/__init__.py b/studio/backend/core/inference/__init__.py index ff8b75d36a..6d44742d9c 100644 --- a/studio/backend/core/inference/__init__.py +++ b/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', ] diff --git a/studio/backend/core/inference/audio_codecs.py b/studio/backend/core/inference/audio_codecs.py new file mode 100644 index 0000000000..5a1fd5d984 --- /dev/null +++ b/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") diff --git a/studio/backend/core/inference/inference.py b/studio/backend/core/inference/inference.py index 329f5d944b..780a399637 100644 --- a/studio/backend/core/inference/inference.py +++ b/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 diff --git a/studio/backend/core/inference/orchestrator.py b/studio/backend/core/inference/orchestrator.py new file mode 100644 index 0000000000..b08ec5bb70 --- /dev/null +++ b/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 diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py new file mode 100644 index 0000000000..3e766b3b01 --- /dev/null +++ b/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(), + }) diff --git a/studio/backend/core/training/__init__.py b/studio/backend/core/training/__init__.py index 8fc2a6c721..0e0299caf6 100644 --- a/studio/backend/core/training/__init__.py +++ b/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', diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 0ff317ecf1..b51fe60491 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -3,6 +3,7 @@ Unsloth Training Backend Integrates Unsloth training capabilities with the FastAPI backend """ import os +import sys # Prevent tokenizer parallelism deadlocks when datasets uses multiprocessing fork os.environ["TOKENIZERS_PARALLELISM"] = "false" @@ -23,19 +24,15 @@ from dataclasses import dataclass import pandas as pd from datasets import Dataset, load_dataset -# Add the parent directory to sys.path to import unsloth modules -#sys.path.append(os.path.join(os.path.dirname(__file__), '..')) -from utils.models import is_vision_model +from utils.models import is_vision_model, detect_audio_type from utils.datasets import format_and_template_dataset from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER from trl import SFTTrainer, SFTConfig -# Import Unsloth trainers -#from unsloth_compiled_cache.UnslothSFTTrainer import _UnslothSFTTrainer as SFTTrainer - logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) + @dataclass class TrainingProgress: """Training progress tracking""" @@ -69,9 +66,15 @@ class UnslothTrainer: self.is_training = False self.should_stop = False self.save_on_stop = True + self.load_in_4bit = True # Track quantization mode for metadata # Model state tracking self.is_vlm = False + self.is_audio = False + self.is_audio_vlm = False # Multimodal model (e.g. Gemma 3N) trained on audio data + self._audio_type = None # 'csm', 'whisper', 'snac', 'bicodec', 'dac' + self._cuda_audio_used = False # Set once after audio CUDA preprocessing; never cleared + self._spark_tts_repo_dir = None # Path to downloaded Spark-TTS repo (for BiCodecTokenizer) self.model_name = None # Training metrics tracking @@ -108,13 +111,224 @@ class UnslothTrainer: except Exception as e: logger.error(f"Error in progress callback: {e}") + def _create_progress_callback(self): + """Create a TrainerCallback for progress tracking. Reused by all training branches.""" + from transformers import TrainerCallback + trainer_ref = self + + class _ProgressCallback(TrainerCallback): + def on_log(self, args, state, control, logs=None, **kwargs): + if not logs: + return + loss_value = logs.get('loss', logs.get('train_loss', 0.0)) + current_step = state.global_step + grad_norm = logs.get('grad_norm', None) + + elapsed_seconds = None + if trainer_ref.training_start_time is not None: + elapsed_seconds = time.time() - trainer_ref.training_start_time + + eta_seconds = None + if elapsed_seconds is not None and current_step > 0: + total_steps = trainer_ref.training_progress.total_steps + if total_steps > 0: + steps_remaining = total_steps - current_step + if steps_remaining > 0: + eta_seconds = (elapsed_seconds / current_step) * steps_remaining + + num_tokens = getattr(state, "num_input_tokens_seen", None) + + trainer_ref._update_progress( + step=current_step, + epoch=round(state.epoch, 2) if state.epoch else 0, + loss=loss_value, + learning_rate=logs.get('learning_rate', 0.0), + elapsed_seconds=elapsed_seconds, + eta_seconds=eta_seconds, + grad_norm=grad_norm, + num_tokens=num_tokens, + eval_loss=logs.get('eval_loss', None), + status_message="", + ) + + def on_epoch_end(self, args, state, control, **kwargs): + trainer_ref._update_progress(epoch=state.epoch, step=state.global_step) + + def on_step_end(self, args, state, control, **kwargs): + if trainer_ref.should_stop: + print(f"Stop detected at step {state.global_step}\n") + control.should_training_stop = True + return control + + return _ProgressCallback() + + def _calculate_total_steps(self, num_samples, batch_size, grad_accum, num_epochs, max_steps): + """Calculate total training steps from dataset size and training params.""" + if max_steps and max_steps > 0: + return max_steps + len_dataloader = math.ceil(num_samples / batch_size) + steps_per_epoch = max(len_dataloader // grad_accum + int(len_dataloader % grad_accum > 0), 1) + return steps_per_epoch * num_epochs + + def _build_audio_training_args(self, training_args, output_dir, *, extra_args=None): + """Build training args dict for audio branches. + + Constructs the common config (batch size, lr, warmup, fp16/bf16, etc.) + and applies per-branch overrides via extra_args. + """ + batch_size = training_args.get('batch_size', 2) + gradient_accumulation_steps = training_args.get('gradient_accumulation_steps', 4) + warmup_steps_val = training_args.get('warmup_steps', 5) + max_steps_val = training_args.get('max_steps', 0) + learning_rate = training_args.get('learning_rate', 2e-4) + weight_decay = training_args.get('weight_decay', 0.001) + lr_scheduler_type = training_args.get('lr_scheduler_type', 'linear') + random_seed = training_args.get('random_seed', 3407) + optim_value = training_args.get('optim', 'adamw_8bit') + + config = { + "per_device_train_batch_size": batch_size, + "gradient_accumulation_steps": gradient_accumulation_steps, + "warmup_steps": warmup_steps_val if warmup_steps_val is not None else 5, + "learning_rate": learning_rate, + "fp16": not is_bfloat16_supported(), + "bf16": is_bfloat16_supported(), + "logging_steps": 1, + "optim": optim_value, + "weight_decay": weight_decay, + "lr_scheduler_type": lr_scheduler_type, + "seed": random_seed, + "output_dir": output_dir, + "report_to": ["wandb"] if training_args.get('enable_wandb', False) else "none", + } + + # max_steps vs epochs + if max_steps_val and max_steps_val > 0: + config["max_steps"] = max_steps_val + else: + config["num_train_epochs"] = training_args.get('num_epochs', 3) + + # save_steps + save_steps_val = training_args.get('save_steps', 0) + if save_steps_val and save_steps_val > 0: + config["save_steps"] = save_steps_val + config["save_strategy"] = "steps" + + # Apply per-branch overrides + if extra_args: + config.update(extra_args) + + return config + + def _finalize_training(self, output_dir, label=""): + """Save model after training and update progress. Used by all training branches.""" + if self.should_stop and self.save_on_stop: + self.trainer.save_model() + self.tokenizer.save_pretrained(output_dir) + self._patch_adapter_config(output_dir) + msg = f"{label} training stopped" if label else "Training stopped" + print(f"\n{msg}. Model saved to {output_dir}\n") + self._update_progress( + is_training=False, + status_message=f"Training stopped. Model saved to {output_dir}", + ) + elif self.should_stop: + msg = f"{label} training cancelled" if label else "Training cancelled" + print(f"\n{msg}.\n") + self._update_progress(is_training=False, status_message="Training cancelled.") + else: + self.trainer.save_model() + self.tokenizer.save_pretrained(output_dir) + self._patch_adapter_config(output_dir) + msg = f"{label} training completed" if label else "Training completed" + print(f"\n{msg}! Model saved to {output_dir}\n") + self._update_progress( + is_training=False, + is_completed=True, + status_message=f"Training completed! Model saved to {output_dir}", + ) + + def _cleanup_audio_artifacts(self): + """Remove sys.path entries and sys.modules from previous audio preprocessing. + + After audio training, cloned repo dirs (OuteTTS, Spark-TTS) remain on + sys.path and heavy audio modules (snac, whisper, sparktts, outetts) stay + in sys.modules. When the next training run calls dataset.map(num_proc=N), + forked child processes inherit this stale state and deadlock. + """ + import sys as _sys + + # Remove cloned audio repo paths from sys.path + base_dir = os.path.dirname(os.path.abspath(__file__)) + audio_paths = [ + os.path.join(base_dir, "inference", "OuteTTS"), # DAC/OuteTTS + ] + # Spark-TTS path is relative to the downloaded repo + if self._spark_tts_repo_dir: + spark_code_dir = os.path.join(os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS") + audio_paths.append(spark_code_dir) + + removed_paths = [] + for path in audio_paths: + if path in _sys.path: + _sys.path.remove(path) + removed_paths.append(path) + + # Remove stale audio modules from sys.modules + prefixes = ('snac', 'whisper', 'sparktts', 'outetts') + removed_modules = [key for key in _sys.modules if key.startswith(prefixes)] + for key in removed_modules: + del _sys.modules[key] + + if removed_paths or removed_modules: + print(f"Cleaned up audio artifacts: {len(removed_paths)} paths, " + f"{len(removed_modules)} modules\n") + + def _resolve_audio_columns(self, dataset, custom_format_mapping: dict = None): + """Resolve audio, text, and speaker columns from user mapping or hardcoded fallback. + + Returns: + dict with keys: audio_col, text_col, speaker_col (speaker_col may be None) + """ + cols = dataset.column_names + + if custom_format_mapping: + audio_col = None + text_col = None + speaker_col = None + for col, role in custom_format_mapping.items(): + if role == "audio": + audio_col = col + elif role == "text": + text_col = col + elif role == "speaker_id": + speaker_col = col + # Use mapping if both required columns exist in the dataset + if audio_col and audio_col in cols and text_col and text_col in cols: + return {"audio_col": audio_col, "text_col": text_col, "speaker_col": speaker_col} + + # Hardcoded fallback (existing behavior) + audio_col = next((c for c in cols if c.lower() in ("audio", "speech")), None) + text_col = next((c for c in cols if c.lower() in ("text", "sentence", "transcript", "transcription")), None) + + speaker_col = None + if "source" in cols: + speaker_col = "source" + elif "speaker_id" in cols: + speaker_col = "speaker_id" + + return {"audio_col": audio_col, "text_col": text_col, "speaker_col": speaker_col} + + def load_model(self, model_name: str, max_seq_length: int = 2048, load_in_4bit: bool = True, hf_token: Optional[str] = None, - is_dataset_multimodal: bool = False) -> bool: + is_dataset_image: bool = False, + is_dataset_audio: bool = False) -> bool: """Load model for training (supports both text and vision models)""" + self.load_in_4bit = load_in_4bit # Store for training_meta.json try: if self.model is not None: del self.model @@ -127,17 +341,48 @@ class UnslothTrainer: print("\nClearing GPU memory before training...") clear_gpu_cache() + # Clean up sys.path and sys.modules from previous audio preprocessing + # to prevent deadlocks when forking worker processes in dataset.map() + self._cleanup_audio_artifacts() + + # Reload Unsloth-patched transformers modeling modules before clearing + # the compiled cache. unsloth_compile_transformers() sets __UNSLOTH_PATCHED__ + # on each modeling module and replaces methods with exec'd code. + # clear_unsloth_compiled_cache() deletes the disk cache, but the flag + # prevents re-compilation — leaving missing cache files. Reloading + # restores original class definitions so Unsloth can re-compile cleanly. + import sys as _sys + import importlib + for _key, _mod in list(_sys.modules.items()): + if 'transformers.models.' in _key and '.modeling_' in _key: + if hasattr(_mod, '__UNSLOTH_PATCHED__'): + try: + importlib.reload(_mod) + except Exception: + pass # Non-critical — Unsloth will handle stale modules + # Remove stale compiled cache so the new model gets a fresh one from utils.cache_cleanup import clear_unsloth_compiled_cache clear_unsloth_compiled_cache() + # Detect audio model type dynamically (config.json + tokenizer) + self._audio_type = detect_audio_type(model_name, hf_token) + # audio_vlm is detected as an audio_type now, handle it separately + if self._audio_type == 'audio_vlm': + self.is_audio = False + self.is_audio_vlm = is_dataset_audio # Only use audio VLM path if dataset has audio + self._audio_type = None + else: + self.is_audio = self._audio_type is not None + self.is_audio_vlm = False - # Detect if this is a vision model AND dataset is multimodal - # A vision-capable model with a text-only dataset should use FastLanguageModel - self.is_vlm = is_vision_model(model_name) and is_dataset_multimodal + # VLM: vision model with image dataset (mutually exclusive with audio paths) + vision = is_vision_model(model_name) if not self.is_audio else False + self.is_vlm = not self.is_audio_vlm and vision and is_dataset_image self.model_name = model_name + self.max_seq_length = max_seq_length - logger.info(f"Model architecture is vision: {is_vision_model(model_name)}") - logger.info(f"Dataset is multimodal: {is_dataset_multimodal}") + logger.info(f"Audio type: {self._audio_type}, is_audio: {self.is_audio}, is_audio_vlm: {self.is_audio_vlm}") + logger.info(f"Dataset has images: {is_dataset_image}, audio: {is_dataset_audio}") logger.info(f"Using VLM path: {self.is_vlm}") # Reset training state for new run @@ -152,19 +397,148 @@ class UnslothTrainer: # Update UI immediately with loading message model_display = model_name.split('/')[-1] if '/' in model_name else model_name + model_type_label = 'audio' if self.is_audio else ('vision' if self.is_vlm else 'text') self._update_progress( - status_message=f"Loading {'vision' if self.is_vlm else 'text'} model... {model_display}" + status_message=f"Loading {model_type_label} model... {model_display}" ) - print(f"\nLoading {'vision' if self.is_vlm else 'text'} model: {model_name}") + print(f"\nLoading {model_type_label} model: {model_name}") # Set HF token if provided if hf_token: os.environ["HF_TOKEN"] = hf_token + # Proactive gated-model check: verify access BEFORE from_pretrained. + # Catches ALL gated/private models (text, vision, audio) globally. + if '/' in model_name: # Only check HF repo IDs, not local paths + try: + from huggingface_hub import model_info as hf_model_info + info = hf_model_info(model_name, token=hf_token or None) + # model_info succeeds even for gated repos (metadata is public), + # but info.gated tells us if files require acceptance/token. + if info.gated and not hf_token: + friendly = ( + f"Access denied for '{model_name}'. This model is gated. " + f"Please add a Hugging Face token with access and try again." + ) + logger.error(f"Model '{model_name}' is gated (gated={info.gated}) and no HF token provided") + self._update_progress(error=friendly, is_training=False) + return False + except Exception as gate_err: + from huggingface_hub.utils import GatedRepoError, RepositoryNotFoundError + if isinstance(gate_err, (GatedRepoError, RepositoryNotFoundError)): + friendly = ( + f"Access denied for '{model_name}'. This model is gated or private. " + f"Please add a Hugging Face token with access and try again." + ) + logger.error(f"Gated model check failed: {gate_err}") + self._update_progress(error=friendly, is_training=False) + return False # Branch based on model type - if self.is_vlm: + if self._audio_type == 'csm': + # CSM: FastModel + auto_model=CsmForConditionalGeneration + load_in_4bit=False + from unsloth import FastModel + from transformers import CsmForConditionalGeneration + self.model, self.tokenizer = FastModel.from_pretrained( + model_name=model_name, + max_seq_length=max_seq_length, + dtype=None, + auto_model=CsmForConditionalGeneration, + load_in_4bit=False, + token=hf_token, + ) + logger.info("Loaded CSM audio model") + + elif self._audio_type == 'whisper': + # Whisper: FastModel + auto_model=WhisperForConditionalGeneration + load_in_4bit=False + from unsloth import FastModel + from transformers import WhisperForConditionalGeneration + self.model, self.tokenizer = FastModel.from_pretrained( + model_name=model_name, + dtype=None, + load_in_4bit=False, + auto_model=WhisperForConditionalGeneration, + whisper_language="English", + whisper_task="transcribe", + token=hf_token, + ) + # Configure generation settings (notebook lines 100-105) + self.model.generation_config.language = "<|en|>" + self.model.generation_config.task = "transcribe" + self.model.config.suppress_tokens = [] + self.model.generation_config.forced_decoder_ids = None + logger.info("Loaded Whisper audio model (FastModel)") + + elif self._audio_type == 'snac': + # Orpheus: language model with audio codec tokens + self.model, self.tokenizer = FastLanguageModel.from_pretrained( + model_name=model_name, + max_seq_length=max_seq_length, + dtype=None, + load_in_4bit=load_in_4bit, + token=hf_token, + ) + logger.info(f"Loaded {self._audio_type} audio model (FastLanguageModel)") + + elif self._audio_type == 'bicodec': + # Spark-TTS: download full repo (contains sparktts package + BiCodec weights), + # then load only the LLM subfolder with FastModel. + # model_name may be: + # "Spark-TTS-0.5B/LLM" (local-style, from YAML mapping) + # "unsloth/Spark-TTS-0.5B" (HF repo ID) + from unsloth import FastModel + from huggingface_hub import snapshot_download + + if model_name.endswith("/LLM"): + # "Spark-TTS-0.5B/LLM" → parent="Spark-TTS-0.5B" + local_dir = model_name.rsplit("/", 1)[0] + hf_repo = f"unsloth/{local_dir}" + llm_path = model_name + else: + # "unsloth/Spark-TTS-0.5B" → local_dir="Spark-TTS-0.5B" + hf_repo = model_name + local_dir = model_name.split("/")[-1] + llm_path = f"{local_dir}/LLM" + + repo_path = snapshot_download(hf_repo, local_dir=local_dir) + self._spark_tts_repo_dir = os.path.abspath(repo_path) # Absolute path for sys.path + llm_path = os.path.join(self._spark_tts_repo_dir, "LLM") + + self.model, self.tokenizer = FastModel.from_pretrained( + model_name=llm_path, + max_seq_length=max_seq_length, + dtype=torch.float32, # Spark-TTS requires float32 + load_in_4bit=False, + token=hf_token, + ) + logger.info("Loaded Spark-TTS (bicodec) model") + + elif self._audio_type == 'dac': + # OuteTTS: uses FastModel (not FastLanguageModel) with load_in_4bit=False + from unsloth import FastModel + self.model, self.tokenizer = FastModel.from_pretrained( + model_name, + max_seq_length=max_seq_length, + load_in_4bit=False, + token=hf_token, + ) + logger.info("Loaded OuteTTS (dac) model (FastModel)") + + elif self.is_audio_vlm: + # Audio VLM: multimodal model trained on audio (e.g. Gemma 3N) + # Uses FastModel (general loader) — returns (model, processor) + from unsloth import FastModel + self.model, self.tokenizer = FastModel.from_pretrained( + model_name=model_name, + max_seq_length=max_seq_length, + dtype=None, + load_in_4bit=load_in_4bit, + token=hf_token, + ) + logger.info("Loaded audio VLM model (FastModel)") + + elif self.is_vlm: # Load vision model - returns (model, tokenizer) self.model, self.tokenizer = FastVisionModel.from_pretrained( model_name=model_name, @@ -201,10 +575,41 @@ class UnslothTrainer: print("Model loaded successfully") return True - except Exception as e: + except OSError as e: + if "could not get source code" in str(e) and not getattr(self, '_source_code_retried', False): + # Unsloth's patching can leave stale state that makes + # inspect.getsource() fail when switching model families + # (e.g. gemma3 → gemma3n). The load always succeeds on a + # second attempt because the failed first call's partial + # imports clean up the stale state as a side effect. + self._source_code_retried = True + print(f"\n'could not get source code' — retrying once...\n") + return self.load_model(model_name, max_seq_length, load_in_4bit, hf_token, + is_dataset_image, is_dataset_audio) + error_msg = str(e) + error_lower = error_msg.lower() + if any(k in error_lower for k in ("gated repo", "access to it at", "401", "403", "unauthorized", "forbidden")): + error_msg = ( + f"Access denied for '{model_name}'. This model is gated or private. " + f"Please add a Hugging Face token with access and try again." + ) logger.error(f"Error loading model: {e}") - self._update_progress(error=str(e), is_training=False) + self._update_progress(error=error_msg, is_training=False) return False + except Exception as e: + error_msg = str(e) + # Catch gated/auth errors and surface a friendly message + error_lower = error_msg.lower() + if any(k in error_lower for k in ("gated repo", "access to it at", "401", "403", "unauthorized", "forbidden")): + error_msg = ( + f"Access denied for '{model_name}'. This model is gated or private. " + f"Please add a Hugging Face token with access and try again." + ) + logger.error(f"Error loading model: {e}") + self._update_progress(error=error_msg, is_training=False) + return False + finally: + self._source_code_retried = False def prepare_model_for_training(self, use_lora: bool = True, @@ -282,8 +687,81 @@ class UnslothTrainer: print(f"Configuring LoRA adapters (r={lora_r}, alpha={lora_alpha})...\n") print(f"Gradient checkpointing: {use_gradient_checkpointing} (type: {type(use_gradient_checkpointing).__name__})\n") - # Branch based on vision vs text - if self.is_vlm: + # Branch based on model type: audio, audio_vlm, vision, or text + if self._audio_type in ('csm', 'bicodec', 'dac') or self.is_audio_vlm: + # Models using FastModel.get_peft_model (codec audio + audio VLM) + from unsloth import FastModel + label = self._audio_type or 'audio_vlm' + print(f"{label} LoRA configuration:") + print(f" - Target modules: {target_modules}") + if self.is_audio_vlm: + print(f" - Finetune vision layers: {finetune_vision_layers}") + print(f" - Finetune language layers: {finetune_language_layers}") + print(f" - Finetune attention modules: {finetune_attention_modules}") + print(f" - Finetune MLP modules: {finetune_mlp_modules}") + print() + + peft_kwargs = dict( + r=lora_r, + target_modules=target_modules, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + bias="none", + use_gradient_checkpointing=use_gradient_checkpointing, + random_state=3407, + use_rslora=use_rslora, + loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + ) + # Audio VLM models support VLM-style layer selection + if self.is_audio_vlm: + peft_kwargs.update( + finetune_vision_layers=finetune_vision_layers, + finetune_language_layers=finetune_language_layers, + finetune_attention_modules=finetune_attention_modules, + finetune_mlp_modules=finetune_mlp_modules, + ) + + self.model = FastModel.get_peft_model(self.model, **peft_kwargs) + + elif self._audio_type == 'whisper': + # Phase 2: Whisper uses FastModel.get_peft_model with task_type=None + from unsloth import FastModel + print(f"Audio model (whisper) LoRA configuration:") + print(f" - Target modules: {target_modules}\n") + + self.model = FastModel.get_peft_model( + self.model, + r=lora_r, + target_modules=target_modules, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + bias="none", + use_gradient_checkpointing=use_gradient_checkpointing, + random_state=3407, + use_rslora=use_rslora, + loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + task_type=None, + ) + + elif self._audio_type == 'snac': + # Orpheus uses FastLanguageModel.get_peft_model + print(f"Audio model ({self._audio_type}) LoRA configuration:") + print(f" - Target modules: {target_modules}\n") + + self.model = FastLanguageModel.get_peft_model( + self.model, + r=lora_r, + target_modules=target_modules, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + bias="none", + use_gradient_checkpointing=use_gradient_checkpointing, + random_state=3407, + use_rslora=use_rslora, + loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + ) + + elif self.is_vlm: # Vision model LoRA print(f"Vision model LoRA configuration:") print(f" - Finetune vision layers: {finetune_vision_layers}") @@ -346,6 +824,937 @@ class UnslothTrainer: self._update_progress(error=error_details) return False + def _apply_csm_forward_fix(self): + """Monkey-patch CsmForConditionalGeneration.forward to fix depth decoder kwargs. + + The original transformers forward passes raw **kwargs (num_items_in_batch, + causal_mask, etc.) from the Trainer/PEFT through to the depth decoder, + causing depth_decoder_loss=None and 'Tensor + NoneType' crash. + + We patch at both instance AND class level for maximum reliability, + and strip non-TransformersKwargs params that Unsloth/PEFT inject. + """ + import types + import torch + import torch.nn as nn + from transformers.models.csm.modeling_csm import ( + CsmForConditionalGeneration, + CsmOutputWithPast, + ) + + base_csm = self.model.base_model.model # CsmForConditionalGeneration + + # Save original forward (the @can_return_tuple wrapped version) + _original_forward = CsmForConditionalGeneration.forward + + # Keys that the depth decoder and its sub-layers actually understand + _TRANSFORMERS_KWARGS = { + 'num_items_in_batch', 'output_hidden_states', 'output_attentions', + 'output_router_logits', 'cu_seq_lens_q', 'cu_seq_lens_k', + 'max_length_q', 'max_length_k', + } + + def _fixed_csm_forward( + self, + input_ids=None, input_values=None, attention_mask=None, + input_values_cutoffs=None, position_ids=None, past_key_values=None, + inputs_embeds=None, labels=None, use_cache=None, + cache_position=None, logits_to_keep=0, **kwargs, + ): + # Strip non-standard kwargs injected by Unsloth/PEFT (causal_mask, + # num_logits_to_keep, task_ids, return_dict, etc.) + output_attentions = kwargs.pop('output_attentions', None) + output_hidden_states = kwargs.pop('output_hidden_states', None) + kwargs.pop('return_dict', None) + kwargs.pop('causal_mask', None) + kwargs.pop('num_logits_to_keep', None) + kwargs.pop('task_ids', None) + + # Only keep recognized TransformersKwargs + clean_kwargs = {k: v for k, v in kwargs.items() if k in _TRANSFORMERS_KWARGS} + + if input_ids is not None and input_ids.ndim == 2: + merged = self._merge_input_ids_with_input_values( + input_ids, input_values, input_values_cutoffs, labels + ) + inputs_embeds = merged["inputs_embeds"] + labels = merged["labels"] + input_ids = None + + backbone_outputs = self.backbone_model( + input_ids=input_ids, attention_mask=attention_mask, + position_ids=position_ids, past_key_values=past_key_values, + inputs_embeds=inputs_embeds, use_cache=use_cache, + cache_position=cache_position, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + **clean_kwargs, + ) + + backbone_hidden_states = backbone_outputs[0] + slice_indices = ( + slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) + else logits_to_keep + ) + backbone_logits = self.lm_head(backbone_hidden_states[:, slice_indices, :]) + + loss = None + backbone_loss = None + depth_decoder_loss = None + depth_decoder_outputs = None + if labels is not None: + backbone_labels = labels[:, :, 0] + backbone_loss = self.loss_function( + logits=backbone_logits, labels=backbone_labels, + vocab_size=self.config.vocab_size, **clean_kwargs, + ) + + train_mask = ~(labels[:, :, 1:] == -100).all(dim=-1) + depth_decoder_input_ids = labels[train_mask][..., :self.config.num_codebooks - 1] + depth_decoder_input_ids = nn.functional.pad( + depth_decoder_input_ids, (1, 0), value=0 + ) + + train_idxs = train_mask.nonzero(as_tuple=True) + backbone_last_hidden_states = backbone_hidden_states[ + train_idxs[0], train_idxs[1] - 1, : + ] + depth_decoder_labels = labels[train_mask] + + # Build clean kwargs for depth decoder + dd_kwargs = clean_kwargs.copy() + # Scale num_items_in_batch for depth decoder (31 codebooks) + if 'num_items_in_batch' in dd_kwargs: + dd_kwargs['num_items_in_batch'] = ( + dd_kwargs['num_items_in_batch'] * (self.config.num_codebooks - 1) + ) + + depth_decoder_outputs = self.depth_decoder( + input_ids=depth_decoder_input_ids, + backbone_last_hidden_state=backbone_last_hidden_states, + use_cache=False, return_dict=True, + labels=depth_decoder_labels, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + **dd_kwargs, + ) + + depth_decoder_loss = depth_decoder_outputs.loss + if depth_decoder_loss is None: + logger.warning( + "CSM depth_decoder_loss is None! " + f"labels shape={depth_decoder_labels.shape}, " + f"train_mask sum={train_mask.sum().item()}" + ) + # Fallback: use only backbone loss instead of crashing + loss = backbone_loss + else: + loss = backbone_loss + depth_decoder_loss + + return CsmOutputWithPast( + loss=loss, backbone_loss=backbone_loss, + depth_decoder_loss=depth_decoder_loss, logits=backbone_logits, + past_key_values=backbone_outputs.past_key_values, + hidden_states=backbone_outputs.hidden_states, + attentions=backbone_outputs.attentions, + depth_decoder_logits=( + depth_decoder_outputs.logits if depth_decoder_outputs else None + ), + depth_decoder_past_key_values=( + depth_decoder_outputs.past_key_values if depth_decoder_outputs else None + ), + depth_decoder_hidden_states=( + depth_decoder_outputs.hidden_states if depth_decoder_outputs else None + ), + depth_decoder_attentions=( + depth_decoder_outputs.attentions if depth_decoder_outputs else None + ), + ) + + # Patch at BOTH instance and class level for maximum reliability. + # Instance-level: catches calls via BaseTuner.forward -> self.model.forward() + base_csm.forward = types.MethodType(_fixed_csm_forward, base_csm) + # Class-level: catches any path that resolves through the class dict + CsmForConditionalGeneration.forward = _fixed_csm_forward + print("Applied CSM forward fix (class + instance level)\n") + + def _preprocess_csm_dataset(self, dataset, custom_format_mapping=None): + """Preprocess dataset for CSM TTS training (exact notebook copy).""" + from transformers import AutoProcessor + from datasets import Audio + import torch + + processor = AutoProcessor.from_pretrained(self.model_name) + + # Strip pad_to_multiple_of from tokenizer init_kwargs — fine-tuned models + # (e.g. keanteng/sesame-csm-elise) save it in tokenizer_config.json, and + # _merge_kwargs leaks it into audio_kwargs where EncodecFeatureExtractor rejects it. + processor.tokenizer.init_kwargs.pop('pad_to_multiple_of', None) + + # Resolve columns from user mapping or hardcoded fallback + resolved = self._resolve_audio_columns(dataset, custom_format_mapping) + audio_col = resolved["audio_col"] + text_col = resolved["text_col"] + speaker_key = resolved["speaker_col"] + + if audio_col is None: + raise ValueError(f"No audio column found in dataset. Columns: {dataset.column_names}") + if text_col is None: + raise ValueError(f"No text column found in dataset. Columns: {dataset.column_names}") + if speaker_key is None: + print("No speaker found, adding default 'source' of 0 for all examples\n") + dataset = dataset.add_column("source", ["0"] * len(dataset)) + speaker_key = "source" + + print(f"CSM preprocessing: audio_col='{audio_col}', text_col='{text_col}', speaker_key='{speaker_key}'\n") + + dataset = dataset.cast_column(audio_col, Audio(sampling_rate=24000)) + + required_keys = ["input_ids", "attention_mask", "labels", "input_values", "input_values_cutoffs"] + + self._update_progress(status_message="Preprocessing CSM dataset...") + processed_examples = [] + skipped = 0 + for idx in range(len(dataset)): + if self.should_stop: + print("Stopped during CSM preprocessing\n") + break + + example = dataset[idx] + try: + conversation = [{ + "role": str(example[speaker_key]), + "content": [ + {"type": "text", "text": example.get(text_col, "")}, + {"type": "audio", "path": example[audio_col]["array"]}, + ], + }] + # NOTE: pad_to_multiple_of intentionally omitted from text_kwargs — + # CsmProcessor._merge_kwargs leaks it to EncodecFeatureExtractor which rejects it. + model_inputs = processor.apply_chat_template( + conversation, + tokenize=True, + return_dict=True, + output_labels=True, + text_kwargs={ + "padding": "max_length", + "max_length": 256, + "padding_side": "right", + }, + audio_kwargs={ + "sampling_rate": 24_000, + "max_length": 240001, + "padding": "max_length", + }, + common_kwargs={"return_tensors": "pt"}, + ) + + out = {} + for k in required_keys: + if k not in model_inputs: + raise KeyError(f"Missing required key '{k}' in model outputs") + out[k] = model_inputs[k][0] + + if not all(isinstance(out[k], torch.Tensor) for k in out): + skipped += 1 + continue + + processed_examples.append(out) + + except Exception as e: + logger.warning(f"Error processing CSM example {idx}: {e}") + skipped += 1 + continue + + if (idx + 1) % 100 == 0: + self._update_progress( + status_message=f"Preprocessing CSM... {idx + 1}/{len(dataset)}" + ) + + if not processed_examples: + raise ValueError( + f"No valid examples after CSM preprocessing (skipped {skipped})" + ) + + result_dataset = Dataset.from_list(processed_examples) + print(f"CSM preprocessing complete: {len(result_dataset)} examples " + f"({skipped} skipped)\n") + return result_dataset + + def _format_audio_vlm_dataset(self, dataset, custom_format_mapping=None): + """Format dataset as audio chat messages for multimodal models (e.g. Gemma 3N). + + Expects columns: audio (Audio), text (str). + Produces: messages column with system/user/assistant chat format. + """ + from datasets import Audio + + resolved = self._resolve_audio_columns(dataset, custom_format_mapping) + audio_col = resolved["audio_col"] + text_col = resolved["text_col"] + if not audio_col or not text_col: + raise ValueError( + f"Audio VLM dataset needs 'audio' and 'text' columns, got: {dataset.column_names}" + ) + + # Store resolved audio column name for the collator closure + self._audio_vlm_audio_col = audio_col + + # Cast audio to 16kHz (standard for speech models) + dataset = dataset.cast_column(audio_col, Audio(sampling_rate=16000)) + + def format_messages(samples): + formatted = {"messages": []} + for idx in range(len(samples[audio_col])): + audio = samples[audio_col][idx]["array"] + label = str(samples[text_col][idx]) + message = [ + {"role": "system", "content": [ + {"type": "text", "text": "You are an assistant that transcribes speech accurately."} + ]}, + {"role": "user", "content": [ + {"type": "audio", "audio": audio}, + {"type": "text", "text": "Please transcribe this audio."} + ]}, + {"role": "assistant", "content": [ + {"type": "text", "text": label} + ]}, + ] + formatted["messages"].append(message) + return formatted + + self._update_progress(status_message="Formatting audio VLM dataset...") + dataset = dataset.map(format_messages, batched=True, batch_size=4, num_proc=safe_num_proc(4)) + print(f"Audio VLM dataset formatted: {len(dataset)} examples\n") + return dataset + + def _preprocess_snac_dataset(self, dataset, custom_format_mapping=None): + """Preprocess dataset for Orpheus TTS training with SNAC codec. + + Mirrors Orpheus_(3B)-TTS.ipynb: encode audio with SNAC (24kHz, 3 hierarchical + layers), interleave 7 codes per frame, wrap with Orpheus special tokens, + train on full sequence (no label masking). + """ + import torch + import torchaudio.transforms as T + + SNAC_MODEL_NAME = "hubertsiuzdak/snac_24khz" + SNAC_SAMPLE_RATE = 24000 + device = "cuda" if torch.cuda.is_available() else "cpu" + max_length = self.max_seq_length or 2048 + tokenizer = self.tokenizer + + # Orpheus special token IDs (hardcoded in tokenizer vocabulary) + START_OF_HUMAN = 128259 + END_OF_HUMAN = 128260 + START_OF_AI = 128261 + END_OF_AI = 128262 + START_OF_SPEECH = 128257 + END_OF_SPEECH = 128258 + END_OF_TEXT = 128009 + AUDIO_OFFSET = 128266 + + resolved = self._resolve_audio_columns(dataset, custom_format_mapping) + audio_col = resolved["audio_col"] + text_col = resolved["text_col"] + speaker_col = resolved["speaker_col"] + has_source = speaker_col is not None + if not audio_col or not text_col: + raise ValueError( + f"SNAC dataset needs 'audio' and 'text' columns, got: {dataset.column_names}" + ) + + # Cast audio column so datasets 4.x AudioDecoder objects are decoded to dicts + from datasets import Audio + dataset = dataset.cast_column(audio_col, Audio(sampling_rate=SNAC_SAMPLE_RATE)) + + # Get dataset sample rate from first example (after cast, always SNAC_SAMPLE_RATE) + first_audio = dataset[0][audio_col] + ds_sample_rate = first_audio.get("sampling_rate", SNAC_SAMPLE_RATE) if isinstance(first_audio, dict) else SNAC_SAMPLE_RATE + + # Load SNAC codec model + self._update_progress(status_message="Loading SNAC codec model...") + print("Loading SNAC codec model...\n") + from snac import SNAC + snac_model = SNAC.from_pretrained(SNAC_MODEL_NAME) + snac_model = snac_model.to(device).eval() + + # Resample transform (created once) + resample_transform = T.Resample(orig_freq=ds_sample_rate, new_freq=SNAC_SAMPLE_RATE) if ds_sample_rate != SNAC_SAMPLE_RATE else None + + self._update_progress(status_message="Encoding audio with SNAC...") + print(f"SNAC preprocessing: audio_col='{audio_col}', text_col='{text_col}', " + f"has_source={has_source}, ds_sample_rate={ds_sample_rate}\n") + + processed_examples = [] + skipped = 0 + for idx in range(len(dataset)): + if self.should_stop: + print("Stopped during SNAC preprocessing\n") + break + + example = dataset[idx] + try: + text = example.get(text_col) + if not text: + skipped += 1 + continue + + audio_data = example.get(audio_col) + if audio_data is None or audio_data.get("array") is None: + skipped += 1 + continue + + # --- Encode audio with SNAC (notebook lines 122-142) --- + waveform = torch.from_numpy(audio_data["array"]).unsqueeze(0).to(dtype=torch.float32) + if resample_transform is not None: + waveform = resample_transform(waveform) + + waveform = waveform.unsqueeze(0).to(device) + with torch.inference_mode(): + codes = snac_model.encode(waveform) + + # Interleave 7 codes per frame with layer offsets (notebook lines 134-142) + all_codes = [] + for i in range(codes[0].shape[1]): + all_codes.append(codes[0][0][i].item() + AUDIO_OFFSET) + all_codes.append(codes[1][0][2*i].item() + AUDIO_OFFSET + 4096) + all_codes.append(codes[2][0][4*i].item() + AUDIO_OFFSET + (2*4096)) + all_codes.append(codes[2][0][(4*i)+1].item() + AUDIO_OFFSET + (3*4096)) + all_codes.append(codes[1][0][(2*i)+1].item() + AUDIO_OFFSET + (4*4096)) + all_codes.append(codes[2][0][(4*i)+2].item() + AUDIO_OFFSET + (5*4096)) + all_codes.append(codes[2][0][(4*i)+3].item() + AUDIO_OFFSET + (6*4096)) + + if len(all_codes) == 0: + skipped += 1 + continue + + # Deduplicate consecutive frames with same first code (notebook lines 185-207) + deduped = all_codes[:7] + for i in range(7, len(all_codes), 7): + if all_codes[i] != deduped[-7]: + deduped.extend(all_codes[i:i+7]) + all_codes = deduped + + # --- Build text tokens (notebook lines 217-224) --- + text_prompt = f"{example[speaker_col]}: {text}" if has_source and example.get(speaker_col) else text + text_ids = tokenizer.encode(text_prompt, add_special_tokens=True) + text_ids.append(END_OF_TEXT) + + # --- Build full input_ids (notebook lines 225-234) --- + input_ids = ( + [START_OF_HUMAN] + + text_ids + + [END_OF_HUMAN] + + [START_OF_AI] + + [START_OF_SPEECH] + + all_codes + + [END_OF_SPEECH] + + [END_OF_AI] + ) + + # Truncate to max_length + input_ids = input_ids[:max_length] + + # Labels = input_ids (no masking — Orpheus trains on full sequence) + labels = list(input_ids) + attention_mask = [1] * len(input_ids) + + processed_examples.append({ + "input_ids": input_ids, + "labels": labels, + "attention_mask": attention_mask, + }) + + except Exception as e: + logger.warning(f"Error processing SNAC example {idx}: {e}") + skipped += 1 + continue + + # Progress update every 100 examples + if (idx + 1) % 100 == 0: + self._update_progress( + status_message=f"Encoding audio... {idx + 1}/{len(dataset)}" + ) + + # Free SNAC model from GPU + print("Freeing SNAC codec model from GPU...\n") + snac_model.to("cpu") + del snac_model + import gc + gc.collect() + torch.cuda.empty_cache() + self._cuda_audio_used = True + + if not processed_examples: + raise ValueError( + f"No valid examples after SNAC preprocessing (skipped {skipped})" + ) + + result_dataset = Dataset.from_list(processed_examples) + print(f"SNAC preprocessing complete: {len(result_dataset)} examples " + f"({skipped} skipped)\n") + return result_dataset + + def _preprocess_bicodec_dataset(self, dataset, custom_format_mapping=None): + """Preprocess dataset for Spark-TTS training with BiCodec tokenizer. + + Mirrors Spark_TTS_(0_5B).ipynb: encode audio with BiCodec (semantic + global tokens), + format as special-token text strings for SFTTrainer with dataset_text_field="text". + """ + import sys + import torch + import numpy as np + import torchaudio.transforms as T + + import subprocess + + device = "cuda" if torch.cuda.is_available() else "cpu" + + # The sparktts Python package lives in the SparkAudio/Spark-TTS GitHub repo, + # NOT in the unsloth/Spark-TTS-0.5B HF model repo. Clone it if needed. + spark_code_dir = os.path.join(os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS") + sparktts_pkg = os.path.join(spark_code_dir, "sparktts") + if not os.path.isdir(sparktts_pkg): + self._update_progress(status_message="Cloning Spark-TTS code repo...") + print(f"Cloning SparkAudio/Spark-TTS to {spark_code_dir}...\n") + 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 + from sparktts.utils.audio import audio_volume_normalize + + # Resolve audio and text columns + resolved = self._resolve_audio_columns(dataset, custom_format_mapping) + audio_col = resolved["audio_col"] + text_col = resolved["text_col"] + speaker_col = resolved["speaker_col"] + has_source = speaker_col is not None + if not audio_col or not text_col: + raise ValueError( + f"BiCodec dataset needs 'audio' and 'text' columns, got: {dataset.column_names}" + ) + + # Cast audio column so datasets 4.x AudioDecoder objects are decoded to dicts. + # Don't resample here — BiCodec's target_sr may differ; the loop handles resampling. + from datasets import Audio + dataset = dataset.cast_column(audio_col, Audio()) + + # Load BiCodec tokenizer + self._update_progress(status_message="Loading BiCodec tokenizer...") + print("Loading BiCodec tokenizer...\n") + audio_tokenizer = BiCodecTokenizer(self._spark_tts_repo_dir, device) + + target_sr = audio_tokenizer.config['sample_rate'] + + self._update_progress(status_message="Encoding audio with BiCodec...") + print(f"BiCodec preprocessing: audio_col='{audio_col}', text_col='{text_col}', " + f"has_source={has_source}, target_sr={target_sr}\n") + + def extract_wav2vec2_features(wavs: torch.Tensor) -> torch.Tensor: + """Extract wav2vec2 features (average of layers 11, 14, 16).""" + if wavs.shape[0] != 1: + raise ValueError(f"Expected batch size 1, but got shape {wavs.shape}") + wav_np = wavs.squeeze(0).cpu().numpy() + + processed = audio_tokenizer.processor( + wav_np, + sampling_rate=16000, + return_tensors="pt", + padding=True, + ) + input_values = processed.input_values.to(audio_tokenizer.feature_extractor.device) + model_output = audio_tokenizer.feature_extractor(input_values) + + if model_output.hidden_states is None: + raise ValueError("Wav2Vec2Model did not return hidden states.") + + feats_mix = ( + model_output.hidden_states[11] + + model_output.hidden_states[14] + + model_output.hidden_states[16] + ) / 3 + return feats_mix + + processed_examples = [] + skipped = 0 + for idx in range(len(dataset)): + if self.should_stop: + print("Stopped during BiCodec preprocessing\n") + break + + example = dataset[idx] + try: + text = example.get(text_col) + if not text: + skipped += 1 + continue + + audio_data = example.get(audio_col) + if audio_data is None or audio_data.get("array") is None: + skipped += 1 + continue + + audio_array = audio_data["array"] + sampling_rate = audio_data.get("sampling_rate", target_sr) + + # Resample if needed + if sampling_rate != target_sr: + resampler = T.Resample(orig_freq=sampling_rate, new_freq=target_sr) + audio_tensor_temp = torch.from_numpy(audio_array).float() + audio_array = resampler(audio_tensor_temp).numpy() + + # Volume normalize if configured + if audio_tokenizer.config.get("volume_normalize", False): + audio_array = audio_volume_normalize(audio_array) + + # Get reference clip + ref_wav_np = audio_tokenizer.get_ref_clip(audio_array) + + # Prepare tensors + audio_tensor = torch.from_numpy(audio_array).unsqueeze(0).float().to(device) + ref_wav_tensor = torch.from_numpy(ref_wav_np).unsqueeze(0).float().to(device) + + # Extract wav2vec2 features + feat = extract_wav2vec2_features(audio_tensor) + + batch = { + "wav": audio_tensor, + "ref_wav": ref_wav_tensor, + "feat": feat.to(device), + } + + # BiCodec tokenize + semantic_token_ids, global_token_ids = audio_tokenizer.model.tokenize(batch) + + global_tokens = "".join( + [f"<|bicodec_global_{i}|>" for i in global_token_ids.squeeze().cpu().numpy()] + ) + semantic_tokens = "".join( + [f"<|bicodec_semantic_{i}|>" for i in semantic_token_ids.squeeze().cpu().numpy()] + ) + + # Format text with source prefix if available + text_content = f"{example[speaker_col]}: {text}" if has_source and example.get(speaker_col) else text + + formatted = "".join([ + "<|task_tts|>", + "<|start_content|>", + text_content, + "<|end_content|>", + "<|start_global_token|>", + global_tokens, + "<|end_global_token|>", + "<|start_semantic_token|>", + semantic_tokens, + "<|end_semantic_token|>", + "<|im_end|>", + ]) + + processed_examples.append({"text": formatted}) + + except Exception as e: + logger.warning(f"Error processing BiCodec example {idx}: {e}") + skipped += 1 + continue + + # Progress update every 100 examples + if (idx + 1) % 100 == 0: + self._update_progress( + status_message=f"Encoding audio with BiCodec... {idx + 1}/{len(dataset)}" + ) + + # Free BiCodec model from GPU + print("Freeing BiCodec tokenizer from GPU...\n") + audio_tokenizer.model.cpu() + audio_tokenizer.feature_extractor.cpu() + del audio_tokenizer + import gc + gc.collect() + torch.cuda.empty_cache() + self._cuda_audio_used = True + + if not processed_examples: + raise ValueError( + f"No valid examples after BiCodec preprocessing (skipped {skipped})" + ) + + result_dataset = Dataset.from_list(processed_examples) + print(f"BiCodec preprocessing complete: {len(result_dataset)} examples " + f"({skipped} skipped)\n") + # Debug: show first example text (truncated) + sample = result_dataset[0]["text"] + print(f"Sample text (first 200 chars): {sample[:200]}...\n") + print(f"Sample text length: {len(sample)} chars\n") + return result_dataset + + def _preprocess_dac_dataset(self, dataset, custom_format_mapping=None): + """Preprocess dataset for OuteTTS training with DAC codec. + + Mirrors Oute_TTS_(1B).ipynb DataCreationV3: uses Whisper for word timings, + OuteTTS AudioProcessor for speaker representations, PromptProcessor for + training prompts. Outputs text strings for SFTTrainer with dataset_text_field="text". + """ + import sys + import io + import tempfile + import torch + import numpy as np + import soundfile as sf + from datasets import Dataset as HFDataset + + device = "cuda" if torch.cuda.is_available() else "cpu" + + # Clone OuteTTS repo (same as audio_codecs._load_dac) + import subprocess + base_dir = os.path.dirname(os.path.abspath(__file__)) + outetts_code_dir = os.path.join(base_dir, "inference", "OuteTTS") + outetts_pkg = os.path.join(outetts_code_dir, "outetts") + if not os.path.isdir(outetts_pkg): + self._update_progress(status_message="Cloning OuteTTS code repo...") + print(f"Cloning edwko/OuteTTS to {outetts_code_dir}...\n") + subprocess.run( + ["git", "clone", "--depth", "1", "https://github.com/edwko/OuteTTS", outetts_code_dir], + check=True, + ) + for fpath in [ + os.path.join(outetts_pkg, "models", "gguf_model.py"), + os.path.join(outetts_pkg, "interface.py"), + os.path.join(outetts_pkg, "__init__.py"), + ]: + if os.path.exists(fpath): + os.remove(fpath) + print(f"Removed {fpath}\n") + + 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.version.v3.prompt_processor import PromptProcessor + from outetts.models.config import ModelConfig as OuteTTSModelConfig + from outetts.utils.preprocessing import text_normalizations + + # Resolve audio and text columns + resolved = self._resolve_audio_columns(dataset, custom_format_mapping) + audio_col = resolved["audio_col"] + text_col = resolved["text_col"] + if not audio_col or not text_col: + raise ValueError( + f"DAC dataset needs 'audio' and 'text' columns, got: {dataset.column_names}" + ) + + # Cast audio to 24kHz (notebook: dataset.cast_column("audio", Audio(sampling_rate=24000))) + from datasets import Audio + dataset = dataset.cast_column(audio_col, Audio(sampling_rate=24000)) + print("Cast audio column to 24kHz\n") + + # Load Whisper for word timings + self._update_progress(status_message="Loading Whisper model for word timings...") + print("Loading Whisper model for word timings...\n") + import whisper + whisper_model = whisper.load_model("turbo", device=device) + + # Load OuteTTS AudioProcessor + PromptProcessor + self._update_progress(status_message="Loading OuteTTS AudioProcessor...") + print("Loading OuteTTS AudioProcessor...\n") + model_tokenizer_path = "OuteAI/Llama-OuteTTS-1.0-1B" + dummy_config = OuteTTSModelConfig( + tokenizer_path=model_tokenizer_path, + device=device, + audio_codec_path=None, + ) + audio_processor = AudioProcessor(config=dummy_config) + prompt_processor = PromptProcessor(model_tokenizer_path) + + self._update_progress(status_message="Preprocessing audio with OuteTTS...") + print(f"DAC preprocessing: audio_col='{audio_col}', text_col='{text_col}'\n") + + processed_examples = [] + skipped = 0 + for idx in range(len(dataset)): + if self.should_stop: + print("Stopped during DAC preprocessing\n") + break + + example = dataset[idx] + try: + text = example.get(text_col) + if not text or not isinstance(text, str): + skipped += 1 + continue + + audio_data = example.get(audio_col) + if audio_data is None or audio_data.get("array") is None: + skipped += 1 + continue + + audio_array = np.array(audio_data["array"], dtype=np.float32) + sampling_rate = audio_data.get("sampling_rate", 24000) + + # Convert to WAV bytes (Whisper needs a file path) + buf = io.BytesIO() + sf.write(buf, audio_array, sampling_rate, format="WAV", subtype="FLOAT") + buf.seek(0) + audio_bytes = buf.getvalue() + + # 1. Get word timings from Whisper + with tempfile.NamedTemporaryFile(suffix=".wav", delete=True) as tmp: + tmp.write(audio_bytes) + tmp.flush() + whisper_result = whisper_model.transcribe(tmp.name, word_timestamps=True) + + normalized_transcript = text_normalizations(text) + words_with_timings = [] + if whisper_result and "segments" in whisper_result: + for segment in whisper_result["segments"]: + for word_info in segment.get("words", []): + cleaned = word_info["word"].strip() + if cleaned: + words_with_timings.append({ + "word": cleaned, + "start": float(word_info["start"]), + "end": float(word_info["end"]), + }) + + if not words_with_timings: + skipped += 1 + continue + + # 2. Create speaker representation with AudioProcessor + speaker_data_dict = { + "audio": {"bytes": audio_bytes}, + "text": normalized_transcript, + "words": words_with_timings, + } + speaker = audio_processor.create_speaker_from_dict(speaker_data_dict) + if speaker is None: + skipped += 1 + continue + + # 3. Get training prompt from PromptProcessor + prompt = prompt_processor.get_training_prompt(speaker) + if prompt: + processed_examples.append({"text": prompt}) + + except Exception as e: + logger.warning(f"Error processing DAC example {idx}: {e}") + skipped += 1 + continue + + if (idx + 1) % 100 == 0: + self._update_progress( + status_message=f"Preprocessing audio with OuteTTS... {idx + 1}/{len(dataset)}" + ) + + # Free Whisper from GPU (notebook: data_processor.whisper_model.to('cpu')) + print("Moving Whisper model to CPU...\n") + whisper_model.to('cpu') + del whisper_model + del audio_processor + del prompt_processor + import gc + gc.collect() + torch.cuda.empty_cache() + self._cuda_audio_used = True + + if not processed_examples: + raise ValueError( + f"No valid examples after DAC preprocessing (skipped {skipped})" + ) + + result_dataset = HFDataset.from_list(processed_examples) + print(f"DAC preprocessing complete: {len(result_dataset)} examples " + f"({skipped} skipped)\n") + sample = result_dataset[0]["text"] + print(f"Sample text (first 200 chars): {sample[:200]}...\n") + return result_dataset + + def _preprocess_whisper_dataset(self, dataset, eval_split=None, custom_format_mapping=None): + """Preprocess dataset for Whisper speech-to-text training. + + Mirrors Whisper.ipynb: extract audio features with Whisper's feature + extractor, tokenize text labels. Returns (train_data, eval_data) where + each is a list of dicts with 'input_features' and 'labels'. + """ + from datasets import Audio + + WHISPER_SAMPLE_RATE = 16000 + + resolved = self._resolve_audio_columns(dataset, custom_format_mapping) + audio_col = resolved["audio_col"] + text_col = resolved["text_col"] + if not audio_col or not text_col: + raise ValueError( + f"Whisper dataset needs 'audio' and 'text' columns, got: {dataset.column_names}" + ) + + # Cast audio to 16kHz (Whisper's expected sample rate) + dataset = dataset.cast_column(audio_col, Audio(sampling_rate=WHISPER_SAMPLE_RATE)) + + # Train/eval split (notebook does dataset.train_test_split) + eval_dataset_raw = None + if eval_split: + splits = dataset.train_test_split(test_size=0.06, seed=42) + dataset = splits["train"] + eval_dataset_raw = splits["test"] + + self._update_progress(status_message="Processing audio for Whisper...") + print(f"Whisper preprocessing: audio_col='{audio_col}', text_col='{text_col}', " + f"samples={len(dataset)}\n") + + def process_split(ds, split_name="train"): + processed = [] + skipped = 0 + for idx in range(len(ds)): + if self.should_stop: + print(f"Stopped during Whisper {split_name} preprocessing\n") + break + + example = ds[idx] + try: + audio_data = example.get(audio_col) + text = example.get(text_col) + if audio_data is None or audio_data.get("array") is None or not text: + skipped += 1 + continue + + # Extract audio features (notebook line 112-115) + features = self.tokenizer.feature_extractor( + audio_data["array"], sampling_rate=audio_data["sampling_rate"] + ) + # Tokenize text (notebook line 116) + tokenized_text = self.tokenizer.tokenizer(text) + + processed.append({ + "input_features": features.input_features[0], + "labels": tokenized_text.input_ids, + }) + except Exception as e: + logger.warning(f"Error processing Whisper {split_name} example {idx}: {e}") + skipped += 1 + continue + + if (idx + 1) % 100 == 0: + self._update_progress( + status_message=f"Processing {split_name} audio... {idx + 1}/{len(ds)}" + ) + + print(f"Whisper {split_name} preprocessing: {len(processed)} examples ({skipped} skipped)\n") + return processed + + train_data = process_split(dataset, "train") + eval_data = process_split(eval_dataset_raw, "eval") if eval_dataset_raw else None + + if not train_data: + raise ValueError("No valid examples after Whisper preprocessing") + + return (train_data, eval_data) + def load_and_format_dataset(self, dataset_source: str, format_type: str = "auto", @@ -485,6 +1894,33 @@ class UnslothTrainer: print("Stopped before applying chat template\n") return None + # ========== AUDIO MODELS: custom preprocessing ========== + if self._audio_type == 'csm': + processed = self._preprocess_csm_dataset(dataset, custom_format_mapping) + return (processed, None) + + elif self._audio_type == 'whisper': + train_data, eval_data = self._preprocess_whisper_dataset( + dataset, eval_split=eval_split, custom_format_mapping=custom_format_mapping + ) + return (train_data, eval_data) + + elif self._audio_type == 'snac': + processed = self._preprocess_snac_dataset(dataset, custom_format_mapping) + return (processed, None) + + elif self._audio_type == 'bicodec': + processed = self._preprocess_bicodec_dataset(dataset, custom_format_mapping) + return ({"dataset": processed, "final_format": "audio_bicodec"}, None) + + elif self._audio_type == 'dac': + processed = self._preprocess_dac_dataset(dataset, custom_format_mapping) + return ({"dataset": processed, "final_format": "audio_dac"}, None) + + elif self.is_audio_vlm: + formatted = self._format_audio_vlm_dataset(dataset, custom_format_mapping) + return (formatted, None) + # ========== FORMAT FIRST ========== print(f"Formatting dataset with format_type='{format_type}'...\n") @@ -552,7 +1988,7 @@ class UnslothTrainer: from datasets import get_dataset_split_names load_kwargs = {"path": dataset_source} if subset: - load_kwargs["name"] = subset + load_kwargs["config_name"] = subset available_splits = get_dataset_split_names(**load_kwargs) print(f"Available splits: {available_splits}\n") @@ -632,6 +2068,22 @@ class UnslothTrainer: self._update_progress(error="Model not loaded") return False + # Pre-import heavy transformers modules on the main thread. + # Unsloth's patched_import hook (deepseek_v3_moe.py) is not thread-safe + # with Python's importlib cache, causing KeyError: 'size' if these are + # first imported inside the worker thread. + import transformers # noqa: F401 – ensures submodules are cached + from transformers import ( # noqa: F401 + Trainer as _HFTrainer, + TrainingArguments as _TrainingArguments, + TrainerCallback as _TrainerCallback, + ) + if self._audio_type == 'whisper': + from transformers import ( # noqa: F401 + Seq2SeqTrainer as _Seq2SeqTrainer, + Seq2SeqTrainingArguments as _Seq2SeqTrainingArguments, + ) + # Start training in separate thread self.training_thread = threading.Thread( target=self._train_worker, @@ -694,6 +2146,107 @@ class UnslothTrainer: output_dir = training_args.get('output_dir', './outputs') os.makedirs(output_dir, exist_ok=True) + # ========== AUDIO TRAINER BRANCH ========== + if self._audio_type == 'csm': + # CSM uses plain HF Trainer (NOT SFTTrainer) + # Needs remove_unused_columns=False for depth decoder (input_values + cutoffs) + from transformers import Trainer as HFTrainer, TrainingArguments + self._apply_csm_forward_fix() + + config = self._build_audio_training_args(training_args, output_dir, extra_args={ + "remove_unused_columns": False, + }) + self.trainer = HFTrainer( + model=self.model, train_dataset=dataset, + args=TrainingArguments(**config), + ) + self.trainer.add_callback(self._create_progress_callback()) + + batch_size = training_args.get('batch_size', 2) + total = self._calculate_total_steps( + len(dataset), batch_size, + training_args.get('gradient_accumulation_steps', 4), + training_args.get('num_epochs', 3), + training_args.get('max_steps', 0), + ) + self._update_progress(total_steps=total, status_message="Starting CSM training...") + print(f"CSM training config: {config}\n") + self.trainer.train() + self._finalize_training(output_dir, "CSM") + return + + elif self._audio_type == 'snac': + # Orpheus: language model with SNAC codec tokens — plain HF Trainer + # DataCollatorForSeq2Seq dynamically pads variable-length sequences per batch + # (text + audio codes vary in length) and pads labels with -100. + from transformers import Trainer as HFTrainer, TrainingArguments, DataCollatorForSeq2Seq + + config = self._build_audio_training_args(training_args, output_dir) + self.trainer = HFTrainer( + model=self.model, train_dataset=dataset, + args=TrainingArguments(**config), + data_collator=DataCollatorForSeq2Seq( + tokenizer=self.tokenizer, padding=True, pad_to_multiple_of=8, + ), + ) + self.trainer.add_callback(self._create_progress_callback()) + + batch_size = training_args.get('batch_size', 2) + total = self._calculate_total_steps( + len(dataset), batch_size, + training_args.get('gradient_accumulation_steps', 4), + training_args.get('num_epochs', 3), + training_args.get('max_steps', 0), + ) + self._update_progress(total_steps=total, status_message="Starting SNAC training...") + print(f"SNAC training config: {config}\n") + self.trainer.train() + self._finalize_training(output_dir, "SNAC") + return + + elif self._audio_type == 'whisper': + # Whisper: Seq2SeqTrainer with custom speech collator + from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments + from utils.datasets import DataCollatorSpeechSeq2SeqWithPadding + + eval_dataset = training_args.get('eval_dataset', None) + extra = {"remove_unused_columns": False, "label_names": ["labels"]} + if eval_dataset: + extra["eval_strategy"] = "steps" + extra["eval_steps"] = training_args.get('eval_steps', 5) + + config = self._build_audio_training_args(training_args, output_dir, extra_args=extra) + + trainer_kwargs = { + "model": self.model, + "train_dataset": dataset, + "data_collator": DataCollatorSpeechSeq2SeqWithPadding(processor=self.tokenizer), + "processing_class": self.tokenizer.feature_extractor, + "args": Seq2SeqTrainingArguments(**config), + } + if eval_dataset: + trainer_kwargs["eval_dataset"] = eval_dataset + + self.trainer = Seq2SeqTrainer(**trainer_kwargs) + self.trainer.add_callback(self._create_progress_callback()) + + batch_size = training_args.get('batch_size', 2) + total = self._calculate_total_steps( + len(dataset), batch_size, + training_args.get('gradient_accumulation_steps', 4), + training_args.get('num_epochs', 3), + training_args.get('max_steps', 0), + ) + self._update_progress(total_steps=total, status_message="Starting Whisper training...") + print(f"Whisper training config: {config}\n") + self.trainer.train() + self._finalize_training(output_dir, "Whisper") + return + + elif self._audio_type is not None and self._audio_type not in ('bicodec', 'dac'): + # bicodec/dac use the standard SFTTrainer text path below + raise NotImplementedError(f"Audio training for '{self._audio_type}' not yet implemented") + # ========== DATA COLLATOR SELECTION ========== # Detect special model types model_name_lower = self.model_name.lower() @@ -738,8 +2291,43 @@ class UnslothTrainer: self._update_progress(error=error_msg, is_training=False) return + elif self.is_audio_vlm: + # Audio VLM collator (e.g. Gemma 3N with audio data) + # Mirrors the collate_fn from Gemma3N_(4B)-Audio notebook + print("Configuring audio VLM data collator...\n") + processor = self.tokenizer # FastModel returns processor as tokenizer + + audio_col_name = getattr(self, '_audio_vlm_audio_col', 'audio') + + def audio_vlm_collate_fn(examples): + texts = [] + audios = [] + for example in examples: + text = processor.apply_chat_template( + example["messages"], tokenize=False, add_generation_prompt=False + ).strip() + texts.append(text) + audios.append(example[audio_col_name]["array"]) + + batch = processor( + text=texts, audio=audios, return_tensors="pt", padding=True + ) + + # Labels = input_ids with special tokens masked + labels = batch["input_ids"].clone() + labels[labels == processor.tokenizer.pad_token_id] = -100 + for attr in ('audio_token_id', 'image_token_id', 'boi_token_id', 'eoi_token_id'): + token_id = getattr(processor.tokenizer, attr, None) + if token_id is not None: + labels[labels == token_id] = -100 + batch["labels"] = labels + return batch + + data_collator = audio_vlm_collate_fn + print("Audio VLM data collator configured\n") + elif self.is_vlm: - # Standard VLM collator + # Standard VLM collator (images) print("Using UnslothVisionDataCollator for vision model\n") from unsloth.trainer import UnslothVisionDataCollator @@ -748,19 +2336,18 @@ class UnslothTrainer: print("Vision data collator configured\n") # ========== TRAINING CONFIGURATION ========== - # Handle epochs vs max_steps properly - max_steps_val = training_args.get('max_steps', 0) - num_epochs_val = training_args.get('num_epochs', 3) - # Handle warmup_steps vs warmup_ratio warmup_steps_val = training_args.get('warmup_steps', None) warmup_ratio_val = training_args.get('warmup_ratio', None) + lr_value = training_args.get('learning_rate', 2e-4) + print(f"[DEBUG] learning_rate from training_args: {lr_value} (type: {type(lr_value).__name__})\n") + config_args = { "per_device_train_batch_size": training_args.get('batch_size', 2), "gradient_accumulation_steps": training_args.get('gradient_accumulation_steps', 4), "num_train_epochs": training_args.get('num_epochs', 3), # Default to epochs - "learning_rate": training_args.get('learning_rate', 2e-4), + "learning_rate": lr_value, "fp16": not is_bfloat16_supported(), "bf16": is_bfloat16_supported(), "logging_steps": 1, @@ -769,8 +2356,16 @@ class UnslothTrainer: "output_dir": output_dir, "report_to": ["wandb"] if training_args.get('enable_wandb', False) else "none", "include_num_input_tokens_seen": True, # Enable token counting - "dataset_num_proc": safe_num_proc(max(1, os.cpu_count() // 4)), + "dataset_num_proc": 1 if (self.is_audio or self.is_audio_vlm or self._cuda_audio_used) else safe_num_proc(max(1, os.cpu_count() // 4)), + "max_seq_length": training_args.get('max_seq_length', 2048), } + + # On Windows with transformers 5.x, disable DataLoader multiprocessing + # to avoid issues with modified sys.path (.venv_t5) in spawned workers. + if sys.platform == "win32": + import transformers as _tf + if _tf.__version__.startswith("5."): + config_args["dataloader_num_workers"] = 0 # Add warmup parameter - use warmup_ratio if provided, otherwise warmup_steps if warmup_ratio_val is not None: @@ -819,9 +2414,10 @@ class UnslothTrainer: optim_value = training_args.get('optim', "adamw_8bit") lr_scheduler_type_value = training_args.get('lr_scheduler_type', "linear") - if self.is_vlm: - # Vision-specific config - print("Configuring vision model training parameters\n") + if self.is_vlm or self.is_audio_vlm: + # Vision / audio VLM config (both need skip_prepare_dataset + remove_unused_columns) + label = "audio VLM" if self.is_audio_vlm else "vision" + print(f"Configuring {label} model training parameters\n") # Use provided values or defaults for vision models optim_value = training_args.get('optim', "adamw_torch_fused") lr_scheduler_type_value = training_args.get('lr_scheduler_type', "cosine") @@ -830,7 +2426,7 @@ class UnslothTrainer: "lr_scheduler_type": lr_scheduler_type_value, "gradient_checkpointing": True, "gradient_checkpointing_kwargs": {"use_reentrant": False}, - "max_grad_norm": 0.3, # Recommended for vision models + "max_grad_norm": 0.3, "remove_unused_columns": False, "dataset_text_field": "", "dataset_kwargs": {"skip_prepare_dataset": True}, @@ -850,14 +2446,39 @@ class UnslothTrainer: config_args["packing"] = packing_enabled print(f"Sequence packing: {'enabled' if packing_enabled else 'disabled'}\n") + # Audio codec overrides — BiCodec/DAC use the text SFTTrainer path + if self._audio_type == 'bicodec': + config_args["packing"] = False + print("Applied BiCodec overrides: packing=False\n") + elif self._audio_type == 'dac': + config_args["packing"] = False + print("Applied DAC overrides: packing=False\n") + print(f"The configuration is: {config_args}") print("Training configuration prepared\n") # ========== TRAINER INITIALIZATION ========== - if self.is_vlm: + if self.is_audio_vlm: + # Audio VLM (e.g. Gemma 3N + audio): raw Dataset from _format_audio_vlm_dataset + # Notebook uses processing_class=processor.tokenizer (text tokenizer only) + train_dataset = dataset if isinstance(dataset, Dataset) else dataset['dataset'] + processing_class = self.tokenizer.tokenizer if hasattr(self.tokenizer, 'tokenizer') else self.tokenizer trainer_kwargs = { "model": self.model, - "train_dataset": dataset['dataset'], + "train_dataset": train_dataset, + "processing_class": processing_class, + "data_collator": data_collator, + "args": SFTConfig(**config_args), + } + if eval_dataset is not None: + trainer_kwargs["eval_dataset"] = eval_dataset + self.trainer = SFTTrainer(**trainer_kwargs) + elif self.is_vlm: + # Image VLM: dataset is dict wrapper from format_and_template_dataset + train_dataset = dataset['dataset'] if isinstance(dataset, dict) else dataset + trainer_kwargs = { + "model": self.model, + "train_dataset": train_dataset, "processing_class": self.tokenizer, "data_collator": data_collator, "args": SFTConfig(**config_args), @@ -896,7 +2517,8 @@ class UnslothTrainer: train_on_responses_enabled = training_args.get('train_on_completions', False) # DeepSeek OCR handles this internally in its collator, so skip - if train_on_responses_enabled and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'): + # Audio VLM handles label masking in its collator, so skip + if train_on_responses_enabled and not self.is_audio_vlm and not self.is_audio and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'): try: print("Configuring train on responses only...\n") @@ -925,7 +2547,7 @@ class UnslothTrainer: train_on_responses_enabled = False # Apply train on responses only if we have valid parts - if train_on_responses_enabled and instruction_part and response_part and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'): + if train_on_responses_enabled and instruction_part and response_part and not self.is_audio_vlm and not self.is_audio and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'): try: from unsloth.chat_templates import train_on_responses_only @@ -933,7 +2555,7 @@ class UnslothTrainer: self.trainer, instruction_part=instruction_part, response_part=response_part, - num_proc=config_args.get("dataset_num_proc", safe_num_proc(max(1, os.cpu_count() // 4))), + num_proc=config_args["dataset_num_proc"], ) print("Train on responses only configured successfully\n") @@ -980,103 +2602,17 @@ class UnslothTrainer: else: print("Training on full sequences (including prompts)\n") - # Add custom callback for progress tracking - from transformers import TrainerCallback - - class ProgressCallback(TrainerCallback): - def __init__(self, trainer_instance): - self.trainer_instance = trainer_instance - - def on_train_begin(self, args, state, control, **kwargs): - """Called at the beginning of training""" - pass - - def on_log(self, args, state, control, logs=None, **kwargs): - """Called when logging occurs""" - if logs: - # Get loss from either 'loss' or 'train_loss' key - loss_value = logs.get('loss', logs.get('train_loss', 0.0)) - current_step = state.global_step - - # Extract grad_norm from logs (available when gradient clipping is enabled) - grad_norm = logs.get('grad_norm', None) - - # Calculate elapsed_seconds - elapsed_seconds = None - if self.trainer_instance.training_start_time is not None: - elapsed_seconds = time.time() - self.trainer_instance.training_start_time - - # Calculate eta_seconds - eta_seconds = None - if elapsed_seconds is not None and current_step > 0: - total_steps = self.trainer_instance.training_progress.total_steps - if total_steps > 0: - steps_remaining = total_steps - current_step - if steps_remaining > 0: - time_per_step = elapsed_seconds / current_step - eta_seconds = time_per_step * steps_remaining - - # Extract num_tokens from TRL SFTTrainer state (real counter) - # Requires include_num_input_tokens_seen=True in SFTConfig - num_tokens = getattr(state, "num_input_tokens_seen", None) - - self.trainer_instance._update_progress( - step=current_step, - epoch=round(state.epoch, 2) if state.epoch else 0, - loss=loss_value, - learning_rate=logs.get('learning_rate', 0.0), - elapsed_seconds=elapsed_seconds, - eta_seconds=eta_seconds, - grad_norm=grad_norm, - num_tokens=num_tokens, - eval_loss=logs.get('eval_loss', None), - status_message="" - ) - - def on_epoch_end(self, args, state, control, **kwargs): - """Called at the end of each epoch""" - self.trainer_instance._update_progress( - epoch=state.epoch, - step=state.global_step - ) - - def on_step_end(self, args, state, control, **kwargs): - """Called at the end of each step""" - # Check if we should stop training - if self.trainer_instance.should_stop: - print(f"Stop detected at step {state.global_step}\n") - control.should_training_stop = True - return control - # ========== PROGRESS TRACKING ========== - progress_callback = ProgressCallback(self) - self.trainer.add_callback(progress_callback) + self.trainer.add_callback(self._create_progress_callback()) - num_samples = len(self.trainer.train_dataset) + num_samples = len(dataset['dataset'] if isinstance(dataset, dict) else dataset) batch_size = training_args.get('batch_size', 2) - grad_accum = training_args.get('gradient_accumulation_steps', 4) - num_epochs = training_args.get('num_epochs', 3) - max_steps_val = training_args.get('max_steps', 0) - - # Step 1: Calculate dataloader length (number of batches) - len_dataloader = math.ceil(num_samples / batch_size) - - # Step 2: Calculate steps per epoch (following transformers logic) - num_update_steps_per_epoch = max( - len_dataloader // grad_accum + int(len_dataloader % grad_accum > 0), - 1 + total_steps = self._calculate_total_steps( + num_samples, batch_size, + training_args.get('gradient_accumulation_steps', 4), + training_args.get('num_epochs', 3), + training_args.get('max_steps', 0), ) - - # Step 3: Determine total steps based on max_steps or epochs - if max_steps_val and max_steps_val > 0: - # Use max_steps if specified - total_steps = max_steps_val - print(f"Progress tracking: {total_steps} steps (max_steps)\n") - else: - # Calculate from epochs - total_steps = num_update_steps_per_epoch * num_epochs - print(f"Progress tracking: {total_steps} steps ({num_epochs} epochs × {num_update_steps_per_epoch} steps/epoch)\n") - self._update_progress(total_steps=total_steps) # ========== START TRAINING ========== @@ -1085,40 +2621,47 @@ class UnslothTrainer: self.trainer.train() # ========== SAVE MODEL ========== - if self.should_stop and self.save_on_stop: - # Stopped by user — save model at current checkpoint - self.trainer.save_model() - self.tokenizer.save_pretrained(output_dir) - print(f"\nTraining stopped. Model saved to {output_dir}\n") - self._update_progress( - is_training=False, - status_message=f"Training stopped. Model saved to {output_dir}", - ) - elif self.should_stop: - # Cancelled by user — don't save - print("\nTraining cancelled.\n") - self._update_progress( - is_training=False, - status_message="Training cancelled.", - ) - else: - # Normal completion - self.trainer.save_model() - self.tokenizer.save_pretrained(output_dir) - print(f"\nTraining completed! Model saved to {output_dir}\n") - self._update_progress( - is_training=False, - is_completed=True, - status_message=f"Training completed! Model saved to {output_dir}", - ) + self._finalize_training(output_dir) except Exception as e: + import traceback logger.error(f"Training error: {e}") + logger.error(f"Full traceback:\n{traceback.format_exc()}") self._update_progress(is_training=False, error=str(e)) finally: self.is_training = False + def _patch_adapter_config(self, output_dir: str) -> None: + """Patch adapter_config.json with unsloth_training_method. + + Values: 'qlora', 'lora', 'FT', 'CPT', 'DPO', 'GRPO', etc. + For LoRA/QLoRA, the distinction comes from load_in_4bit. + """ + config_path = os.path.join(output_dir, "adapter_config.json") + if not os.path.exists(config_path): + logger.info("No adapter_config.json found — skipping training method patch") + return + + try: + with open(config_path, "r") as f: + config = json.load(f) + + # Determine the training method + if self.load_in_4bit: + method = "qlora" + else: + method = "lora" + + config["unsloth_training_method"] = method + logger.info(f"Patching adapter_config.json with unsloth_training_method='{method}'") + + with open(config_path, "w") as f: + json.dump(config, f, indent=2) + + except Exception as e: + logger.warning(f"Failed to patch adapter_config.json: {e}") + def stop_training(self, save: bool = True): """Stop ongoing training""" print(f"\nStopping training (save={save})...") diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 153f4335e3..5323e73ae1 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -1,620 +1,535 @@ """ -Training backend for FastAPI integration -""" -import matplotlib.pyplot as plt -from typing import Any, Generator, Tuple -import logging -import math +Training backend — subprocess orchestrator. -from .trainer import get_trainer, TrainingProgress -from utils.hardware import clear_gpu_cache +Each training job runs in a fresh subprocess (mp.get_context("spawn")), +solving the transformers version-switching problem. The old in-process +UnslothTrainer singleton is only used inside the subprocess (worker.py). + +This file orchestrates the subprocess lifecycle, pumps events from the +worker's mp.Queue, and exposes the same API surface to routes/training.py. + +Pattern follows core/data_recipe/jobs/manager.py. +""" +import math +import multiprocessing as mp +import queue +import threading +import time +import logging +from dataclasses import dataclass, field +from pathlib import Path +from typing import Optional, Tuple, Any + +import matplotlib.pyplot as plt logger = logging.getLogger(__name__) +_CTX = mp.get_context("spawn") + # Plot styling constants -PLOT_WIDTH = 8 # Inches -PLOT_HEIGHT = 3.5 # Inches +PLOT_WIDTH = 8 +PLOT_HEIGHT = 3.5 + + +@dataclass +class TrainingProgress: + """Mirror of trainer.TrainingProgress — kept here so the parent process + never needs to import the heavy ML modules.""" + epoch: float = 0 + step: int = 0 + total_steps: int = 0 + loss: float = 0.0 + learning_rate: float = 0.0 + is_training: bool = False + is_completed: bool = False + error: Optional[str] = None + status_message: str = "Ready to train" + elapsed_seconds: Optional[float] = None + eta_seconds: Optional[float] = None + grad_norm: Optional[float] = None + num_tokens: Optional[int] = None + eval_loss: Optional[float] = None class TrainingBackend: """ - Training orchestration backend. - Handles both text and vision models, LoRA and full finetuning. + Training orchestration backend — subprocess-based. + Launches a fresh subprocess per training job, communicates via mp.Queue. """ def __init__(self): - self.trainer = get_trainer() + # Subprocess state + self._proc: Optional[mp.Process] = None + self._event_queue: Any = None + self._stop_queue: Any = None + self._pump_thread: Optional[threading.Thread] = None + self._lock = threading.Lock() - # Training Metrics - self.loss_history = [] - self.lr_history = [] - self.step_history = [] - self.grad_norm_history = [] - self.grad_norm_step_history = [] - self.eval_loss_history = [] - self.eval_step_history = [] + # Progress state (updated by pump thread from subprocess events) + self._progress = TrainingProgress() + self._should_stop = False + + # Training Metrics (consumed by routes for SSE and /metrics) + self.loss_history: list = [] + self.lr_history: list = [] + self.step_history: list = [] + self.grad_norm_history: list = [] + self.grad_norm_step_history: list = [] + self.eval_loss_history: list = [] + self.eval_step_history: list = [] + self.eval_enabled: bool = False + self.current_theme: str = "light" + + # Job metadata + self.current_job_id: Optional[str] = None + self._output_dir: Optional[str] = None + + logger.info("TrainingBackend initialized (subprocess mode)") + + # ------------------------------------------------------------------ + # Public API (called by routes/training.py) + # ------------------------------------------------------------------ + + def start_training(self, **kwargs) -> bool: + """Spawn a subprocess to run the full training pipeline. + + All kwargs are serialized into a config dict and sent to the worker. + Returns True if the subprocess was started successfully. + """ + with self._lock: + if self._proc is not None and self._proc.is_alive(): + logger.warning("Training subprocess already running") + return False + + # Join prior pump thread to prevent it from consuming events + # from the new job's queue (it reads self._event_queue dynamically). + if self._pump_thread is not None and self._pump_thread.is_alive(): + self._pump_thread.join(timeout=5.0) + if self._pump_thread.is_alive(): + logger.warning("Previous pump thread did not exit within 5s") + self._pump_thread = None + + # Reset state + self._should_stop = False + self._progress = TrainingProgress(is_training=True, status_message="Initializing training...") + self.loss_history.clear() + self.lr_history.clear() + self.step_history.clear() + self.grad_norm_history.clear() + self.grad_norm_step_history.clear() + self.eval_loss_history.clear() + self.eval_step_history.clear() self.eval_enabled = False - self.current_theme = "light" + self._output_dir = None - self.trainer.add_progress_callback(self._on_progress_update) + # Resolve project root (studio/backend/core/training/ → project root) + project_root = str(Path(__file__).resolve().parent.parent.parent.parent.parent) - logger.info("TrainingBackend initialized") + # Build config dict for the subprocess + config = { + "project_root": project_root, + "model_name": kwargs["model_name"], + "training_type": kwargs.get("training_type", "LoRA/QLoRA"), + "hf_token": kwargs.get("hf_token", ""), + "load_in_4bit": kwargs.get("load_in_4bit", True), + "max_seq_length": kwargs.get("max_seq_length", 2048), + "hf_dataset": kwargs.get("hf_dataset", ""), + "local_datasets": kwargs.get("local_datasets"), + "format_type": kwargs.get("format_type", ""), + "subset": kwargs.get("subset"), + "train_split": kwargs.get("train_split", "train"), + "eval_split": kwargs.get("eval_split"), + "eval_steps": kwargs.get("eval_steps", 0.00), + "dataset_slice_start": kwargs.get("dataset_slice_start"), + "dataset_slice_end": kwargs.get("dataset_slice_end"), + "custom_format_mapping": kwargs.get("custom_format_mapping"), + "is_dataset_image": kwargs.get("is_dataset_image", False), + "is_dataset_audio": kwargs.get("is_dataset_audio", False), + "num_epochs": kwargs.get("num_epochs", 3), + "learning_rate": kwargs.get("learning_rate", "2e-4"), + "batch_size": kwargs.get("batch_size", 2), + "gradient_accumulation_steps": kwargs.get("gradient_accumulation_steps", 4), + "warmup_steps": kwargs.get("warmup_steps"), + "warmup_ratio": kwargs.get("warmup_ratio"), + "max_steps": kwargs.get("max_steps", 0), + "save_steps": kwargs.get("save_steps", 0), + "weight_decay": kwargs.get("weight_decay", 0.01), + "random_seed": kwargs.get("random_seed", 3407), + "packing": kwargs.get("packing", False), + "optim": kwargs.get("optim", "adamw_8bit"), + "lr_scheduler_type": kwargs.get("lr_scheduler_type", "linear"), + "use_lora": kwargs.get("use_lora", True), + "lora_r": kwargs.get("lora_r", 16), + "lora_alpha": kwargs.get("lora_alpha", 16), + "lora_dropout": kwargs.get("lora_dropout", 0.0), + "target_modules": kwargs.get("target_modules"), + "gradient_checkpointing": kwargs.get("gradient_checkpointing", "unsloth"), + "use_rslora": kwargs.get("use_rslora", False), + "use_loftq": kwargs.get("use_loftq", False), + "train_on_completions": kwargs.get("train_on_completions", False), + "finetune_vision_layers": kwargs.get("finetune_vision_layers", True), + "finetune_language_layers": kwargs.get("finetune_language_layers", True), + "finetune_attention_modules": kwargs.get("finetune_attention_modules", True), + "finetune_mlp_modules": kwargs.get("finetune_mlp_modules", True), + "enable_wandb": kwargs.get("enable_wandb", False), + "wandb_token": kwargs.get("wandb_token"), + "wandb_project": kwargs.get("wandb_project", "unsloth-training"), + "enable_tensorboard": kwargs.get("enable_tensorboard", False), + "tensorboard_dir": kwargs.get("tensorboard_dir", "runs"), + } - def _on_progress_update(self, progress: TrainingProgress): - """Callback for progress updates""" - if progress.step >= 0 and progress.loss > 0: - self.loss_history.append(progress.loss) - self.lr_history.append(progress.learning_rate) - self.step_history.append(progress.step) - if progress.step >= 0 and progress.grad_norm is not None: - try: - grad_norm = float(progress.grad_norm) - except (TypeError, ValueError): - grad_norm = None - if grad_norm is not None and math.isfinite(grad_norm): - self.grad_norm_history.append(grad_norm) - self.grad_norm_step_history.append(progress.step) - if progress.eval_loss is not None: - self.eval_loss_history.append(progress.eval_loss) - self.eval_step_history.append(progress.step) + # Derive load_in_4bit from training_type + if config["training_type"] != "LoRA/QLoRA": + config["load_in_4bit"] = False - def start_training(self, - # Model parameters - model_name: str, - training_type: str, # NEW: "LoRA/QLoRA" or "Full Finetuning" - hf_token: str, - load_in_4bit: bool, - max_seq_length: int, + # Spawn subprocess + from .worker import run_training_process - # Dataset parameters - hf_dataset: str, - local_datasets: list, - format_type: str, # CHANGED: was data_template + self._event_queue = _CTX.Queue() + self._stop_queue = _CTX.Queue() - # Training parameters - num_epochs: int, - learning_rate: str, - batch_size: int, - gradient_accumulation_steps: int, - warmup_steps: int, # May be None even without default - warmup_ratio: float, # May be None even without default - max_steps: int, - save_steps: int, - weight_decay: float, - random_seed: int, - packing: bool, - optim: str, - lr_scheduler_type: str, + self._proc = _CTX.Process( + target=run_training_process, + kwargs={ + "event_queue": self._event_queue, + "stop_queue": self._stop_queue, + "config": config, + }, + daemon=True, + ) + self._proc.start() + logger.info("Training subprocess started (pid=%s)", self._proc.pid) - # LoRA parameters - use_lora: bool, # Should be derived from training_type - lora_r: int, - lora_alpha: int, - lora_dropout: float, - target_modules: list, - gradient_checkpointing: str, - use_rslora: bool, - use_loftq: bool, - train_on_completions: bool, + # Start event pump thread + self._pump_thread = threading.Thread(target=self._pump_loop, daemon=True) + self._pump_thread.start() - # NEW: Vision-specific LoRA parameters - finetune_vision_layers: bool, - finetune_language_layers: bool, - finetune_attention_modules: bool, - finetune_mlp_modules: bool, - - # Logging parameters - enable_wandb: bool, - wandb_token: str, - wandb_project: str, - enable_tensorboard: bool, - tensorboard_dir: str, - - # Optional parameters - custom_format_mapping: dict = None, - subset: str = None, - train_split: str = "train", - eval_split: str = None, - eval_steps: float = 0.00, - is_dataset_multimodal: bool = False, - dataset_slice_start: int = None, - dataset_slice_end: int = None) -> bool: - """ - Start training. - - Returns: - True if training started successfully, False otherwise. - """ - try: - # Wait for any previous training thread to finish - old_thread = getattr(self.trainer, "training_thread", None) - if old_thread and old_thread.is_alive(): - logger.info("Waiting for previous training thread to finish...") - old_thread.join(timeout=30) - - # Explicitly free old SFTTrainer and CUDA resources before loading new model. - # Without this, forked multiprocessing workers (num_proc tokenization) inherit - # stale CUDA state from the previous run, causing extreme slowdowns or crashes. - if self.trainer.trainer is not None: - logger.info("Cleaning up previous SFTTrainer...") - self.trainer.trainer = None - if self.trainer.model is not None: - self.trainer.model = None - if self.trainer.tokenizer is not None: - self.trainer.tokenizer = None - # Flush all pending async CUDA ops so forked tokenization processes - # don't inherit stale async state that causes pool join to hang. - import torch as _torch - if _torch.cuda.is_available(): - _torch.cuda.synchronize() - import gc - gc.collect() - clear_gpu_cache() - - # Reset stop flag and clear history - self.trainer.should_stop = False - self.trainer.save_on_stop = True - self.loss_history = [] - self.lr_history = [] - self.step_history = [] - self.grad_norm_history = [] - self.grad_norm_step_history = [] - self.eval_loss_history = [] - self.eval_step_history = [] - self.eval_enabled = False - import time - output_dir = f"./outputs/{model_name.replace('/', '_')}_{int(time.time())}" - - # Derive use_lora from training_type - use_lora_actual = (training_type == "LoRA/QLoRA") - if use_lora_actual: print("using Lora") - else: print("using full finetuning") - logger.info(f"Starting training - Type: {training_type}, Model: {model_name}") - - # ========== LOAD MODEL ========== - logger.info("Loading model...") - success = self.trainer.load_model( - model_name=model_name, - max_seq_length=max_seq_length, - load_in_4bit=load_in_4bit if use_lora_actual else False, # Only 4bit for LoRA - hf_token=hf_token if hf_token.strip() else None, - is_dataset_multimodal=is_dataset_multimodal, - ) - - if not success or self.trainer.should_stop: - logger.error("Failed to load model or stopped by user") - return False - - # ========== PREPARE MODEL FOR TRAINING ========== - if use_lora_actual: - logger.info("Preparing model with LoRA...") - success = self.trainer.prepare_model_for_training( - use_lora=True, - # Vision-specific parameters - finetune_vision_layers=finetune_vision_layers, - finetune_language_layers=finetune_language_layers, - finetune_attention_modules=finetune_attention_modules, - finetune_mlp_modules=finetune_mlp_modules, - # Standard LoRA parameters - target_modules=target_modules, - lora_r=lora_r, - lora_alpha=lora_alpha, - lora_dropout=lora_dropout, - use_gradient_checkpointing=gradient_checkpointing, - use_rslora=use_rslora, - use_loftq=use_loftq - ) - else: - logger.info("Preparing model for full finetuning...") - success = self.trainer.prepare_model_for_training( - use_lora=False # Full finetuning - ) - - if not success or self.trainer.should_stop: - logger.error("Failed to prepare model or stopped by user") - return False - - # ========== LOAD DATASET ========== - logger.info("Loading dataset...") - #breakpoint() - dataset_result = self.trainer.load_and_format_dataset( - dataset_source=hf_dataset if hf_dataset.strip() else None, - format_type=format_type, - local_datasets=local_datasets if local_datasets else None, - custom_format_mapping=custom_format_mapping, - subset=subset, - train_split=train_split, - eval_split=eval_split, - eval_steps=eval_steps, - dataset_slice_start=dataset_slice_start, - dataset_slice_end=dataset_slice_end, - ) - - # Unpack: load_and_format_dataset returns (dataset, eval_dataset) - if isinstance(dataset_result, tuple): - dataset, eval_dataset = dataset_result - else: - dataset = dataset_result - eval_dataset = None - - # Track whether eval is enabled for status reporting - self.eval_enabled = eval_dataset is not None - - if dataset is None or self.trainer.should_stop: - logger.error("Failed to load dataset or stopped by user") - return False - - # ========== START TRAINING ========== - # Convert learning rate string to float - try: - lr_value = float(learning_rate) - except ValueError: - logger.error(f"Invalid learning rate: {learning_rate}") - self.trainer._update_progress( - error=f"Invalid learning rate: {learning_rate}", - is_training=False - ) - return - - logger.info("Starting training worker thread...") - success = self.trainer.start_training( - dataset=dataset, - eval_dataset=eval_dataset, - eval_steps=eval_steps, - output_dir=output_dir, - num_epochs=num_epochs, - learning_rate=lr_value, - batch_size=batch_size, - gradient_accumulation_steps=gradient_accumulation_steps, - warmup_steps=warmup_steps, - warmup_ratio=warmup_ratio, - max_steps=max_steps if max_steps > 0 else 0, - save_steps=save_steps if save_steps > 0 else 0, - weight_decay=weight_decay, - random_seed=random_seed, - packing=packing, - train_on_completions=train_on_completions, - enable_wandb=enable_wandb, - wandb_project=wandb_project, - wandb_token=wandb_token if wandb_token.strip() else None, - enable_tensorboard=enable_tensorboard, - tensorboard_dir=tensorboard_dir, - max_seq_length=max_seq_length, - optim=optim, - lr_scheduler_type=lr_scheduler_type, - ) - - if not success: - logger.error("Failed to start training") - return False - - return True - - except Exception as e: - logger.error(f"Error in start_training: {e}", exc_info=True) - self.trainer._update_progress( - error=str(e), - is_training=False - ) - return False + return True def stop_training(self, save: bool = True) -> bool: - """ - Stop ongoing training. + """Send stop signal to the training subprocess.""" + self._should_stop = True + with self._lock: + if self._stop_queue is not None: + try: + self._stop_queue.put({"type": "stop", "save": save}) + except (OSError, ValueError): + pass + # Update progress immediately for responsive UI + self._progress.status_message = ( + "Stopping training and saving checkpoint..." + if save else "Cancelling training..." + ) + return True - Args: - save: If True, save the model at the current checkpoint. + def is_training_active(self) -> bool: + """Check if training is currently active.""" + with self._lock: + # Subprocess alive = active + if self._proc is not None and self._proc.is_alive(): + return True + + # Stop was requested and process exited → inactive + if self._should_stop: + return False + + # Check progress state + p = self._progress + if p.is_training: + return True + if p.is_completed or p.error: + return False + + # Check status message for activity indicators + status_lower = (p.status_message or "").lower() + if any(k in status_lower for k in ["cancelled", "canceled", "stopped", "completed", "ready to train"]): + return False + if any(k in status_lower for k in ["loading", "preparing", "training", "configuring", "tokenizing", "starting", "importing"]): + return True - Returns: - True if training was successfully stopped. - """ - try: - logger.info(f"Stopping training (save={save})...") - self.trainer.stop_training(save=save) - return True - except Exception as e: - logger.error(f"Error stopping training: {e}") return False def get_training_status(self, theme: str = "light") -> Tuple: - """ - Get current training status and loss plot. + """Get current training status and loss plot.""" + with self._lock: + progress = self._progress - Args: - theme: "light" or "dark" for plot styling + if not (progress.is_training or progress.is_completed or progress.error): + return (None, progress) - Returns: - Tuple of (plot, progress) - """ + plot = self._create_loss_plot(progress, theme) + return (plot, progress) - try: - progress = self.trainer.get_training_progress() - - # If not training and not completed, return no updates - if not (progress.is_training or progress.is_completed or progress.error): - return (None, progress) - - # Generate plot - plot = self._create_loss_plot(progress, theme) - return (plot, progress) - - except Exception as e: - logger.error(f"Error getting training status: {e}") - return (None, None) - - def refresh_plot_for_theme(self, theme: str) -> plt.Figure: - """ - Refresh plot with new theme. - - Args: - theme: "light" or "dark" - - Returns: - Updated matplotlib figure - """ + def refresh_plot_for_theme(self, theme: str) -> Optional[plt.Figure]: + """Refresh plot with new theme.""" if theme and isinstance(theme, str) and theme in ['light', 'dark']: self.current_theme = theme - - # Always generate plot if we have loss history if self.loss_history: - progress = self.trainer.get_training_progress() + with self._lock: + progress = self._progress return self._create_loss_plot(progress, self.current_theme) - return None - def is_training_active(self) -> bool: - """ - Check if training is currently active (from load_model start to completion/error). - - Returns: - True if training is in progress, False otherwise - """ + # ------------------------------------------------------------------ + # Compatibility shims — routes/training.py accesses these + # ------------------------------------------------------------------ + + class _TrainerShim: + """Minimal shim so routes that access backend.trainer.* still work.""" + def __init__(self, backend: "TrainingBackend"): + self._backend = backend + self.should_stop = False + + @property + def training_progress(self): + return self._backend._progress + + @training_progress.setter + def training_progress(self, value): + self._backend._progress = value + + def get_training_progress(self): + return self._backend._progress + + def _update_progress(self, **kwargs): + with self._backend._lock: + for key, value in kwargs.items(): + if hasattr(self._backend._progress, key): + setattr(self._backend._progress, key, value) + + @property + def trainer(self): + """Compatibility shim for routes that access backend.trainer.*""" + return self._TrainerShim(self) + + # ------------------------------------------------------------------ + # Event pump (background thread) + # ------------------------------------------------------------------ + + def _pump_loop(self) -> None: + """Background thread: consume events from subprocess → update state.""" + while True: + if self._proc is None or self._event_queue is None: + return + + # Try to read an event + event = self._read_queue(self._event_queue, timeout_sec=0.25) + if event is not None: + self._handle_event(event) + continue + + # No event — check if process is still alive + if self._proc.is_alive(): + continue + + # Process exited — drain remaining events + for e in self._drain_queue(self._event_queue): + self._handle_event(e) + + # Mark as done if no explicit complete/error was received + with self._lock: + if self._progress.is_training: + if self._should_stop: + self._progress.is_training = False + self._progress.status_message = "Training stopped." + else: + self._progress.is_training = False + self._progress.error = self._progress.error or "Training process exited unexpectedly" + return + + def _handle_event(self, event: dict) -> None: + """Apply a subprocess event to local state.""" + etype = event.get("type") + + with self._lock: + if etype == "progress": + self._progress.step = event.get("step", self._progress.step) + self._progress.epoch = event.get("epoch", self._progress.epoch) + self._progress.loss = event.get("loss", self._progress.loss) + self._progress.learning_rate = event.get("learning_rate", self._progress.learning_rate) + self._progress.total_steps = event.get("total_steps", self._progress.total_steps) + self._progress.elapsed_seconds = event.get("elapsed_seconds") + self._progress.eta_seconds = event.get("eta_seconds") + self._progress.grad_norm = event.get("grad_norm") + self._progress.num_tokens = event.get("num_tokens") + self._progress.eval_loss = event.get("eval_loss") + self._progress.is_training = True + status = event.get("status_message", "") + if status: + self._progress.status_message = status + + # Update metric histories + step = event.get("step", 0) + loss = event.get("loss", 0.0) + lr = event.get("learning_rate", 0.0) + if step >= 0 and loss > 0: + self.loss_history.append(loss) + self.lr_history.append(lr) + self.step_history.append(step) + + grad_norm = event.get("grad_norm") + if grad_norm is not None: + try: + gn = float(grad_norm) + except (TypeError, ValueError): + gn = None + if gn is not None and math.isfinite(gn): + self.grad_norm_history.append(gn) + self.grad_norm_step_history.append(step) + + eval_loss = event.get("eval_loss") + if eval_loss is not None: + self.eval_loss_history.append(eval_loss) + self.eval_step_history.append(step) + self.eval_enabled = True + + elif etype == "status": + self._progress.status_message = event.get("message", "") + self._progress.is_training = True + + elif etype == "complete": + self._progress.is_training = False + self._progress.is_completed = True + self._output_dir = event.get("output_dir") + msg = event.get("status_message", "Training completed") + self._progress.status_message = msg + + elif etype == "error": + self._progress.is_training = False + self._progress.error = event.get("error", "Unknown error") + logger.error("Training error: %s", event.get("error")) + stack = event.get("stack", "") + if stack: + logger.error("Stack trace:\n%s", stack) + + @staticmethod + def _read_queue(q: Any, timeout_sec: float) -> Optional[dict]: try: - training_thread = getattr(self.trainer, "training_thread", None) - if training_thread and training_thread.is_alive(): - return True + return q.get(timeout=timeout_sec) + except queue.Empty: + return None + except (EOFError, OSError, ValueError): + return None - # Stop requested and worker already exited => inactive. - # This allows UI to show stopped state + "Back to configuration". - if getattr(self.trainer, "should_stop", False): - return False + @staticmethod + def _drain_queue(q: Any) -> list: + events = [] + while True: + try: + events.append(q.get_nowait()) + except queue.Empty: + return events + except (EOFError, OSError, ValueError): + return events - progress = self.trainer.get_training_progress() - # Training is active if is_training is True - # Also check if we're in loading/preparation phase (status_message indicates activity) - is_active = progress.is_training - # Also consider it active if we have a status message indicating loading/preparation - # but haven't completed or errored yet - if not is_active and not progress.is_completed and not progress.error: - status = progress.status_message or "" - status_lower = status.lower() - if any( - keyword in status_lower - for keyword in ["cancelled", "canceled", "stopped", "completed", "ready to train"] - ): - return False - if any( - keyword in status_lower - for keyword in [ - "loading", - "preparing", - "training", - "configuring", - "tokenizing", - "starting", - ] - ): - is_active = True - return is_active - except Exception as e: - logger.error(f"Error checking training state: {e}") - return False + # ------------------------------------------------------------------ + # Plot generation (unchanged from original) + # ------------------------------------------------------------------ def _create_loss_plot(self, progress: TrainingProgress, theme: str = "light") -> plt.Figure: - """ - Create training loss plot with theme-aware styling. + """Create training loss plot with theme-aware styling.""" + plt.close('all') - Args: - progress: Current training progress - theme: "light" or "dark" + LIGHT_STYLE = { + "facecolor": "#ffffff", + "grid_color": "#d1d5db", + "line": "#16b88a", + "text": "#1f2937", + "empty_text": "#6b7280" + } + DARK_STYLE = { + "facecolor": "#292929", + "grid_color": "#404040", + "line": "#4ade80", + "text": "#e5e7eb", + "empty_text": "#9ca3af" + } - Returns: - Matplotlib figure - """ - plt.close('all') + style = LIGHT_STYLE if theme == "light" else DARK_STYLE - # Theme-specific styling - LIGHT_STYLE = { - "facecolor": "#ffffff", - "grid_color": "#d1d5db", - "line": "#16b88a", - "text": "#1f2937", - "empty_text": "#6b7280" - } - DARK_STYLE = { - "facecolor": "#292929", - "grid_color": "#404040", - "line": "#4ade80", - "text": "#e5e7eb", - "empty_text": "#9ca3af" - } + fig, ax = plt.subplots(figsize=(PLOT_WIDTH, PLOT_HEIGHT)) + fig.patch.set_facecolor(style["facecolor"]) + ax.set_facecolor(style["facecolor"]) - style = LIGHT_STYLE if theme == "light" else DARK_STYLE + if self.loss_history: + steps = self.step_history + losses = self.loss_history + scatter_color = "#60a5fa" + ax.scatter(steps, losses, s=16, alpha=0.6, color=scatter_color, + linewidths=0, label="Training Loss (raw)") - fig, ax = plt.subplots(figsize=(PLOT_WIDTH, PLOT_HEIGHT)) - fig.patch.set_facecolor(style["facecolor"]) - ax.set_facecolor(style["facecolor"]) + MA_WINDOW = 20 + window = min(MA_WINDOW, len(losses)) - if self.loss_history: - steps = self.step_history - losses = self.loss_history - scatter_color = "#60a5fa" - # Scatter plot for raw loss points - ax.scatter( - steps, - losses, - s=16, - alpha=0.6, - color=scatter_color, - linewidths=0, - label="Training Loss (raw)", - ) + if window >= 2: + cumsum = [0.0] + for v in losses: + cumsum.append(cumsum[-1] + float(v)) - # Moving average line overlay (trailing window) - MA_WINDOW = 20 # adjust smoothing aggressiveness - window = min(MA_WINDOW, len(losses)) + ma = [] + for i in range(len(losses)): + start = max(0, i - window + 1) + denom = i - start + 1 + ma.append((cumsum[i + 1] - cumsum[start]) / denom) - if window >= 2: - cumsum = [0.0] - for v in losses: - cumsum.append(cumsum[-1] + float(v)) + ax.plot(steps, ma, color=style["line"], linewidth=2.5, alpha=0.95, + label=f"Moving Avg ({ma[-1]:.4f})") - ma = [] - for i in range(len(losses)): - start = max(0, i - window + 1) - denom = i - start + 1 - ma.append((cumsum[i + 1] - cumsum[start]) / denom) + leg = ax.legend(frameon=False, fontsize=9) + for t in leg.get_texts(): + t.set_color(style["text"]) - ax.plot( - steps, - ma, - color=style["line"], - linewidth=2.5, - alpha=0.95, - label=f"Moving Avg ({ma[-1]:.4f})", - ) + ax.set_xlabel('Steps', fontsize=10, color=style["text"]) + ax.set_ylabel('Loss', fontsize=10, color=style["text"]) - leg = ax.legend(frameon=False, fontsize=9) - for t in leg.get_texts(): - t.set_color(style["text"]) - - ax.set_xlabel('Steps', fontsize=10, color=style["text"]) - ax.set_ylabel('Loss', fontsize=10, color=style["text"]) - - # Build status message for title - if progress.error: - title = f"Error: {progress.error}" - elif progress.is_completed: - title = f"Training completed! Final loss: {progress.loss:.4f}" - elif progress.status_message: - title = progress.status_message - elif progress.step > 0: - title = f"Epoch: {progress.epoch} | Step: {progress.step}/{progress.total_steps} | Loss: {progress.loss:.4f}" - else: - title = "Training Loss" - - ax.set_title(title, fontsize=11, fontweight='bold', - pad=10, color=style["text"]) - - # Style grid and spines - ax.grid(True, alpha=0.4, linestyle='--', color=style["grid_color"]) - ax.tick_params(colors=style["text"], which='both') - ax.spines['top'].set_visible(False) - ax.spines['right'].set_visible(False) - ax.spines['bottom'].set_color(style["text"]) - ax.spines['left'].set_color(style["text"]) + if progress.error: + title = f"Error: {progress.error}" + elif progress.is_completed: + title = f"Training completed! Final loss: {progress.loss:.4f}" + elif progress.status_message: + title = progress.status_message + elif progress.step > 0: + title = f"Epoch: {progress.epoch} | Step: {progress.step}/{progress.total_steps} | Loss: {progress.loss:.4f}" else: - display_msg = progress.status_message if progress.status_message else 'Waiting for training data...' - ax.text(0.5, 0.5, display_msg, - ha='center', va='center', fontsize=16, - color=style["empty_text"], - transform=ax.transAxes) - ax.set_xticks([]) - ax.set_yticks([]) - for spine in ax.spines.values(): - spine.set_visible(False) + title = "Training Loss" - fig.tight_layout() - return fig + ax.set_title(title, fontsize=11, fontweight='bold', pad=10, color=style["text"]) + ax.grid(True, alpha=0.4, linestyle='--', color=style["grid_color"]) + ax.tick_params(colors=style["text"], which='both') + ax.spines['top'].set_visible(False) + ax.spines['right'].set_visible(False) + ax.spines['bottom'].set_color(style["text"]) + ax.spines['left'].set_color(style["text"]) + else: + display_msg = progress.status_message if progress.status_message else 'Waiting for training data...' + ax.text(0.5, 0.5, display_msg, ha='center', va='center', fontsize=16, + color=style["empty_text"], transform=ax.transAxes) + ax.set_xticks([]) + ax.set_yticks([]) + for spine in ax.spines.values(): + spine.set_visible(False) + + fig.tight_layout() + return fig def _transfer_to_inference_backend(self) -> bool: + """Transfer model to inference backend. + + With subprocess-based training, the model lives in the subprocess + and is freed when it exits. Inference must load from the saved + checkpoint on disk. This is a no-op placeholder. """ - Transfer the trained model to InferenceBackend. - Called automatically when training completes. - """ - print("=" * 60) - print("DEBUG: _transfer_to_inference_backend() CALLED") - print("=" * 60) - - try: - from ..inference import get_inference_backend - - session = self.current_training_session - - # Check if already transferred - if session.get('transferred', False): - print("DEBUG: Already transferred, returning True") - logger.info("Model already transferred, skipping") - return True - - # Validate session data - if not session.get('base_model_name') or not session.get('output_dir'): - logger.warning("Training session incomplete, cannot transfer") - logger.warning(f"Session data: {session}") - return False - - inference_backend = get_inference_backend() - - base_model_name = session['base_model_name'] - output_dir = session['output_dir'] - is_lora = session['is_lora'] - is_vlm = session['is_vlm'] - - logger.info(f"=" * 60) - logger.info(f"TRANSFERRING MODEL TO INFERENCE BACKEND") - logger.info(f"=" * 60) - logger.info(f" Base model: {base_model_name}") - logger.info(f" Output dir: {output_dir}") - logger.info(f" Is LoRA: {is_lora}") - logger.info(f" Is VLM: {is_vlm}") - - # Transfer the model object directly from trainer memory. - # If is_lora is True, self.trainer.model is a PeftModel (Base + Adapter). - # If is_lora is False, it is the finetuned Base Model. - inference_backend.models[base_model_name] = { - "model": self.trainer.model, - "tokenizer": self.trainer.tokenizer, - "is_vision": is_vlm, - "is_lora": is_lora, - "model_path": base_model_name, - "base_model": None, - "loaded_adapters": {}, - # Unsloth/PEFT training keeps the active adapter named 'default' in memory - "active_adapter": "default" if is_lora else None, - } - - # For vision models, also transfer processor - if is_vlm: - if hasattr(self.trainer, 'tokenizer'): - inference_backend.models[base_model_name]["processor"] = self.trainer.tokenizer - logger.info(" Transferred processor for vision model") - - # Load chat template info - inference_backend._load_chat_template_info(base_model_name) - - # If it was LoRA, register the output path. - # This ensures the Eval UI dropdown (which lists files) knows that - # the model currently in memory corresponds to this specific output directory. - if is_lora: - inference_backend.models[base_model_name]["last_trained_adapter"] = output_dir - logger.info(f"Marked trained LoRA adapter: {output_dir}") - - # Set as active model - inference_backend.active_model_name = base_model_name - logger.info(f"Set active model: {base_model_name}") - - return True - - except Exception as e: - logger.error(f"Error transferring model to inference backend: {e}") - import traceback - traceback.print_exc() - return False + logger.info( + "_transfer_to_inference_backend: subprocess training — " + "model must be loaded from disk (output_dir=%s)", self._output_dir + ) + return False # ========== GLOBAL INSTANCE ========== _training_backend = None + def get_training_backend() -> TrainingBackend: """Get global training backend instance""" global _training_backend diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py new file mode 100644 index 0000000000..4c82347936 --- /dev/null +++ b/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(), + }) diff --git a/studio/backend/main.py b/studio/backend/main.py index 331962fb67..696d8a4d10 100644 --- a/studio/backend/main.py +++ b/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"]) diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py index e615180b7f..19197d99ad 100644 --- a/studio/backend/models/datasets.py +++ b/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 diff --git a/studio/backend/models/inference.py b/studio/backend/models/inference.py index 3a908dcfc3..0eb7a7edaf 100644 --- a/studio/backend/models/inference.py +++ b/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=( diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index 8c7d0c037d..8d5a2bfc0a 100644 --- a/studio/backend/models/models.py +++ b/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") diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index b6b30989bd..5cc8141bac 100644 --- a/studio/backend/models/training.py +++ b/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") diff --git a/studio/backend/requirements/extras-no-deps.txt b/studio/backend/requirements/extras-no-deps.txt index b78b479e51..29bd421a0a 100644 --- a/studio/backend/requirements/extras-no-deps.txt +++ b/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 diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 19d70d654f..b7aa72727a 100644 --- a/studio/backend/routes/datasets.py +++ b/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, diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 6616c9fbd8..7f18f50eaa 100644 --- a/studio/backend/routes/export.py +++ b/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, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 512eb05526..52e5bfad72 100644 --- a/studio/backend/routes/inference.py +++ b/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} diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 2c7f0e846f..1590947fdd 100644 --- a/studio/backend/routes/models.py +++ b/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, ) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 497daaedd3..d6726d69df 100644 --- a/studio/backend/routes/training.py +++ b/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( diff --git a/studio/backend/utils/datasets/__init__.py b/studio/backend/utils/datasets/__init__.py index b47db737c3..2e78057237 100644 --- a/studio/backend/utils/datasets/__init__.py +++ b/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 diff --git a/studio/backend/utils/datasets/data_collators.py b/studio/backend/utils/datasets/data_collators.py index f453eaea1b..41062f6a6f 100644 --- a/studio/backend/utils/datasets/data_collators.py +++ b/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: """ diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 9e1f54a75c..9c78d1a49a 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/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 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, diff --git a/studio/backend/utils/datasets/format_detection.py b/studio/backend/utils/datasets/format_detection.py index 9283ea5d55..2337833bef 100644 --- a/studio/backend/utils/datasets/format_detection.py +++ b/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: diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 92e65cf67c..11d5f54539 100644 --- a/studio/backend/utils/models/__init__.py +++ b/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', diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 5404a198aa..f301e42a78 100644 --- a/studio/backend/utils/models/model_config.py +++ b/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: '' 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(' 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 diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py new file mode 100644 index 0000000000..8efbab46a2 --- /dev/null +++ b/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) diff --git a/studio/frontend/package.json b/studio/frontend/package.json index 63df4cc07d..9a607f1c1f 100644 --- a/studio/frontend/package.json +++ b/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", diff --git a/studio/frontend/src/components/assistant-ui/audio-player.tsx b/studio/frontend/src/components/assistant-ui/audio-player.tsx new file mode 100644 index 0000000000..8c5d19abae --- /dev/null +++ b/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 = ({ src }) => { + const audioRef = useRef(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) => { + 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 ( +
+
+ ); +}; diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx index a3f07ab885..252faea6c3 100644 --- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx +++ b/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 ; } +const AUDIO_PLAYER_RE = //; const MarkdownTextImpl = () => { const { text } = useMessagePartText(); const status = useSmoothStatus(); + const audioMatch = text.match(AUDIO_PLAYER_RE); + if (audioMatch) { + return ; + } + return (
= ({ 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 ( +
+
+ + {audioName} + +
+
+ ); +}; + const Composer: FC = () => { return ( + { ); }; +const ComposerAudioUpload: FC = () => { + const audioInputRef = useRef(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 ( + <> + { + const file = e.target.files?.[0]; + if (file) handleAudioFile(file); + e.target.value = ""; + }} + /> + audioInputRef.current?.click()} + aria-label="Upload audio" + > + + + + ); +}; + const ComposerAction: FC = () => { return (
- +
+ + +
@@ -342,6 +424,19 @@ const AssistantActionBar: FC = () => { ); }; +const UserMessageAudio: FC = () => { + const audioName = useAuiState(({ message }) => sentAudioNames.get(message.id)); + if (!audioName) return null; + return ( +
+
+ + {audioName} +
+
+ ); +}; + const UserMessage: FC = () => { return ( { data-role="user" > +
diff --git a/studio/frontend/src/components/ui/chart.tsx b/studio/frontend/src/components/ui/chart.tsx index 071148982f..427cede136 100644 --- a/studio/frontend/src/components/ui/chart.tsx +++ b/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(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 (
- - {children} - + {containerSize ? ( + + {children} + + ) : null}
); @@ -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> & - React.ComponentProps<"div"> & { - hideLabel?: boolean; - hideIndicator?: boolean; - indicator?: "line" | "dot" | "dashed"; - nameKey?: string; - labelKey?: string; - }) { + formatter, + color, + nameKey, + labelKey, +}: Partial> & + 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 & { - hideIcon?: boolean; - nameKey?: string; - }) { - const { config } = useChart(); +const ChartLegend = RechartsPrimitive.Legend; + +function ChartLegendContent({ + className, + hideIcon = false, + payload, + verticalAlign = "bottom", + nameKey, +}: React.ComponentProps<"div"> & + Pick & { + hideIcon?: boolean; + nameKey?: string; + }) { + const { config } = useChart(); if (!payload?.length) { return null; diff --git a/studio/frontend/src/components/ui/slider.tsx b/studio/frontend/src/components/ui/slider.tsx index 4a235e674d..5d75b88e6d 100644 --- a/studio/frontend/src/components/ui/slider.tsx +++ b/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) { - const _values = React.useMemo( - () => - Array.isArray(value) - ? value - : Array.isArray(defaultValue) - ? defaultValue - : [min, max], - [value, defaultValue, min, max], - ); - - return ( - - - - - {Array.from({ length: _values.length }, (_, index) => ( - - ))} - - ); -} - -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) { + const isControlled = Array.isArray(value); + const [uncontrolledValues, setUncontrolledValues] = + React.useState(() => + 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 ( + + + + + {isSingleThumbHorizontal && ( +
+ )} + {Array.from({ length: values.length }, (_, index) => ( + + ))} + + ); +} + +export { Slider }; diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index af8f10156f..8abe71fd62 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/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[0]["messages"]; type RunMessage = RunMessages[number]; +/** Tracks which user messages were sent with an audio file (messageId → filename). */ +export const sentAudioNames = new Map(); + function collectTextParts(message: RunMessage): string[] { const textParts = message.content .filter((part) => part.type === "text") @@ -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 { @@ -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: ``, + }, + ], + }; + } 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, diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 5d5a9551ef..3ac20221bc 100644 --- a/studio/frontend/src/features/chat/api/chat-api.ts +++ b/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 { - 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 { + 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; +} diff --git a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts b/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts index fece047cd8..eb59d5a23d 100644 --- a/studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts +++ b/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), }; } diff --git a/studio/frontend/src/features/chat/runtime-provider.tsx b/studio/frontend/src/features/chat/runtime-provider.tsx index ce553dcdbf..46e3c697ad 100644 --- a/studio/frontend/src/features/chat/runtime-provider.tsx +++ b/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({ diff --git a/studio/frontend/src/features/chat/shared-composer.tsx b/studio/frontend/src/features/chat/shared-composer.tsx index 6b3fc29d9e..fb4eb8fa69 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/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([]); + const [pendingAudio, setPendingAudio] = useState<{ name: string; base64: string } | null>(null); const [dragging, setDragging] = useState(false); const textareaRef = useRef(null); const fileInputRef = useRef(null); + const audioInputRef = useRef(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 (
- {pendingImages.length > 0 && ( + {(pendingImages.length > 0 || pendingAudio) && (
{pendingImages.map(({ id, file }) => ( removePendingImage(id)} /> ))} + {pendingAudio && ( +
+ + {pendingAudio.name} + +
+ )}
)}