diff --git a/install_python_stack.py b/install_python_stack.py index c9b73034b4..77f1b1cea5 100644 --- a/install_python_stack.py +++ b/install_python_stack.py @@ -171,6 +171,13 @@ def install_python_stack() -> int: req=REQ_ROOT / "extras.txt", ) + # 3b. Extra dependencies (no-deps) — audio model support etc. + pip_install( + "Installing extras (no-deps)", + "--no-deps", "--no-cache-dir", + req=REQ_ROOT / "extras-no-deps.txt", + ) + # 4. Overrides (torchao, transformers) — force-reinstall pip_install( "Installing torchao + transformers overrides", @@ -234,8 +241,11 @@ def install_python_stack() -> int: [sys.executable, str(SINGLE_ENV / "patch_metadata.py")], ) - # 12. Final check - run("Running pip check", [sys.executable, "-m", "pip", "check"], quiet=False) + # 12. Final check (silent — third-party conflicts are expected) + subprocess.run( + [sys.executable, "-m", "pip", "check"], + stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, + ) print(_green("✅ Python dependencies installed")) return 0 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/export/export.py b/studio/backend/core/export/export.py index c0f53f2e0e..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 ) 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 index 5db7e56357..b08ec5bb70 100644 --- a/studio/backend/core/inference/orchestrator.py +++ b/studio/backend/core/inference/orchestrator.py @@ -312,6 +312,9 @@ class InferenceOrchestrator: "is_vision": model_info.get("is_vision", False), "is_lora": model_info.get("is_lora", False), "display_name": model_info.get("display_name", model_name), + "is_audio": model_info.get("is_audio", False), + "audio_type": model_info.get("audio_type"), + "has_audio_input": model_info.get("has_audio_input", False), } self.loading_models.discard(model_name) logger.info("Model '%s' loaded successfully in subprocess", model_name) @@ -545,6 +548,203 @@ class InferenceOrchestrator: except RuntimeError: pass + # ------------------------------------------------------------------ + # Audio generation — TTS, ASR, audio input + # ------------------------------------------------------------------ + + def generate_audio_response( + self, + text: str, + temperature: float = 0.6, + top_p: float = 0.95, + top_k: int = 50, + min_p: float = 0.0, + max_new_tokens: int = 2048, + repetition_penalty: float = 1.1, + use_adapter: Optional[Union[bool, str]] = None, + ) -> Tuple[bytes, int]: + """Generate TTS audio. Returns (wav_bytes, sample_rate). + + Blocking — sends command and waits for the complete audio response. + """ + if not self._ensure_subprocess_alive(): + raise RuntimeError("Inference subprocess is not running") + if not self.active_model_name: + raise RuntimeError("No active model") + + import uuid + request_id = str(uuid.uuid4()) + + cmd = { + "type": "generate_audio", + "request_id": request_id, + "text": text, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "max_new_tokens": max_new_tokens, + "repetition_penalty": repetition_penalty, + } + if use_adapter is not None: + cmd["use_adapter"] = use_adapter + + self._send_cmd(cmd) + + # Wait for audio_done or audio_error + deadline = time.monotonic() + 120.0 + while time.monotonic() < deadline: + remaining = max(0.1, deadline - time.monotonic()) + resp = self._read_resp(timeout=min(remaining, 1.0)) + + if resp is None: + if not self._ensure_subprocess_alive(): + raise RuntimeError("Inference subprocess crashed during audio generation") + continue + + rtype = resp.get("type", "") + + if rtype == "audio_done": + wav_bytes = base64.b64decode(resp["wav_base64"]) + sample_rate = resp["sample_rate"] + return wav_bytes, sample_rate + + if rtype == "audio_error": + raise RuntimeError(resp.get("error", "Audio generation failed")) + + if rtype == "error": + raise RuntimeError(resp.get("error", "Unknown error")) + + if rtype == "status": + continue + + raise RuntimeError("Timeout waiting for audio generation (120s)") + + def generate_whisper_response( + self, + audio_array, + cancel_event=None, + ) -> Generator[str, None, None]: + """Whisper ASR — sends audio to subprocess, yields text.""" + yield from self._generate_audio_input_inner( + audio_array=audio_array, + audio_type="whisper", + messages=[], + system_prompt="", + cancel_event=cancel_event, + ) + + def generate_audio_input_response( + self, + messages, + system_prompt, + audio_array, + temperature: float = 0.7, + top_p: float = 0.9, + top_k: int = 40, + min_p: float = 0.0, + max_new_tokens: int = 512, + repetition_penalty: float = 1.1, + cancel_event=None, + ) -> Generator[str, None, None]: + """Audio input generation (e.g. Gemma 3n) — streams text tokens.""" + yield from self._generate_audio_input_inner( + audio_array=audio_array, + audio_type=None, # worker will use generate_audio_input_response + messages=messages, + system_prompt=system_prompt, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + max_new_tokens=max_new_tokens, + repetition_penalty=repetition_penalty, + cancel_event=cancel_event, + ) + + def _generate_audio_input_inner( + self, + audio_array, + audio_type: Optional[str] = None, + messages: list = None, + system_prompt: str = "", + temperature: float = 0.7, + top_p: float = 0.9, + top_k: int = 40, + min_p: float = 0.0, + max_new_tokens: int = 512, + repetition_penalty: float = 1.1, + cancel_event=None, + ) -> Generator[str, None, None]: + """Shared inner logic for audio input generation (Whisper + ASR).""" + if not self._ensure_subprocess_alive(): + yield "Error: Inference subprocess is not running" + return + if not self.active_model_name: + yield "Error: No active model" + return + + with self._gen_lock: + import uuid + request_id = str(uuid.uuid4()) + + # Convert numpy array to list for mp.Queue serialization + audio_data = audio_array.tolist() if hasattr(audio_array, 'tolist') else list(audio_array) + + cmd = { + "type": "generate_audio_input", + "request_id": request_id, + "audio_data": audio_data, + "audio_type": audio_type, + "messages": messages or [], + "system_prompt": system_prompt, + "temperature": temperature, + "top_p": top_p, + "top_k": top_k, + "min_p": min_p, + "max_new_tokens": max_new_tokens, + "repetition_penalty": repetition_penalty, + } + + try: + self._send_cmd(cmd) + except RuntimeError as exc: + yield f"Error: {exc}" + return + + # Yield tokens — same pattern as _generate_locked + while True: + resp = self._read_resp(timeout=30.0) + + if resp is None: + if not self._ensure_subprocess_alive(): + yield "Error: Inference subprocess crashed during audio input generation" + return + continue + + rtype = resp.get("type", "") + + if rtype == "status": + continue + + if rtype == "error" and not resp.get("request_id"): + yield f"Error: {resp.get('error', 'Unknown error')}" + return + + if rtype == "token": + if cancel_event is not None and cancel_event.is_set(): + self._cancel_generation() + self._drain_until_gen_done(timeout=5.0) + return + yield resp.get("text", "") + + elif rtype == "gen_done": + return + + elif rtype == "gen_error": + yield f"Error: {resp.get('error', 'Unknown error')}" + return + # ------------------------------------------------------------------ # Local helpers (no subprocess needed) # ------------------------------------------------------------------ diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 7683d5e314..3e766b3b01 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -165,6 +165,9 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None: "is_vision": mc.is_vision, "is_lora": mc.is_lora, "is_gguf": False, + "is_audio": getattr(mc, "is_audio", False), + "audio_type": getattr(mc, "audio_type", None), + "has_audio_input": getattr(mc, "has_audio_input", False), } _send_response(resp_queue, { "type": "loaded", @@ -267,6 +270,110 @@ def _handle_generate( }) +def _handle_generate_audio( + backend, + cmd: dict, + resp_queue: Any, +) -> None: + """Handle TTS audio generation — returns WAV bytes + sample_rate.""" + request_id = cmd.get("request_id", "") + try: + wav_bytes, sample_rate = backend.generate_audio_response( + text=cmd["text"], + temperature=cmd.get("temperature", 0.6), + top_p=cmd.get("top_p", 0.95), + top_k=cmd.get("top_k", 50), + min_p=cmd.get("min_p", 0.0), + max_new_tokens=cmd.get("max_new_tokens", 2048), + repetition_penalty=cmd.get("repetition_penalty", 1.1), + use_adapter=cmd.get("use_adapter"), + ) + + # Send WAV bytes as base64 (bytes can't go through mp.Queue directly) + _send_response(resp_queue, { + "type": "audio_done", + "request_id": request_id, + "wav_base64": base64.b64encode(wav_bytes).decode("ascii"), + "sample_rate": sample_rate, + "ts": time.time(), + }) + + except Exception as exc: + logger.error("Audio generation error: %s", exc, exc_info=True) + _send_response(resp_queue, { + "type": "audio_error", + "request_id": request_id, + "error": str(exc), + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + + +def _handle_generate_audio_input( + backend, + cmd: dict, + resp_queue: Any, + cancel_event, +) -> None: + """Handle audio input generation (ASR/Whisper) — streams text tokens back.""" + request_id = cmd.get("request_id", "") + + try: + import numpy as np + + # Decode audio array from list (numpy arrays can't go through mp.Queue) + audio_array = np.array(cmd["audio_data"], dtype=np.float32) + + audio_type = cmd.get("audio_type") + + if audio_type == "whisper": + generator = backend.generate_whisper_response( + audio_array=audio_array, + cancel_event=cancel_event, + ) + else: + generator = backend.generate_audio_input_response( + messages=cmd.get("messages", []), + system_prompt=cmd.get("system_prompt", ""), + audio_array=audio_array, + temperature=cmd.get("temperature", 0.7), + top_p=cmd.get("top_p", 0.9), + top_k=cmd.get("top_k", 40), + min_p=cmd.get("min_p", 0.0), + max_new_tokens=cmd.get("max_new_tokens", 512), + repetition_penalty=cmd.get("repetition_penalty", 1.1), + cancel_event=cancel_event, + ) + + for text_chunk in generator: + if cancel_event.is_set(): + logger.info("Audio input generation cancelled for request %s", request_id) + break + + _send_response(resp_queue, { + "type": "token", + "request_id": request_id, + "text": text_chunk, + "ts": time.time(), + }) + + _send_response(resp_queue, { + "type": "gen_done", + "request_id": request_id, + "ts": time.time(), + }) + + except Exception as exc: + logger.error("Audio input generation error: %s", exc, exc_info=True) + _send_response(resp_queue, { + "type": "gen_error", + "request_id": request_id, + "error": str(exc), + "stack": traceback.format_exc(limit=20), + "ts": time.time(), + }) + + def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None: """Handle an unload command.""" model_name = cmd.get("model_name", "") @@ -414,6 +521,14 @@ def run_inference_process( backend.unload_model(backend.active_model_name) _handle_load(backend, cmd, resp_queue) + elif cmd_type == "generate_audio": + cancel_event.clear() + _handle_generate_audio(backend, cmd, resp_queue) + + elif cmd_type == "generate_audio_input": + cancel_event.clear() + _handle_generate_audio_input(backend, cmd, resp_queue, cancel_event) + elif cmd_type == "unload": _handle_unload(backend, cmd, resp_queue) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index aa4b551dd4..d50431b035 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -23,19 +23,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""" @@ -73,6 +69,11 @@ class UnslothTrainer: # 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 @@ -109,12 +110,222 @@ 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: @@ -129,17 +340,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 @@ -154,19 +396,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, @@ -203,10 +574,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, @@ -284,8 +686,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}") @@ -348,6 +823,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", @@ -467,6 +1873,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") @@ -534,7 +1967,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") @@ -614,6 +2047,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, @@ -676,6 +2125,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() @@ -720,8 +2270,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 @@ -730,19 +2315,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, @@ -751,7 +2335,8 @@ 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 @@ -808,9 +2393,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") @@ -819,7 +2405,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}, @@ -839,14 +2425,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), @@ -885,7 +2496,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") @@ -914,7 +2526,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 @@ -922,7 +2534,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") @@ -969,103 +2581,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 ========== @@ -1074,37 +2600,12 @@ 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) - self._patch_adapter_config(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) - self._patch_adapter_config(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: diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 642b1e8e27..5323e73ae1 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -143,7 +143,8 @@ class TrainingBackend: "dataset_slice_start": kwargs.get("dataset_slice_start"), "dataset_slice_end": kwargs.get("dataset_slice_end"), "custom_format_mapping": kwargs.get("custom_format_mapping"), - "is_dataset_multimodal": kwargs.get("is_dataset_multimodal", False), + "is_dataset_image": kwargs.get("is_dataset_image", False), + "is_dataset_audio": kwargs.get("is_dataset_audio", False), "num_epochs": kwargs.get("num_epochs", 3), "learning_rate": kwargs.get("learning_rate", "2e-4"), "batch_size": kwargs.get("batch_size", 2), diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index e4b833cb9a..4c82347936 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -190,7 +190,8 @@ def run_training_process( max_seq_length=config["max_seq_length"], load_in_4bit=config["load_in_4bit"], hf_token=hf_token, - is_dataset_multimodal=config.get("is_dataset_multimodal", False), + is_dataset_image=config.get("is_dataset_image", False), + is_dataset_audio=config.get("is_dataset_audio", False), ) if not success or trainer.should_stop: if trainer.should_stop: diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py index 18f6ec224b..0fd3db7955 100644 --- a/studio/backend/models/datasets.py +++ b/studio/backend/models/datasets.py @@ -27,11 +27,14 @@ class CheckFormatResponse(BaseModel): requires_manual_mapping: bool detected_format: str columns: List[str] - is_multimodal: bool = False + is_image: bool = False + is_audio: bool = False multimodal_columns: Optional[List[str]] = None suggested_mapping: Optional[Dict[str, str]] = None detected_image_column: Optional[str] = None + detected_audio_column: Optional[str] = None detected_text_column: Optional[str] = None + detected_speaker_column: Optional[str] = None preview_samples: Optional[List[Dict]] = None total_rows: Optional[int] = None warning: Optional[str] = None 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 4a475ff8c2..4669f05a93 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -187,7 +187,7 @@ def check_format( # Run lightweight format check on the preview slice result = check_dataset_format(preview_slice, is_vlm=request.is_vlm) - logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_multimodal={result.get('is_multimodal', False)}") + logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_image={result.get('is_image', False)}") # Generate preview samples preview_samples = None @@ -230,11 +230,14 @@ def check_format( requires_manual_mapping=result["requires_manual_mapping"], detected_format=result["detected_format"], columns=result["columns"], - is_multimodal=result.get("is_multimodal", False), + is_image=result.get("is_image", False), + is_audio=result.get("is_audio", False), multimodal_columns=result.get("multimodal_columns"), suggested_mapping=result.get("suggested_mapping"), detected_image_column=result.get("detected_image_column"), + detected_audio_column=result.get("detected_audio_column"), detected_text_column=result.get("detected_text_column"), + detected_speaker_column=result.get("detected_speaker_column"), preview_samples=preview_samples, total_rows=total_rows, warning=warning, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index aa4b4ed0cd..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__) @@ -241,6 +246,9 @@ async def load_model( is_vision=config.is_vision, is_lora=config.is_lora, is_gguf=False, + is_audio=config.is_audio, + audio_type=config.audio_type, + has_audio_input=config.has_audio_input, inference=inference_config, ) @@ -387,14 +395,23 @@ async def get_status( backend = get_inference_backend() is_vision = False + is_audio = False + audio_type = None + has_audio_input = False if backend.active_model_name: model_info = backend.models.get(backend.active_model_name, {}) is_vision = model_info.get("is_vision", False) + is_audio = model_info.get("is_audio", False) + audio_type = model_info.get("audio_type") + has_audio_input = model_info.get("has_audio_input", False) return InferenceStatusResponse( active_model=backend.active_model_name, is_vision=is_vision, is_gguf=False, + is_audio=is_audio, + audio_type=audio_type, + has_audio_input=has_audio_input, loading=list(getattr(backend, 'loading_models', set())), loaded=list(backend.models.keys()), ) @@ -407,11 +424,118 @@ async def get_status( ) +# ===================================================================== +# Audio (TTS) Generation (/audio/generate) +# ===================================================================== + + +@router.post("/audio/generate") +async def generate_audio(payload: ChatCompletionRequest, request: Request): + """ + Generate audio (TTS) from the latest user message. + Returns a JSON response with base64-encoded WAV audio. + Only works when an audio model is loaded. + """ + import base64 + + backend = get_inference_backend() + if not backend.active_model_name: + raise HTTPException(status_code=400, detail="No model loaded.") + + model_info = backend.models.get(backend.active_model_name, {}) + if not model_info.get("is_audio"): + raise HTTPException(status_code=400, detail="Active model is not an audio model.") + + # Extract text from the last user message + _, chat_messages, _ = _extract_content_parts(payload.messages) + if not chat_messages: + raise HTTPException(status_code=400, detail="No messages provided.") + + last_user_msg = next( + (m for m in reversed(chat_messages) if m["role"] == "user"), None + ) + if not last_user_msg: + raise HTTPException(status_code=400, detail="No user message found.") + + text = last_user_msg["content"] + + try: + wav_bytes, sample_rate = await asyncio.get_event_loop().run_in_executor( + None, + lambda: backend.generate_audio_response( + text=text, + temperature=payload.temperature, + top_p=payload.top_p, + top_k=payload.top_k, + min_p=payload.min_p, + max_new_tokens=payload.max_tokens or 2048, + repetition_penalty=payload.repetition_penalty, + use_adapter=payload.use_adapter, + ), + ) + + audio_b64 = base64.b64encode(wav_bytes).decode("ascii") + completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" + + return JSONResponse(content={ + "id": completion_id, + "object": "chat.completion.audio", + "model": backend.active_model_name, + "audio": { + "data": audio_b64, + "format": "wav", + "sample_rate": sample_rate, + }, + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": f"[Generated audio from: \"{text[:100]}\"]", + }, + "finish_reason": "stop", + }], + }) + + except Exception as e: + logger.error(f"Audio generation error: {e}", exc_info=True) + raise HTTPException(status_code=500, detail=str(e)) + + # ===================================================================== # OpenAI-Compatible Chat Completions (/chat/completions) # ===================================================================== +def _decode_audio_base64(b64: str) -> np.ndarray: + """Decode base64 audio (any format) → float32 numpy array at 16kHz.""" + import torch + import torchaudio + import tempfile + import os + + raw = base64.b64decode(b64) + # torchaudio.load needs a file path or file-like object with format hint + # Write to a temp file so torchaudio can auto-detect the format + with tempfile.NamedTemporaryFile(suffix=".audio", delete=False) as tmp: + tmp.write(raw) + tmp_path = tmp.name + try: + waveform, sr = torchaudio.load(tmp_path) + finally: + os.unlink(tmp_path) + + # Convert to mono if stereo + if waveform.shape[0] > 1: + waveform = waveform.mean(dim=0, keepdim=True) + + # Resample to 16kHz if needed + if sr != 16000: + resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=16000) + waveform = resampler(waveform) + + return waveform.squeeze(0).numpy() + + def _extract_content_parts( messages: list, ) -> tuple[str, list[dict], "Optional[str]"]: @@ -501,6 +625,92 @@ async def openai_chat_completions( ) model_name = backend.active_model_name or payload.model + # ── Audio TTS path: auto-route to audio generation ──── + # (Whisper is ASR not TTS — handled below in audio input path) + model_info = backend.models.get(backend.active_model_name, {}) + if model_info.get("is_audio") and model_info.get("audio_type") != "whisper": + return await generate_audio(payload, request) + + # ── Whisper without audio: return clear error ── + if model_info.get("audio_type") == "whisper" and not payload.audio_base64: + raise HTTPException( + status_code=400, + detail="Whisper models require audio input. Please upload an audio file.", + ) + + # ── Audio INPUT path: decode WAV and route to audio input generation ── + if payload.audio_base64 and model_info.get("has_audio_input"): + audio_array = _decode_audio_base64(payload.audio_base64) + system_prompt, chat_messages, _ = _extract_content_parts(payload.messages) + cancel_event = threading.Event() + completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" + created = int(time.time()) + + def audio_input_generate(): + if model_info.get("audio_type") == "whisper": + return backend.generate_whisper_response( + audio_array=audio_array, + cancel_event=cancel_event, + ) + return backend.generate_audio_input_response( + messages=chat_messages, + system_prompt=system_prompt, + audio_array=audio_array, + temperature=payload.temperature, + top_p=payload.top_p, + top_k=payload.top_k, + min_p=payload.min_p, + max_new_tokens=payload.max_tokens or 512, + repetition_penalty=payload.repetition_penalty, + cancel_event=cancel_event, + ) + + if payload.stream: + async def audio_input_stream(): + try: + first_chunk = ChatCompletionChunk( + id=completion_id, created=created, model=model_name, + choices=[ChunkChoice(delta=ChoiceDelta(role="assistant"), finish_reason=None)], + ) + yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n" + + for chunk_text in audio_input_generate(): + if await request.is_disconnected(): + cancel_event.set() + return + if chunk_text: + chunk = ChatCompletionChunk( + id=completion_id, created=created, model=model_name, + choices=[ChunkChoice(delta=ChoiceDelta(content=chunk_text), finish_reason=None)], + ) + yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n" + + final_chunk = ChatCompletionChunk( + id=completion_id, created=created, model=model_name, + choices=[ChunkChoice(delta=ChoiceDelta(), finish_reason="stop")], + ) + yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n" + yield "data: [DONE]\n\n" + except asyncio.CancelledError: + cancel_event.set() + raise + except Exception as e: + logger.error(f"Error during audio input streaming: {e}", exc_info=True) + yield f"data: {json.dumps({'error': {'message': str(e), 'type': 'server_error'}})}\n\n" + + return StreamingResponse( + audio_input_stream(), + media_type="text/event-stream", + headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"}, + ) + else: + full_text = "".join(audio_input_generate()) + response = ChatCompletion( + id=completion_id, created=created, model=model_name, + choices=[CompletionChoice(message=CompletionMessage(content=full_text), finish_reason="stop")], + ) + return JSONResponse(content=response.model_dump()) + # ── Parse messages (handles multimodal content parts) ───── system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts( payload.messages diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 5f545dc133..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,40 +268,44 @@ async def list_models( @router.get("/config/{model_name:path}") async def get_model_config( model_name: str, + hf_token: Optional[str] = Query(None), current_subject: str = Depends(get_current_subject), ): """ Get configuration for a specific model. - + This endpoint wraps the backend load_model_defaults function. """ try: logger.info(f"Getting model config for: {model_name}") + from utils.models.model_config import detect_audio_type # Load model defaults from backend config_dict = load_model_defaults(model_name) + # Detect model capabilities (pass HF token for gated models) is_vision = is_vision_model(model_name) - + audio_type = detect_audio_type(model_name, hf_token=hf_token) + # Check if it's a LoRA adapter is_lora = False base_model = None - - # Try to create ModelConfig to get more info try: model_config = ModelConfig.from_identifier(model_name) is_lora = model_config.is_lora base_model = model_config.base_model if is_lora else None except Exception: - # If ModelConfig creation fails, use defaults pass - - logger.info(f"Model config result for {model_name}: is_vision={is_vision}, is_lora={is_lora}, base_model={base_model}") + + logger.info(f"Model config result for {model_name}: is_vision={is_vision}, audio_type={audio_type}, is_lora={is_lora}") return ModelDetails( id=model_name, model_name=model_name, config=config_dict, is_vision=is_vision, is_lora=is_lora, + is_audio=audio_type is not None, + audio_type=audio_type, + has_audio_input=is_audio_input_type(audio_type), base_model=base_model, ) diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index 84315291a9..d6726d69df 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -184,7 +184,8 @@ async def start_training( "finetune_language_layers": request.finetune_language_layers, "finetune_attention_modules": request.finetune_attention_modules, "finetune_mlp_modules": request.finetune_mlp_modules, - "is_dataset_multimodal": request.is_dataset_multimodal, + "is_dataset_image": request.is_dataset_image, + "is_dataset_audio": request.is_dataset_audio, "enable_wandb": request.enable_wandb, "wandb_token": request.wandb_token or "", "wandb_project": request.wandb_project or "", 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 e7b92651be..f301e42a78 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -216,12 +216,18 @@ MODEL_NAME_MAPPING = { "unsloth/Nemotron-3-Nano-30B-A3B", ], "unsloth_orpheus-3b-0.1-ft.yaml": [ + "unsloth/orpheus-3b-0.1-ft", "unsloth/orpheus-3b-0.1-ft-unsloth-bnb-4bit", "canopylabs/orpheus-3b-0.1-ft", "unsloth/orpheus-3b-0.1-ft-bnb-4bit", ], "OuteAI_Llama-OuteTTS-1.0-1B.yaml": [ "OuteAI/Llama-OuteTTS-1.0-1B", + "unsloth/Llama-OuteTTS-1.0-1B", + "unsloth/llama-outetts-1.0-1b", + "OuteAI/OuteTTS-1.0-0.6B", + "unsloth/OuteTTS-1.0-0.6B", + "unsloth/outetts-1.0-0.6b", ], "unsloth_PaddleOCR-VL.yaml": [ "unsloth/PaddleOCR-VL", @@ -320,9 +326,11 @@ MODEL_NAME_MAPPING = { ], "sesame_csm-1b.yaml": [ "sesame/csm-1b", + "unsloth/csm-1b", ], "Spark-TTS-0.5B_LLM.yaml": [ "Spark-TTS-0.5B/LLM", + "unsloth/Spark-TTS-0.5B", ], "unsloth_tinyllama-bnb-4bit.yaml": [ "unsloth/tinyllama", @@ -507,6 +515,13 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: try: config = load_model_config(model_name, use_auth=True, token=hf_token) + # Exclude audio-only models that share ForConditionalGeneration suffix + # (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration) + _audio_only_model_types = {'csm', 'whisper'} + model_type = getattr(config, 'model_type', None) + if model_type in _audio_only_model_types: + return False + # Check 1: Architecture class name patterns if hasattr(config, 'architectures'): is_vlm = any( @@ -545,6 +560,115 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: return False +VALID_AUDIO_TYPES = ('snac', 'csm', 'bicodec', 'dac', 'whisper', 'audio_vlm') + +# Cache detection results per session to avoid repeated API calls +_audio_detection_cache: Dict[str, Optional[str]] = {} + +# Tokenizer token patterns → audio_type (all 6 types detected from tokenizer_config.json) +_AUDIO_TOKEN_PATTERNS = { + 'csm': lambda tokens: '<|AUDIO|>' in tokens and '<|audio_eos|>' in tokens, + 'whisper': lambda tokens: '<|startoftranscript|>' in tokens, + 'audio_vlm': lambda tokens: '' 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: """Check if a GGUF filename is a vision projection (mmproj) file.""" return "mmproj" in filename.lower() @@ -1028,6 +1152,23 @@ def load_model_defaults(model_name: str) -> Dict[str, Any]: logger.info(f"Loaded model defaults from {config_path} (via mapping)") return config + # If model_name is a local path (e.g. /home/.../Spark-TTS-0.5B/LLM from + # adapter_config.json), try matching the last 1-2 path components against + # the registry (e.g. "Spark-TTS-0.5B/LLM"). + if model_name not in _REVERSE_MODEL_MAPPING and (model_name.startswith("/") or model_name.startswith(".")): + parts = Path(model_name).parts + for depth in [2, 1]: + if len(parts) >= depth: + suffix = "/".join(parts[-depth:]) + if suffix in _REVERSE_MODEL_MAPPING: + canonical_file = _REVERSE_MODEL_MAPPING[suffix] + for config_path in defaults_dir.rglob(canonical_file): + if config_path.is_file(): + with open(config_path, 'r', encoding='utf-8') as f: + config = yaml.safe_load(f) or {} + logger.info(f"Loaded model defaults from {config_path} (via path suffix '{suffix}')") + return config + # Try exact model name match (for backward compatibility) model_filename = model_name.replace("/", "_") + ".yaml" # Search in subfolders and root @@ -1064,6 +1205,9 @@ class ModelConfig: is_vision: bool # Is this a vision model? is_lora: bool # Is this a lora adapter? is_gguf: bool = False # Is this a GGUF model? + is_audio: bool = False # Is this a TTS audio model? + audio_type: Optional[str] = None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac' + has_audio_input: bool = False # Accepts audio input (ASR/speech understanding) gguf_file: Optional[str] = None # Full path to the .gguf file (local mode) gguf_mmproj_file: Optional[str] = None # Full path to the mmproj .gguf file (vision projection) gguf_hf_repo: Optional[str] = None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF") @@ -1100,6 +1244,9 @@ class ModelConfig: # Check if base model is vision is_vision = is_vision_model(base_model, hf_token=hf_token) + # Check if base model is audio + audio_type = detect_audio_type(base_model, hf_token=hf_token) + display_name = lora_path_obj.name identifier = lora_path # Use path as identifier for local LoRAs @@ -1111,6 +1258,9 @@ class ModelConfig: is_cached=True, # Local LoRAs are always "cached" is_vision=is_vision, is_lora=True, + is_audio=audio_type is not None and audio_type != 'audio_vlm', + audio_type=audio_type, + has_audio_input=is_audio_input_type(audio_type), base_model=base_model, ) @@ -1288,12 +1438,16 @@ class ModelConfig: if not base_model: logger.warning(f"Could not determine base model for LoRA '{path}'") return None - vision = is_vision_model(base_model, hf_token=hf_token) + check_model = base_model else: - vision = is_vision_model(identifier, hf_token=hf_token) - + check_model = identifier + + vision = is_vision_model(check_model, hf_token=hf_token) + audio_type_val = detect_audio_type(check_model, hf_token=hf_token) + has_audio_in = is_audio_input_type(audio_type_val) + display_name = Path(path).name if is_local else identifier.split("/")[-1] - + return cls( identifier=identifier, display_name=display_name, @@ -1302,6 +1456,9 @@ class ModelConfig: is_cached=is_model_cached(identifier) if not is_local else True, is_vision=vision, is_lora=is_lora, + is_audio=audio_type_val is not None and audio_type_val != 'audio_vlm', + audio_type=audio_type_val, + has_audio_input=has_audio_in, base_model=base_model, ) @@ -1381,4 +1538,3 @@ class ModelConfig: is_lora=is_lora, base_model=base_model, # This will be None for base models, and populated for LoRAs ) - pass diff --git a/studio/frontend/package.json b/studio/frontend/package.json index 3d6881cc53..e24843a301 100644 --- a/studio/frontend/package.json +++ b/studio/frontend/package.json @@ -32,10 +32,10 @@ "@radix-ui/react-select": "^2.2.6", "@radix-ui/react-separator": "^1.1.8", "@radix-ui/react-slot": "^1.2.4", - "@streamdown/cjk": "^1.0.2", - "@streamdown/code": "^1.0.2", - "@streamdown/math": "^1.0.2", - "@streamdown/mermaid": "^1.0.2", + "@streamdown/cjk": "1.0.2", + "@streamdown/code": "1.0.2", + "@streamdown/math": "1.0.2", + "@streamdown/mermaid": "1.0.2", "@tailwindcss/vite": "^4.1.18", "@tanstack/react-router": "^1.159.10", "@tanstack/react-table": "^8.21.3", @@ -66,7 +66,7 @@ "remark-gfm": "^4.0.1", "shadcn": "^3.8.4", "sonner": "^2.0.7", - "streamdown": "^2.3.0", + "streamdown": "2.3.0", "tailwind-merge": "^3.4.0", "tailwindcss": "^4.1.18", "tw-animate-css": "^1.4.0", 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/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 4276e618ed..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, @@ -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/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} + +
+ )}
)}