From c48437848d20f81575b77ed5ccedf82ee2883f7f Mon Sep 17 00:00:00 2001 From: Manan17 Date: Sun, 1 Mar 2026 02:30:31 +0000 Subject: [PATCH] revamping up the code and adding inference --- .../gemma/unsloth_gemma-3n-E4B-it.yaml | 2 + .../gemma/unsloth_gemma-3n-E4B.yaml | 2 + studio/backend/core/inference/OuteTTS | 1 + studio/backend/core/inference/audio_codecs.py | 280 ++++ studio/backend/core/inference/inference.py | 361 +++++ .../backend/core/training/inference/OuteTTS | 1 + studio/backend/core/training/trainer.py | 1265 +++++++---------- studio/backend/models/inference.py | 7 + studio/backend/models/models.py | 3 + studio/backend/routes/inference.py | 197 +++ studio/backend/routes/models.py | 5 +- studio/backend/utils/models/model_config.py | 60 +- .../components/assistant-ui/audio-player.tsx | 114 ++ .../components/assistant-ui/markdown-text.tsx | 7 + .../src/components/assistant-ui/thread.tsx | 100 +- .../src/features/chat/api/chat-adapter.ts | 78 +- .../src/features/chat/api/chat-api.ts | 20 + .../chat/hooks/use-chat-model-runtime.ts | 13 +- .../src/features/chat/shared-composer.tsx | 86 +- .../chat/stores/chat-runtime-store.ts | 10 + .../frontend/src/features/chat/types/api.ts | 26 + .../src/features/chat/types/runtime.ts | 3 + 22 files changed, 1861 insertions(+), 780 deletions(-) create mode 160000 studio/backend/core/inference/OuteTTS create mode 100644 studio/backend/core/inference/audio_codecs.py create mode 160000 studio/backend/core/training/inference/OuteTTS create mode 100644 studio/frontend/src/components/assistant-ui/audio-player.tsx 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/core/inference/OuteTTS b/studio/backend/core/inference/OuteTTS new file mode 160000 index 0000000000..59d896747a --- /dev/null +++ b/studio/backend/core/inference/OuteTTS @@ -0,0 +1 @@ +Subproject commit 59d896747aa0a6a207837e7da2d6921805eae684 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 1147c281b7..284c4fa392 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,92 @@ class InferenceBackend: self.models[model_name] = { "is_vision": config.is_vision, "is_lora": config.is_lora, + "is_audio": config.is_audio, + "audio_type": config.audio_type, + "has_audio_input": config.has_audio_input, "model_path": config.path, "base_model": config.base_model if config.is_lora else None, "loaded_adapters": {}, "active_adapter": None, } + # ── Audio model loading path ────────────────────────── + if config.is_audio: + audio_type = config.audio_type + adapter_info = " (LoRA adapter)" if config.is_lora else "" + logger.info(f"Loading audio ({audio_type}) model{adapter_info}: {model_name}") + log_gpu_memory(f"Before loading {model_name}") + + if audio_type == "csm": + from unsloth import FastModel + from transformers import CsmForConditionalGeneration + model, processor = FastModel.from_pretrained( + config.path, + auto_model=CsmForConditionalGeneration, + load_in_4bit=False, + token=hf_token if hf_token and hf_token.strip() else None, + ) + FastModel.for_inference(model) + self.models[model_name]["model"] = model + self.models[model_name]["tokenizer"] = processor + self.models[model_name]["processor"] = processor + elif audio_type == "bicodec": + import os + from unsloth import FastModel + from huggingface_hub import snapshot_download + + # Spark-TTS: download full repo, then load from /LLM subfolder + hf_repo = config.path + local_dir = hf_repo.split("/")[-1] + repo_path = snapshot_download(hf_repo, local_dir=local_dir) + abs_repo_path = os.path.abspath(repo_path) + llm_path = os.path.join(abs_repo_path, "LLM") + logger.info(f"Spark-TTS: downloaded repo to {repo_path}, loading LLM from {llm_path}") + + model, tokenizer = FastModel.from_pretrained( + llm_path, + dtype=torch.float32, + load_in_4bit=False, + token=hf_token if hf_token and hf_token.strip() else None, + ) + FastModel.for_inference(model) + self.models[model_name]["model"] = model + self.models[model_name]["tokenizer"] = tokenizer + self.models[model_name]["model_repo_path"] = abs_repo_path + elif audio_type == "dac": + # OuteTTS uses FastModel (not FastLanguageModel) + from unsloth import FastModel + model, tokenizer = FastModel.from_pretrained( + config.path, + max_seq_length=max_seq_length, + load_in_4bit=False, + token=hf_token if hf_token and hf_token.strip() else None, + ) + FastModel.for_inference(model) + self.models[model_name]["model"] = model + self.models[model_name]["tokenizer"] = tokenizer + else: + # SNAC (Orpheus) uses FastLanguageModel + model, tokenizer = FastLanguageModel.from_pretrained( + model_name=config.path, + max_seq_length=max_seq_length, + load_in_4bit=False, + token=hf_token if hf_token and hf_token.strip() else None, + ) + FastLanguageModel.for_inference(model) + self.models[model_name]["model"] = model + self.models[model_name]["tokenizer"] = tokenizer + + # Load the external codec for this audio type + model_repo_path = self.models[model_name].get("model_repo_path") + self._audio_codec_manager.load_codec(audio_type, self.device, model_repo_path=model_repo_path) + + self.active_model_name = model_name + self.loading_models.discard(model_name) + logger.info(f"Successfully loaded audio model: {model_name}") + log_gpu_memory(f"After loading {model_name}") + return True + model_type = "vision" if config.is_vision else "text" adapter_info = " (LoRA adapter)" if self.models[model_name]["is_lora"] else "" logger.info(f"Loading {model_type} model{adapter_info}: {model_name}") @@ -186,6 +268,10 @@ class InferenceBackend: """ if model_name in self.models: try: + # If this was an audio model, clean up codecs + if self.models[model_name].get("is_audio"): + self._audio_codec_manager.unload() + logger.info(f"Unloading model '{model_name}' from memory.") # Delete the model entry from our registry del self.models[model_name] @@ -801,6 +887,122 @@ class InferenceBackend: yield f"Error: {str(e)}" pass + def generate_audio_input_response(self, messages, system_prompt, audio_array, + temperature, top_p, top_k, min_p, + max_new_tokens, repetition_penalty, + cancel_event=None) -> Generator[str, None, None]: + """Handle audio input (ASR) generation — accepts audio numpy array, streams text output. + + Uses processor.apply_chat_template with audio embedded in messages (Gemma 3n pattern). + """ + import threading + import numpy as np + + model_info = self.models[self.active_model_name] + model = model_info["model"] + processor = model_info.get("processor") or model_info.get("tokenizer") + raw_tokenizer = getattr(processor, "tokenizer", processor) + + # Extract last user text + user_text = "Transcribe this audio." + if messages: + for msg in reversed(messages): + if msg["role"] == "user" and msg.get("content"): + user_text = msg["content"] + break + + # Build messages in Gemma 3n format — audio goes INTO apply_chat_template + audio_messages = [] + if system_prompt: + audio_messages.append({"role": "system", "content": [{"type": "text", "text": system_prompt}]}) + audio_messages.append({ + "role": "user", + "content": [ + {"type": "audio", "audio": audio_array}, + {"type": "text", "text": user_text}, + ], + }) + + # apply_chat_template handles audio embedding + tokenization in one step + inputs = processor.apply_chat_template( + audio_messages, + add_generation_prompt=True, + tokenize=True, + return_dict=True, + return_tensors="pt", + ).to(self.device) + + try: + from transformers import TextIteratorStreamer + from queue import Empty + + streamer = TextIteratorStreamer( + raw_tokenizer, + skip_prompt=True, + skip_special_tokens=True, + timeout=0.2, + ) + + generation_kwargs = dict( + **inputs, + streamer=streamer, + max_new_tokens=max_new_tokens, + use_cache=True, + do_sample=temperature > 0, + temperature=temperature, + top_p=top_p, + top_k=top_k, + min_p=min_p, + ) + + err: dict[str, str] = {} + + def generate_fn(): + with self._generation_lock: + try: + model.generate(**generation_kwargs) + except Exception as e: + err["msg"] = str(e) + logger.error(f"Audio input generation error in thread: {e}") + finally: + try: + streamer.end() + except Exception: + pass + + thread = threading.Thread(target=generate_fn) + thread.start() + + output = "" + try: + while True: + if cancel_event is not None and cancel_event.is_set(): + break + try: + new_token = next(streamer) + except StopIteration: + break + except Empty: + if not thread.is_alive(): + break + continue + if new_token: + output += new_token + yield new_token + finally: + if cancel_event is not None: + cancel_event.set() + thread.join(timeout=10) + if thread.is_alive(): + logger.warning("Audio input generation thread did not exit after cancel/join timeout") + + if err.get("msg"): + yield f"Error: {err['msg']}" + + except Exception as e: + logger.error(f"Audio input generation error: {e}") + yield f"Error: {str(e)}" + def generate_stream(self, prompt: str, temperature: float = 0.7, @@ -925,6 +1127,165 @@ class InferenceBackend: # ... other helper methods (format_chat_prompt, _clean_generated_text, etc.) pass + # ── Audio (TTS) Generation ──────────────────────────────────── + + def generate_audio_response( + self, + text: str, + temperature: float = 0.6, + top_p: float = 0.95, + top_k: int = 50, + min_p: float = 0.0, + max_new_tokens: int = 2048, + repetition_penalty: float = 1.1, + use_adapter: Optional[Union[bool, str]] = None, + ) -> Tuple[bytes, int]: + """ + Generate audio from text for TTS models. + Returns (wav_bytes, sample_rate). + Blocking — generates complete audio before returning. + """ + if not self.active_model_name: + raise RuntimeError("No active model") + + model_info = self.models[self.active_model_name] + audio_type = model_info.get("audio_type") + model = model_info["model"] + tokenizer = model_info.get("tokenizer") + + if not audio_type: + raise RuntimeError(f"Model {self.active_model_name} is not an audio model") + + top_k = self._normalize_top_k(top_k) + + with self._generation_lock: + if use_adapter is not None: + self._apply_adapter_state(use_adapter) + + if audio_type == "snac": + return self._generate_snac(model, tokenizer, text, temperature, top_p, max_new_tokens, repetition_penalty) + elif audio_type == "csm": + processor = model_info.get("processor", tokenizer) + return self._generate_csm(model, processor, text, max_new_tokens) + elif audio_type == "bicodec": + return self._generate_bicodec(model, tokenizer, text, temperature, top_k, max_new_tokens) + elif audio_type == "dac": + return self._generate_dac(model, tokenizer, text, temperature, top_k, top_p, min_p, max_new_tokens, repetition_penalty) + else: + raise RuntimeError(f"Unknown audio_type: {audio_type}") + + def _generate_snac(self, model, tokenizer, text, temperature, top_p, max_new_tokens, repetition_penalty): + """Generate audio using SNAC codec (Orpheus).""" + device = model.device + start_token = torch.tensor([[128259]], device=device) # START_OF_HUMAN + end_tokens = torch.tensor([[128009, 128260]], device=device) # EOT, END_OF_HUMAN + text_ids = tokenizer(text, return_tensors="pt").input_ids.to(device) + input_ids = torch.cat([start_token, text_ids, end_tokens], dim=1) + attention_mask = torch.ones_like(input_ids) + + generated = model.generate( + input_ids=input_ids, + attention_mask=attention_mask, + max_new_tokens=max_new_tokens, + do_sample=True, + temperature=temperature, + top_p=top_p, + repetition_penalty=repetition_penalty, + eos_token_id=128258, # END_OF_SPEECH + use_cache=True, + ) + return self._audio_codec_manager.decode_snac(generated, str(device)) + + def _generate_csm(self, model, processor, text, max_new_tokens): + """Generate audio using CSM (Sesame).""" + speaker_id = 0 + inputs = processor(f"[{speaker_id}]{text}", add_special_tokens=True, return_tensors="pt").to(model.device) + audio_values = model.generate(**inputs, max_new_tokens=max_new_tokens, output_audio=True) + return self._audio_codec_manager.decode_csm(audio_values) + + def _generate_bicodec(self, model, tokenizer, text, temperature, top_k, max_new_tokens): + """Generate audio using BiCodec (Spark-TTS).""" + prompt = "<|task_tts|><|start_content|>" + text + "<|end_content|><|start_global_token|>" + inputs = tokenizer([prompt], return_tensors="pt").to(model.device) + generated = model.generate( + **inputs, + max_new_tokens=max_new_tokens, + do_sample=True, + temperature=temperature, + top_k=top_k, + eos_token_id=tokenizer.eos_token_id, + pad_token_id=tokenizer.pad_token_id, + ) + new_tokens = generated[:, inputs.input_ids.shape[1]:] + decoded_text = tokenizer.batch_decode(new_tokens, skip_special_tokens=False)[0] + return self._audio_codec_manager.decode_bicodec(decoded_text, str(model.device)) + + def _generate_dac(self, model, tokenizer, text, temperature, top_k, top_p, min_p, max_new_tokens, repetition_penalty): + """Generate audio using DAC (OuteTTS). Follows Oute_TTS_(1B).ipynb exactly.""" + # Monkey-patch RepetitionPenaltyLogitsProcessor with a 64-token penalty + # window (same as the OuteTTS notebook) to avoid degenerate repetition. + self._patch_repetition_penalty_processor() + + prompt = "<|im_start|>\n<|text_start|>" + text + "<|text_end|>\n<|audio_start|><|global_features_start|>\n" + with torch.inference_mode(): + with torch.amp.autocast('cuda', dtype=model.dtype): + inputs = tokenizer([prompt], return_tensors="pt").to(model.device) + generated = model.generate( + **inputs, + temperature=temperature, + top_k=top_k, + top_p=top_p, + min_p=min_p, + repetition_penalty=repetition_penalty, + max_new_tokens=max_new_tokens, + ) + decoded_text = tokenizer.batch_decode(generated, skip_special_tokens=False)[0] + return self._audio_codec_manager.decode_dac(decoded_text, str(model.device)) + + _repetition_penalty_patched = False + + @classmethod + def _patch_repetition_penalty_processor(cls): + """ + Monkey-patch transformers' RepetitionPenaltyLogitsProcessor with a + 64-token sliding window variant (from the OuteTTS notebook). + Only applied once per process. + """ + if cls._repetition_penalty_patched: + return + cls._repetition_penalty_patched = True + + from transformers import LogitsProcessor + import transformers.generation.utils as generation_utils + + class RepetitionPenaltyLogitsProcessorPatch(LogitsProcessor): + def __init__(self, penalty: float): + self.penalty_last_n = 64 + if not isinstance(penalty, float) or penalty <= 0: + raise ValueError(f"`penalty` has to be a positive float, but is {penalty}") + self.penalty = penalty + + @torch.no_grad() + def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor) -> torch.FloatTensor: + if self.penalty_last_n == 0 or self.penalty == 1.0: + return scores + batch_size, seq_len = input_ids.shape + vocab_size = scores.shape[-1] + for b in range(batch_size): + start_index = max(0, seq_len - self.penalty_last_n) + window_indices = input_ids[b, start_index:] + if window_indices.numel() == 0: + continue + for token_id in set(window_indices.tolist()): + if token_id >= vocab_size: + continue + logit = scores[b, token_id] + scores[b, token_id] = logit * self.penalty if logit <= 0 else logit / self.penalty + return scores + + generation_utils.RepetitionPenaltyLogitsProcessor = RepetitionPenaltyLogitsProcessorPatch + logger.info("Patched RepetitionPenaltyLogitsProcessor with 64-token window for OuteTTS") + def format_chat_prompt(self, messages: list, system_prompt: str = None) -> str: if not self.active_model_name or self.active_model_name not in self.models: logger.error("No active model available") diff --git a/studio/backend/core/training/inference/OuteTTS b/studio/backend/core/training/inference/OuteTTS new file mode 160000 index 0000000000..59d896747a --- /dev/null +++ b/studio/backend/core/training/inference/OuteTTS @@ -0,0 +1 @@ +Subproject commit 59d896747aa0a6a207837e7da2d6921805eae684 diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index d77effc0d4..b9cfa6d130 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -22,19 +22,20 @@ 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.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__) +# Process-level flag: set True after CUDA-heavy audio preprocessing (Whisper/DAC/BiCodec). +# Once CUDA has been used for audio processing, fork-based multiprocessing (num_proc>1) +# deadlocks because forked children inherit CUDA's internal thread locks. +# This flag is never reset — once contaminated, the process stays contaminated. +_CUDA_AUDIO_PREPROCESSING_DONE = False + @dataclass class TrainingProgress: """Training progress tracking""" @@ -111,6 +112,177 @@ 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) + 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) + 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. @@ -178,6 +350,26 @@ 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() @@ -191,6 +383,7 @@ class UnslothTrainer: # VLM: vision model with image dataset (mutually exclusive with audio VLM) self.is_vlm = not self.is_audio and not self.is_audio_vlm and is_vision_model(model_name) and is_dataset_multimodal self.model_name = model_name + self.max_seq_length = max_seq_length logger.info(f"Audio type: {self._audio_type}") if not self.is_audio: @@ -302,8 +495,15 @@ class UnslothTrainer: logger.info("Loaded Spark-TTS (bicodec) model") elif self._audio_type == 'dac': - # Phase 2: OuteTTS - raise NotImplementedError(f"Audio model type '{self._audio_type}' not yet implemented") + # 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) @@ -865,7 +1065,7 @@ class UnslothTrainer: SNAC_MODEL_NAME = "hubertsiuzdak/snac_24khz" SNAC_SAMPLE_RATE = 24000 device = "cuda" if torch.cuda.is_available() else "cpu" - max_length = getattr(self, '_max_seq_length', 2048) or 2048 + max_length = self.max_seq_length or 2048 tokenizer = self.tokenizer # Orpheus special token IDs (hardcoded in tokenizer vocabulary) @@ -1001,8 +1201,13 @@ class UnslothTrainer: print("Freeing SNAC codec model from GPU...\n") snac_model.to("cpu") del snac_model + import gc + gc.collect() torch.cuda.empty_cache() + global _CUDA_AUDIO_PREPROCESSING_DONE + _CUDA_AUDIO_PREPROCESSING_DONE = True + if not processed_examples: raise ValueError( f"No valid examples after SNAC preprocessing (skipped {skipped})" @@ -1185,8 +1390,14 @@ class UnslothTrainer: 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() + global _CUDA_AUDIO_PREPROCESSING_DONE + _CUDA_AUDIO_PREPROCESSING_DONE = True + if not processed_examples: raise ValueError( f"No valid examples after BiCodec preprocessing (skipped {skipped})" @@ -1201,6 +1412,193 @@ class UnslothTrainer: 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() + + # Mark process as CUDA-contaminated from audio preprocessing. + # Fork-based multiprocessing (num_proc>1) will deadlock after this + # because forked children inherit CUDA's internal thread locks from + # Whisper/DAC processing that can't be released. + global _CUDA_AUDIO_PREPROCESSING_DONE + _CUDA_AUDIO_PREPROCESSING_DONE = 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. @@ -1404,10 +1802,13 @@ class UnslothTrainer: elif self._audio_type == 'bicodec': processed = self._preprocess_bicodec_dataset(dataset, custom_format_mapping) - return (processed, None) + return ({"dataset": processed, "final_format": "audio_bicodec"}, None) - elif self._audio_type in ('xcodec2', 'dac'): - # Phase 2: remaining codec-to-text models + elif self._audio_type == 'dac': + processed = self._preprocess_dac_dataset(dataset, custom_format_mapping) + return ({"dataset": processed, "final_format": "audio_dac"}, None) + + elif self._audio_type == 'xcodec2': raise NotImplementedError(f"Audio dataset preprocessing for '{self._audio_type}' not yet implemented") elif self.is_audio_vlm: @@ -1632,671 +2033,98 @@ class UnslothTrainer: # ========== AUDIO TRAINER BRANCH ========== if self._audio_type == 'csm': - # CSM uses plain HF Trainer with TrainingArguments (NOT SFTTrainer) - # Dataset is already preprocessed — just pass it directly - from transformers import Trainer as HFTrainer, TrainingArguments, TrainerCallback - - # --- Fix: Unsloth's forward patch for CsmForConditionalGeneration fails to - # apply on transformers>=4.54 due to type annotation mismatches (Optional[], - # list vs List, Unpack[TransformersKwargs] vs KWARGS_TYPE). The original - # forward passes **kwargs (containing num_items_in_batch, return_dict, etc.) - # directly to the depth decoder, which causes depth_decoder_loss=None. - # We replicate the critical fixes from the Unsloth patched forward here. + # 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() - 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') - - csm_training_args = { - "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", - # CSM needs input_values + input_values_cutoffs for depth decoder loss; - # without this, Trainer strips them and depth_decoder_loss becomes None + config = self._build_audio_training_args(training_args, output_dir, extra_args={ "remove_unused_columns": False, - } - - # max_steps vs epochs - if max_steps_val and max_steps_val > 0: - csm_training_args["max_steps"] = max_steps_val - print(f"CSM training for {max_steps_val} steps\n") - else: - csm_training_args["num_train_epochs"] = training_args.get('num_epochs', 3) - print(f"CSM training for {csm_training_args['num_train_epochs']} epochs\n") - - # save_steps - save_steps_val = training_args.get('save_steps', 0) - if save_steps_val and save_steps_val > 0: - csm_training_args["save_steps"] = save_steps_val - csm_training_args["save_strategy"] = "steps" - - # The dataset for CSM is a plain Dataset (not a dict) - train_ds = dataset - - print(f"CSM training config: {csm_training_args}\n") - + }) self.trainer = HFTrainer( - model=self.model, - train_dataset=train_ds, - args=TrainingArguments(**csm_training_args), + model=self.model, train_dataset=dataset, + args=TrainingArguments(**config), ) - print("CSM Trainer initialized\n") + self.trainer.add_callback(self._create_progress_callback()) - # Progress callback (same as standard) - class ProgressCallback(TrainerCallback): - def __init__(self, trainer_instance): - self.trainer_instance = trainer_instance - - def on_log(self, args, state, control, logs=None, **kwargs): - if logs: - 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 self.trainer_instance.training_start_time is not None: - elapsed_seconds = time.time() - self.trainer_instance.training_start_time - - 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 - - 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): - self.trainer_instance._update_progress( - epoch=state.epoch, - step=state.global_step - ) - - def on_step_end(self, args, state, control, **kwargs): - if self.trainer_instance.should_stop: - print(f"Stop detected at step {state.global_step}\n") - control.should_training_stop = True - return control - - self.trainer.add_callback(ProgressCallback(self)) - - # Calculate total steps - num_samples = len(train_ds) - grad_accum = training_args.get('gradient_accumulation_steps', 4) - num_epochs = training_args.get('num_epochs', 3) - len_dataloader = math.ceil(num_samples / batch_size) - num_update_steps_per_epoch = max( - len_dataloader // grad_accum + int(len_dataloader % grad_accum > 0), 1 + 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), ) - - if max_steps_val and max_steps_val > 0: - total_steps = max_steps_val - else: - total_steps = num_update_steps_per_epoch * num_epochs - - self._update_progress(total_steps=total_steps) - print(f"CSM progress tracking: {total_steps} total steps\n") - - # Train - self._update_progress(status_message="Starting CSM training...") - print("Starting CSM training...\n") + self._update_progress(total_steps=total, status_message="Starting CSM training...") + print(f"CSM training config: {config}\n") self.trainer.train() - - # Save - if self.should_stop and self.save_on_stop: - self.trainer.save_model() - self.tokenizer.save_pretrained(output_dir) - print(f"\nCSM training 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: - print("\nCSM training cancelled.\n") - self._update_progress( - is_training=False, - status_message="Training cancelled.", - ) - else: - self.trainer.save_model() - self.tokenizer.save_pretrained(output_dir) - print(f"\nCSM training 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}", - ) - return # Exit _train_worker for CSM + self._finalize_training(output_dir, "CSM") + return elif self._audio_type == 'snac': - # Orpheus: language model with SNAC codec tokens - # Dataset is already preprocessed — use plain HF Trainer (same as CSM) - from transformers import Trainer as HFTrainer, TrainingArguments, TrainerCallback + # Orpheus: language model with SNAC codec tokens — plain HF Trainer + from transformers import Trainer as HFTrainer, TrainingArguments + + config = self._build_audio_training_args(training_args, output_dir) + 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) - 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') - - snac_training_args = { - "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: - snac_training_args["max_steps"] = max_steps_val - print(f"snac training for {max_steps_val} steps\n") - else: - snac_training_args["num_train_epochs"] = training_args.get('num_epochs', 3) - print(f"snac training for {snac_training_args['num_train_epochs']} epochs\n") - - # save_steps - save_steps_val = training_args.get('save_steps', 0) - if save_steps_val and save_steps_val > 0: - snac_training_args["save_steps"] = save_steps_val - snac_training_args["save_strategy"] = "steps" - - train_ds = dataset - - print(f"snac training config: {snac_training_args}\n") - - self.trainer = HFTrainer( - model=self.model, - train_dataset=train_ds, - args=TrainingArguments(**snac_training_args), + 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), ) - print("snac Trainer initialized\n") - - # Progress callback (same as CSM) - class ProgressCallback(TrainerCallback): - def __init__(self, trainer_instance): - self.trainer_instance = trainer_instance - - def on_log(self, args, state, control, logs=None, **kwargs): - if logs: - 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 self.trainer_instance.training_start_time is not None: - elapsed_seconds = time.time() - self.trainer_instance.training_start_time - - 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 - - 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): - self.trainer_instance._update_progress( - epoch=state.epoch, - step=state.global_step - ) - - def on_step_end(self, args, state, control, **kwargs): - if self.trainer_instance.should_stop: - print(f"Stop detected at step {state.global_step}\n") - control.should_training_stop = True - return control - - self.trainer.add_callback(ProgressCallback(self)) - - # Calculate total steps - num_samples = len(train_ds) - grad_accum = training_args.get('gradient_accumulation_steps', 4) - num_epochs = training_args.get('num_epochs', 3) - len_dataloader = math.ceil(num_samples / batch_size) - num_update_steps_per_epoch = max( - len_dataloader // grad_accum + int(len_dataloader % grad_accum > 0), 1 - ) - - if max_steps_val and max_steps_val > 0: - total_steps = max_steps_val - else: - total_steps = num_update_steps_per_epoch * num_epochs - - self._update_progress(total_steps=total_steps) - print(f"snac progress tracking: {total_steps} total steps\n") - - # Train - self._update_progress(status_message="Starting snac training...") - print("Starting snac training...\n") + self._update_progress(total_steps=total, status_message="Starting SNAC training...") + print(f"SNAC training config: {config}\n") self.trainer.train() - - # Save - if self.should_stop and self.save_on_stop: - self.trainer.save_model() - self.tokenizer.save_pretrained(output_dir) - print(f"\nsnac training 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: - print("\nsnac training cancelled.\n") - self._update_progress( - is_training=False, - status_message="Training cancelled.", - ) - else: - self.trainer.save_model() - self.tokenizer.save_pretrained(output_dir) - print(f"\nsnac training 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}", - ) - return # Exit _train_worker for snac + self._finalize_training(output_dir, "SNAC") + return elif self._audio_type == 'whisper': # Whisper: Seq2SeqTrainer with custom speech collator - from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments, TrainerCallback + from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments from utils.datasets import DataCollatorSpeechSeq2SeqWithPadding - batch_size = training_args.get('batch_size', 1) - 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', 1e-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') eval_dataset = training_args.get('eval_dataset', None) - eval_steps_val = training_args.get('eval_steps', 5) - - whisper_training_args = { - "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", - "remove_unused_columns": False, - "label_names": ["labels"], - } - - # Eval config + extra = {"remove_unused_columns": False, "label_names": ["labels"]} if eval_dataset: - whisper_training_args["eval_strategy"] = "steps" - whisper_training_args["eval_steps"] = eval_steps_val + extra["eval_strategy"] = "steps" + extra["eval_steps"] = training_args.get('eval_steps', 5) - # max_steps vs epochs - if max_steps_val and max_steps_val > 0: - whisper_training_args["max_steps"] = max_steps_val - print(f"Whisper training for {max_steps_val} steps\n") - else: - whisper_training_args["num_train_epochs"] = training_args.get('num_epochs', 3) - print(f"Whisper training for {whisper_training_args['num_train_epochs']} epochs\n") - - # save_steps - save_steps_val = training_args.get('save_steps', 0) - if save_steps_val and save_steps_val > 0: - whisper_training_args["save_steps"] = save_steps_val - whisper_training_args["save_strategy"] = "steps" - - train_ds = dataset - data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=self.tokenizer) - - print(f"Whisper training config: {whisper_training_args}\n") + config = self._build_audio_training_args(training_args, output_dir, extra_args=extra) trainer_kwargs = { "model": self.model, - "train_dataset": train_ds, - "data_collator": data_collator, + "train_dataset": dataset, + "data_collator": DataCollatorSpeechSeq2SeqWithPadding(processor=self.tokenizer), "processing_class": self.tokenizer.feature_extractor, - "args": Seq2SeqTrainingArguments(**whisper_training_args), + "args": Seq2SeqTrainingArguments(**config), } if eval_dataset: trainer_kwargs["eval_dataset"] = eval_dataset self.trainer = Seq2SeqTrainer(**trainer_kwargs) - print("Whisper Seq2SeqTrainer initialized\n") - - # Progress callback (same as CSM/SNAC) - class ProgressCallback(TrainerCallback): - def __init__(self, trainer_instance): - self.trainer_instance = trainer_instance - - def on_log(self, args, state, control, logs=None, **kwargs): - if logs: - 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 self.trainer_instance.training_start_time is not None: - elapsed_seconds = time.time() - self.trainer_instance.training_start_time - - 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 - - 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): - self.trainer_instance._update_progress( - epoch=state.epoch, - step=state.global_step - ) - - def on_step_end(self, args, state, control, **kwargs): - if self.trainer_instance.should_stop: - print(f"Stop detected at step {state.global_step}\n") - control.should_training_stop = True - return control - - self.trainer.add_callback(ProgressCallback(self)) - - # Calculate total steps - num_samples = len(train_ds) - grad_accum = training_args.get('gradient_accumulation_steps', 4) - num_epochs = training_args.get('num_epochs', 3) - len_dataloader = math.ceil(num_samples / batch_size) - num_update_steps_per_epoch = max( - len_dataloader // grad_accum + int(len_dataloader % grad_accum > 0), 1 - ) - - if max_steps_val and max_steps_val > 0: - total_steps = max_steps_val - else: - total_steps = num_update_steps_per_epoch * num_epochs - - self._update_progress(total_steps=total_steps) - print(f"Whisper progress tracking: {total_steps} total steps\n") - - # Train - self._update_progress(status_message="Starting Whisper training...") - print("Starting Whisper training...\n") - self.trainer.train() - - # Save - if self.should_stop and self.save_on_stop: - self.trainer.save_model() - self.tokenizer.save_pretrained(output_dir) - print(f"\nWhisper training 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: - print("\nWhisper training cancelled.\n") - self._update_progress( - is_training=False, - status_message="Training cancelled.", - ) - else: - self.trainer.save_model() - self.tokenizer.save_pretrained(output_dir) - print(f"\nWhisper training 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}", - ) - return # Exit _train_worker for Whisper - - elif self._audio_type == 'bicodec': - # Spark-TTS: SFTTrainer with dataset_text_field="text" - # Dataset is already preprocessed to text strings with BiCodec tokens - from transformers import TrainerCallback + self.trainer.add_callback(self._create_progress_callback()) 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') - max_seq_length = training_args.get('max_seq_length', 2048) - - print(f"BiCodec training params: lr={learning_rate}, warmup={warmup_steps_val}, " - f"max_steps={max_steps_val}, batch={batch_size}, max_seq_len={max_seq_length}\n") - - bicodec_training_args = { - "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": False, # Spark-TTS requires full float32 - "bf16": False, # Spark-TTS requires full float32 - "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: - bicodec_training_args["max_steps"] = max_steps_val - print(f"BiCodec training for {max_steps_val} steps\n") - else: - bicodec_training_args["num_train_epochs"] = training_args.get('num_epochs', 3) - print(f"BiCodec training for {bicodec_training_args['num_train_epochs']} epochs\n") - - # save_steps - save_steps_val = training_args.get('save_steps', 0) - if save_steps_val and save_steps_val > 0: - bicodec_training_args["save_steps"] = save_steps_val - bicodec_training_args["save_strategy"] = "steps" - - train_ds = dataset - - print(f"BiCodec training config: {bicodec_training_args}\n") - - self.trainer = SFTTrainer( - model=self.model, - tokenizer=self.tokenizer, - train_dataset=train_ds, - dataset_text_field="text", - max_seq_length=max_seq_length, - packing=False, - args=SFTConfig(**bicodec_training_args), + 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), ) - print("BiCodec SFTTrainer initialized\n") - - # Progress callback (same pattern as CSM/SNAC) - class ProgressCallback(TrainerCallback): - def __init__(self, trainer_instance): - self.trainer_instance = trainer_instance - - def on_log(self, args, state, control, logs=None, **kwargs): - if logs: - 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 self.trainer_instance.training_start_time is not None: - elapsed_seconds = time.time() - self.trainer_instance.training_start_time - - 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 - - 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): - self.trainer_instance._update_progress( - epoch=state.epoch, - step=state.global_step - ) - - def on_step_end(self, args, state, control, **kwargs): - if self.trainer_instance.should_stop: - print(f"Stop detected at step {state.global_step}\n") - control.should_training_stop = True - return control - - self.trainer.add_callback(ProgressCallback(self)) - - # Calculate total steps - num_samples = len(train_ds) - grad_accum = training_args.get('gradient_accumulation_steps', 4) - num_epochs = training_args.get('num_epochs', 3) - len_dataloader = math.ceil(num_samples / batch_size) - num_update_steps_per_epoch = max( - len_dataloader // grad_accum + int(len_dataloader % grad_accum > 0), 1 - ) - - if max_steps_val and max_steps_val > 0: - total_steps = max_steps_val - else: - total_steps = num_update_steps_per_epoch * num_epochs - - self._update_progress(total_steps=total_steps) - print(f"BiCodec progress tracking: {total_steps} total steps\n") - - # Train - self._update_progress(status_message="Starting BiCodec training...") - print("Starting BiCodec training...\n") + 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 - # Save - if self.should_stop and self.save_on_stop: - self.trainer.save_model() - self.tokenizer.save_pretrained(output_dir) - print(f"\nBiCodec training 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: - print("\nBiCodec training cancelled.\n") - self._update_progress( - is_training=False, - status_message="Training cancelled.", - ) - else: - self.trainer.save_model() - self.tokenizer.save_pretrained(output_dir) - print(f"\nBiCodec training 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}", - ) - return # Exit _train_worker for BiCodec - - elif self._audio_type is not None: - # Remaining audio types not yet implemented + 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 ========== @@ -2388,19 +2216,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, @@ -2410,6 +2237,7 @@ class UnslothTrainer: "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)), + "max_seq_length": training_args.get('max_seq_length', 2048), } # Add warmup parameter - use warmup_ratio if provided, otherwise warmup_steps @@ -2491,6 +2319,14 @@ 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") @@ -2555,7 +2391,7 @@ class UnslothTrainer: # DeepSeek OCR handles this internally in its collator, so skip # Audio VLM handles label masking in its collator, so skip - if train_on_responses_enabled and not self.is_audio_vlm and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'): + 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") @@ -2584,15 +2420,23 @@ 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 self.is_audio_vlm 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 + # After CUDA-heavy audio preprocessing (Whisper/DAC/BiCodec/SNAC), + # fork-based multiprocessing deadlocks because children inherit + # CUDA's internal thread locks. Use single-process mode instead. + toro_num_proc = config_args.get("dataset_num_proc", safe_num_proc(max(1, os.cpu_count() // 4))) + if _CUDA_AUDIO_PREPROCESSING_DONE: + toro_num_proc = 1 + print("Using single-process train_on_responses (CUDA audio preprocessing detected)\n") + self.trainer = train_on_responses_only( 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=toro_num_proc, ) print("Train on responses only configured successfully\n") @@ -2639,103 +2483,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(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 ========== @@ -2744,32 +2502,7 @@ class UnslothTrainer: self.trainer.train() # ========== SAVE MODEL ========== - if self.should_stop and self.save_on_stop: - # Stopped by user — save model at current checkpoint - self.trainer.save_model() - self.tokenizer.save_pretrained(output_dir) - print(f"\nTraining stopped. Model saved to {output_dir}\n") - self._update_progress( - is_training=False, - status_message=f"Training stopped. Model saved to {output_dir}", - ) - elif self.should_stop: - # Cancelled by user — don't save - print("\nTraining cancelled.\n") - self._update_progress( - is_training=False, - status_message="Training cancelled.", - ) - else: - # Normal completion - self.trainer.save_model() - self.tokenizer.save_pretrained(output_dir) - print(f"\nTraining completed! Model saved to {output_dir}\n") - self._update_progress( - is_training=False, - is_completed=True, - status_message=f"Training completed! Model saved to {output_dir}", - ) + self._finalize_training(output_dir) except Exception as e: import traceback 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/routes/inference.py b/studio/backend/routes/inference.py index 8d1c6667c0..3fa4e43da5 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__) @@ -184,6 +189,9 @@ async def load_model( is_vision=config.is_vision, is_lora=config.is_lora, is_gguf=False, + is_audio=config.is_audio, + audio_type=config.audio_type, + has_audio_input=config.has_audio_input, inference=inference_config, ) @@ -330,14 +338,23 @@ async def get_status( backend = get_inference_backend() is_vision = False + is_audio = False + audio_type = None + has_audio_input = False if backend.active_model_name: model_info = backend.models.get(backend.active_model_name, {}) is_vision = model_info.get("is_vision", False) + is_audio = model_info.get("is_audio", False) + audio_type = model_info.get("audio_type") + has_audio_input = model_info.get("has_audio_input", False) return InferenceStatusResponse( active_model=backend.active_model_name, is_vision=is_vision, is_gguf=False, + is_audio=is_audio, + audio_type=audio_type, + has_audio_input=has_audio_input, loading=list(getattr(backend, 'loading_models', set())), loaded=list(backend.models.keys()), ) @@ -350,11 +367,118 @@ async def get_status( ) +# ===================================================================== +# Audio (TTS) Generation (/audio/generate) +# ===================================================================== + + +@router.post("/audio/generate") +async def generate_audio(payload: ChatCompletionRequest, request: Request): + """ + Generate audio (TTS) from the latest user message. + Returns a JSON response with base64-encoded WAV audio. + Only works when an audio model is loaded. + """ + import base64 + + backend = get_inference_backend() + if not backend.active_model_name: + raise HTTPException(status_code=400, detail="No model loaded.") + + model_info = backend.models.get(backend.active_model_name, {}) + if not model_info.get("is_audio"): + raise HTTPException(status_code=400, detail="Active model is not an audio model.") + + # Extract text from the last user message + _, chat_messages, _ = _extract_content_parts(payload.messages) + if not chat_messages: + raise HTTPException(status_code=400, detail="No messages provided.") + + last_user_msg = next( + (m for m in reversed(chat_messages) if m["role"] == "user"), None + ) + if not last_user_msg: + raise HTTPException(status_code=400, detail="No user message found.") + + text = last_user_msg["content"] + + try: + wav_bytes, sample_rate = await asyncio.get_event_loop().run_in_executor( + None, + lambda: backend.generate_audio_response( + text=text, + temperature=payload.temperature, + top_p=payload.top_p, + top_k=payload.top_k, + min_p=payload.min_p, + max_new_tokens=payload.max_tokens or 2048, + repetition_penalty=payload.repetition_penalty, + use_adapter=payload.use_adapter, + ), + ) + + audio_b64 = base64.b64encode(wav_bytes).decode("ascii") + completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" + + return JSONResponse(content={ + "id": completion_id, + "object": "chat.completion.audio", + "model": backend.active_model_name, + "audio": { + "data": audio_b64, + "format": "wav", + "sample_rate": sample_rate, + }, + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": f"[Generated audio from: \"{text[:100]}\"]", + }, + "finish_reason": "stop", + }], + }) + + except Exception as e: + logger.error(f"Audio generation error: {e}", exc_info=True) + raise HTTPException(status_code=500, detail=str(e)) + + # ===================================================================== # OpenAI-Compatible Chat Completions (/chat/completions) # ===================================================================== +def _decode_audio_base64(b64: str) -> np.ndarray: + """Decode base64 audio (any format) → float32 numpy array at 16kHz.""" + import torch + import torchaudio + import tempfile + import os + + raw = base64.b64decode(b64) + # torchaudio.load needs a file path or file-like object with format hint + # Write to a temp file so torchaudio can auto-detect the format + with tempfile.NamedTemporaryFile(suffix=".audio", delete=False) as tmp: + tmp.write(raw) + tmp_path = tmp.name + try: + waveform, sr = torchaudio.load(tmp_path) + finally: + os.unlink(tmp_path) + + # Convert to mono if stereo + if waveform.shape[0] > 1: + waveform = waveform.mean(dim=0, keepdim=True) + + # Resample to 16kHz if needed + if sr != 16000: + resampler = torchaudio.transforms.Resample(orig_freq=sr, new_freq=16000) + waveform = resampler(waveform) + + return waveform.squeeze(0).numpy() + + def _extract_content_parts( messages: list, ) -> tuple[str, list[dict], "Optional[str]"]: @@ -444,6 +568,79 @@ async def openai_chat_completions( ) model_name = backend.active_model_name or payload.model + # ── Audio TTS path: auto-route to audio generation ──── + model_info = backend.models.get(backend.active_model_name, {}) + if model_info.get("is_audio"): + return await generate_audio(payload, request) + + # ── Audio INPUT path: decode WAV and route to audio input generation ── + if payload.audio_base64 and model_info.get("has_audio_input"): + audio_array = _decode_audio_base64(payload.audio_base64) + system_prompt, chat_messages, _ = _extract_content_parts(payload.messages) + cancel_event = threading.Event() + completion_id = f"chatcmpl-{uuid.uuid4().hex[:12]}" + created = int(time.time()) + + def audio_input_generate(): + return backend.generate_audio_input_response( + messages=chat_messages, + system_prompt=system_prompt, + audio_array=audio_array, + temperature=payload.temperature, + top_p=payload.top_p, + top_k=payload.top_k, + min_p=payload.min_p, + max_new_tokens=payload.max_tokens or 512, + repetition_penalty=payload.repetition_penalty, + cancel_event=cancel_event, + ) + + if payload.stream: + async def audio_input_stream(): + try: + first_chunk = ChatCompletionChunk( + id=completion_id, created=created, model=model_name, + choices=[ChunkChoice(delta=ChoiceDelta(role="assistant"), finish_reason=None)], + ) + yield f"data: {first_chunk.model_dump_json(exclude_none=True)}\n\n" + + for chunk_text in audio_input_generate(): + if await request.is_disconnected(): + cancel_event.set() + return + if chunk_text: + chunk = ChatCompletionChunk( + id=completion_id, created=created, model=model_name, + choices=[ChunkChoice(delta=ChoiceDelta(content=chunk_text), finish_reason=None)], + ) + yield f"data: {chunk.model_dump_json(exclude_none=True)}\n\n" + + final_chunk = ChatCompletionChunk( + id=completion_id, created=created, model=model_name, + choices=[ChunkChoice(delta=ChoiceDelta(), finish_reason="stop")], + ) + yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n" + yield "data: [DONE]\n\n" + except asyncio.CancelledError: + cancel_event.set() + raise + except Exception as e: + logger.error(f"Error during audio input streaming: {e}", exc_info=True) + yield f"data: {json.dumps({'error': {'message': str(e), 'type': 'server_error'}})}\n\n" + + return StreamingResponse( + audio_input_stream(), + media_type="text/event-stream", + headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"}, + ) + else: + full_text = "".join(audio_input_generate()) + response = ChatCompletion( + id=completion_id, created=created, model=model_name, + choices=[CompletionChoice(message=CompletionMessage(content=full_text), finish_reason="stop")], + ) + return JSONResponse(content=response.model_dump()) + # ── Parse messages (handles multimodal content parts) ───── system_prompt, chat_messages, extracted_image_b64 = _extract_content_parts( payload.messages diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 2c7f0e846f..8072c7e1f3 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -225,7 +225,10 @@ async def list_models( id=model_name, name=model_name.split("/")[-1] if "/" in model_name else model_name, is_vision=model_data.get("is_vision", False), - is_lora=model_data.get("is_lora", False) + is_lora=model_data.get("is_lora", False), + is_audio=model_data.get("is_audio", False), + audio_type=model_data.get("audio_type"), + has_audio_input=model_data.get("has_audio_input", False), ) loaded_models.append(model_info) diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 3dde2b3a63..6144bbe09c 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -220,6 +220,11 @@ MODEL_NAME_MAPPING = { ], "OuteAI_Llama-OuteTTS-1.0-1B.yaml": [ "OuteAI/Llama-OuteTTS-1.0-1B", + "unsloth/Llama-OuteTTS-1.0-1B", + "unsloth/llama-outetts-1.0-1b", + "OuteAI/OuteTTS-1.0-0.6B", + "unsloth/OuteTTS-1.0-0.6B", + "unsloth/outetts-1.0-0.6b", ], "unsloth_PaddleOCR-VL.yaml": [ "unsloth/PaddleOCR-VL", @@ -402,6 +407,10 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: logger.info(f"Model {model_name} detected as VLM: has img_processor") return True + # Check 4: Exclude audio models that have ForConditionalGeneration but aren't VLMs + # (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration) + # These are handled by is_audio_model() instead + # Check 4: Has image_token_index (common in VLMs for image placeholder tokens) if hasattr(config, 'image_token_index'): logger.info(f"Model {model_name} detected as VLM: has image_token_index") @@ -425,6 +434,38 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: pass +def is_audio_model(model_name: str) -> Optional[str]: + """ + Check if a model is a TTS audio model by looking up its YAML config. + + Returns the audio_type string ('snac', 'csm', 'bicodec', 'dac') or None. + """ + try: + defaults = load_model_defaults(model_name) + audio_type = defaults.get('audio_type') + if audio_type and isinstance(audio_type, str) and audio_type in ('snac', 'csm', 'bicodec', 'dac'): + logger.info(f"Model {model_name} detected as audio model: audio_type={audio_type}") + return audio_type + return None + except Exception as e: + logger.debug(f"Could not determine if {model_name} is audio model: {e}") + return None + + +def has_audio_input_model(model_name: str) -> bool: + """ + Check if a model accepts audio input (ASR/speech understanding) by looking up its YAML config. + + Returns True if the model has 'audio_input: true' in its defaults. + """ + try: + defaults = load_model_defaults(model_name) + return bool(defaults.get('audio_input')) + except Exception as e: + logger.debug(f"Could not determine if {model_name} has audio input: {e}") + return False + + def detect_gguf_model(path: str) -> Optional[str]: """ Check if the given local path is or contains a GGUF model file. @@ -909,6 +950,9 @@ class ModelConfig: is_vision: bool # Is this a vision model? is_lora: bool # Is this a lora adapter? is_gguf: bool = False # Is this a GGUF model? + is_audio: bool = False # Is this a TTS audio model? + audio_type: Optional[str] = None # Audio codec type: 'snac', 'csm', 'bicodec', 'dac' + has_audio_input: bool = False # Accepts audio input (ASR/speech understanding) gguf_file: Optional[str] = None # Full path to the .gguf file (local mode) gguf_hf_repo: Optional[str] = None # HF repo ID for -hf mode (e.g. "unsloth/gemma-3-4b-it-GGUF") gguf_variant: Optional[str] = None # Quantization variant (e.g. "Q4_K_M") @@ -944,6 +988,9 @@ class ModelConfig: # Check if base model is vision is_vision = is_vision_model(base_model, hf_token=hf_token) + # Check if base model is audio + audio_type = is_audio_model(base_model) + display_name = lora_path_obj.name identifier = lora_path # Use path as identifier for local LoRAs @@ -955,6 +1002,8 @@ class ModelConfig: is_cached=True, # Local LoRAs are always "cached" is_vision=is_vision, is_lora=True, + is_audio=audio_type is not None, + audio_type=audio_type, base_model=base_model, ) @@ -1105,11 +1154,15 @@ class ModelConfig: logger.warning(f"Could not determine base model for LoRA '{path}'") return None vision = is_vision_model(base_model, hf_token=hf_token) + audio_type_val = is_audio_model(base_model) + has_audio_in = has_audio_input_model(base_model) else: vision = is_vision_model(identifier, hf_token=hf_token) - + audio_type_val = is_audio_model(identifier) + has_audio_in = has_audio_input_model(identifier) + display_name = Path(path).name if is_local else identifier.split("/")[-1] - + return cls( identifier=identifier, display_name=display_name, @@ -1118,6 +1171,9 @@ class ModelConfig: is_cached=is_model_cached(identifier) if not is_local else True, is_vision=vision, is_lora=is_lora, + is_audio=audio_type_val is not None, + audio_type=audio_type_val, + has_audio_input=has_audio_in, base_model=base_model, ) 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 +165,50 @@ const ComposerAnimated: FC = () => { ); }; +const AUDIO_ACCEPT = "audio/wav,audio/mpeg,audio/webm,audio/ogg,audio/flac,audio/mp4"; +const MAX_AUDIO_SIZE = 50 * 1024 * 1024; + +function fileToBase64(file: File): Promise { + return new Promise((resolve, reject) => { + const reader = new FileReader(); + reader.onload = () => { + const result = reader.result as string; + const commaIndex = result.indexOf(","); + resolve(commaIndex >= 0 ? result.slice(commaIndex + 1) : result); + }; + reader.onerror = () => reject(new Error("Failed to read file")); + reader.readAsDataURL(file); + }); +} + +const PendingAudioChip: FC = () => { + const audioName = useChatRuntimeStore((s) => s.pendingAudioName); + const clearPendingAudio = useChatRuntimeStore((s) => s.clearPendingAudio); + if (!audioName) return null; + return ( +
+
+ + {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 (
- +
+ + +
diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index af8f10156f..1c2d46f2db 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 { @@ -92,6 +92,25 @@ function findLatestUserImageBase64(messages: RunMessages): string | undefined { return undefined; } +function findLatestUserAudioBase64(messages: RunMessages): string | undefined { + // Check message content parts (from compare view's CompareMessagePart with type: "audio") + for (let i = messages.length - 1; i >= 0; i -= 1) { + const message = messages[i]; + if (!message || message.role !== "user") continue; + + for (const part of message.content ?? []) { + if (part.type === "audio" && "audio" in part) { + const raw = (part as { type: "audio"; audio: string }).audio; + if (raw) return raw.startsWith("data:") ? raw.split(",")[1] : raw; + } + } + } + + // Check the runtime store (from main composer's audio upload) + const pendingAudio = useChatRuntimeStore.getState().pendingAudioBase64; + return pendingAudio ?? undefined; +} + async function resolveUseAdapter( threadId: string | undefined, ): Promise { @@ -135,8 +154,64 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { }); } const imageBase64 = findLatestUserImageBase64(messages); + const audioBase64 = findLatestUserAudioBase64(messages); + // Clear pending audio from store after extracting (consumed on send) + if (audioBase64) { + runtime.clearPendingAudio(); + } const useAdapter = await resolveUseAdapter(unstable_threadId); + // ── Audio model path (non-streaming) ───────────────────── + const activeModel = runtime.models.find( + (m) => m.id === params.checkpoint, + ); + if (activeModel?.isAudio) { + const threadKey = unstable_threadId || "__default"; + runtime.setThreadRunning(threadKey, true); + try { + yield { + content: [{ type: "text" as const, text: "Generating audio..." }], + }; + + const result = await generateAudio( + { + model: params.checkpoint, + messages: outboundMessages, + stream: false, + temperature: params.temperature, + top_p: params.topP, + max_tokens: params.maxTokens, + top_k: params.topK, + min_p: params.minP, + repetition_penalty: params.repetitionPenalty, + ...(useAdapter === undefined ? {} : { use_adapter: useAdapter }), + }, + abortSignal, + ); + + const audioUrl = `data:audio/wav;base64,${result.audio.data}`; + yield { + content: [ + { + type: "text" as const, + text: ``, + }, + ], + }; + } catch (err) { + if (!abortSignal.aborted) { + toast.error("Audio generation failed", { + description: + err instanceof Error ? err.message : "Unknown error", + }); + } + throw err; + } finally { + runtime.setThreadRunning(threadKey, false); + } + return; + } + const threadKey = unstable_threadId || "__default"; let waitingFirstChunk = true; let firstTokenSettled = false; @@ -194,6 +269,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter { min_p: params.minP, repetition_penalty: params.repetitionPenalty, image_base64: imageBase64, + audio_base64: audioBase64, ...(useAdapter === undefined ? {} : { use_adapter: useAdapter }), }, abortSignal, diff --git a/studio/frontend/src/features/chat/api/chat-api.ts b/studio/frontend/src/features/chat/api/chat-api.ts index 5d5a9551ef..0720141349 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..ea5be5ea3f 100644 --- a/studio/frontend/src/features/chat/shared-composer.tsx +++ b/studio/frontend/src/features/chat/shared-composer.tsx @@ -1,7 +1,8 @@ import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button"; import { Button } from "@/components/ui/button"; import { useAui } from "@assistant-ui/react"; -import { ArrowUpIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react"; +import { ArrowUpIcon, HeadphonesIcon, MicIcon, PlusIcon, SquareIcon, XIcon } from "lucide-react"; +import { useChatRuntimeStore } from "./stores/chat-runtime-store"; import { type KeyboardEvent, type MutableRefObject, @@ -17,7 +18,8 @@ import { export type CompareMessagePart = | { type: "text"; text: string } - | { type: "image"; image: string }; + | { type: "image"; image: string } + | { type: "audio"; audio: string }; export interface CompareHandle { append: (content: CompareMessagePart[]) => void; @@ -27,6 +29,8 @@ export interface CompareHandle { const IMAGE_ACCEPT = "image/jpeg,image/png,image/webp,image/gif"; const MAX_IMAGE_SIZE = 20 * 1024 * 1024; +const AUDIO_ACCEPT = "audio/wav,audio/mpeg,audio/webm,audio/ogg,audio/flac,audio/mp4"; +const MAX_AUDIO_SIZE = 50 * 1024 * 1024; function fileToBase64DataURL(file: File): Promise { return new Promise((resolve, reject) => { @@ -182,9 +186,18 @@ export function SharedComposer({ const [text, setText] = useState(""); const [running, setRunning] = useState(false); const [pendingImages, setPendingImages] = useState([]); + 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 +217,27 @@ export function SharedComposer({ const next: PendingImage[] = []; for (let i = 0; i < files.length; i++) { const file = files[i]; - if (!file?.type.match(/^image\/(jpeg|png|webp|gif)$/i)) continue; + if (!file) continue; + // Handle audio files + if (file.type.match(/^audio\//i) && file.size <= MAX_AUDIO_SIZE) { + const reader = new FileReader(); + reader.onload = () => { + const result = reader.result as string; + const commaIndex = result.indexOf(","); + const base64 = commaIndex >= 0 ? result.slice(commaIndex + 1) : result; + setPendingAudio({ name: file.name, base64 }); + setPendingAudioStore(base64, file.name); + }; + reader.readAsDataURL(file); + continue; + } + // Handle image files + if (!file.type.match(/^image\/(jpeg|png|webp|gif)$/i)) continue; if (file.size > MAX_IMAGE_SIZE) continue; next.push({ id: crypto.randomUUID(), file }); } setPendingImages((prev) => [...prev, ...next]); - }, []); + }, [setPendingAudioStore]); const removePendingImage = useCallback((id: string) => { setPendingImages((prev) => prev.filter((p) => p.id !== id)); @@ -217,7 +245,7 @@ export function SharedComposer({ async function send() { const msg = text.trim(); - if (!msg && pendingImages.length === 0) return; + if (!msg && pendingImages.length === 0 && !pendingAudio) return; const content: CompareMessagePart[] = []; for (const { file } of pendingImages) { @@ -228,6 +256,9 @@ export function SharedComposer({ // skip failed image } } + if (pendingAudio) { + content.push({ type: "audio", audio: pendingAudio.base64 }); + } if (msg) { content.push({ type: "text", text: msg }); } @@ -238,6 +269,8 @@ export function SharedComposer({ } setText(""); setPendingImages([]); + setPendingAudio(null); + clearPendingAudioStore(); textareaRef.current?.focus(); } @@ -257,7 +290,7 @@ export function SharedComposer({ } } - const canSend = (text.trim().length > 0 || pendingImages.length > 0) && !running; + const canSend = (text.trim().length > 0 || pendingImages.length > 0 || pendingAudio !== null) && !running; return (
- {pendingImages.length > 0 && ( + {(pendingImages.length > 0 || pendingAudio) && (
{pendingImages.map(({ id, file }) => ( removePendingImage(id)} /> ))} + {pendingAudio && ( +
+ + {pendingAudio.name} + +
+ )}
)}