diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 8208f71694..e2d03e53b5 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -22,7 +22,7 @@ from dataclasses import dataclass import pandas as pd from datasets import Dataset, load_dataset -from utils.models import is_vision_model +from utils.models import is_vision_model, detect_audio_type from utils.datasets import format_and_template_dataset from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER from trl import SFTTrainer, SFTConfig @@ -69,7 +69,7 @@ class UnslothTrainer: self.is_vlm = False self.is_audio = False self.is_audio_vlm = False # Multimodal model (e.g. Gemma 3N) trained on audio data - self._audio_type = None # 'csm', 'whisper', 'snac', 'xcodec2', 'bicodec', 'dac' + self._audio_type = None # 'csm', 'whisper', 'snac', 'bicodec', 'dac' self._cuda_audio_used = False # Set once after audio CUDA preprocessing; never cleared self._spark_tts_repo_dir = None # Path to downloaded Spark-TTS repo (for BiCodecTokenizer) self.model_name = None @@ -314,17 +314,6 @@ class UnslothTrainer: return {"audio_col": audio_col, "text_col": text_col, "speaker_col": speaker_col} - def _resolve_audio_type(self, model_name: str) -> Optional[str]: - """Resolve audio_type from YAML model config. Returns None for non-audio models.""" - try: - from utils.models.model_config import load_model_defaults - defaults = load_model_defaults(model_name) - audio_type = defaults.get('audio_type') - if audio_type and isinstance(audio_type, str): - return audio_type - except Exception as e: - logger.warning(f"Could not resolve audio_type for {model_name}: {e}") - return None def load_model(self, model_name: str, @@ -369,23 +358,26 @@ class UnslothTrainer: # Remove stale compiled cache so the new model gets a fresh one from utils.cache_cleanup import clear_unsloth_compiled_cache clear_unsloth_compiled_cache() - # Detect audio model type from YAML config - self._audio_type = self._resolve_audio_type(model_name) - self.is_audio = self._audio_type is not None + # Detect audio model type dynamically (config.json + tokenizer) + self._audio_type = detect_audio_type(model_name, hf_token) + # audio_vlm is detected as an audio_type now, handle it separately + if self._audio_type == 'audio_vlm': + self.is_audio = False + self.is_audio_vlm = is_dataset_audio # Only use audio VLM path if dataset has audio + self._audio_type = None + else: + self.is_audio = self._audio_type is not None + self.is_audio_vlm = False - # Audio VLM: multimodal model (e.g. Gemma 3N) trained on audio data - # Uses FastModel + SFTTrainer with audio collator - self.is_audio_vlm = not self.is_audio and is_vision_model(model_name) and is_dataset_audio - # 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_image + # VLM: vision model with image dataset (mutually exclusive with audio paths) + vision = is_vision_model(model_name) if not self.is_audio else False + self.is_vlm = not self.is_audio_vlm and vision and is_dataset_image self.model_name = model_name self.max_seq_length = max_seq_length - logger.info(f"Audio type: {self._audio_type}") - if not self.is_audio: - logger.info(f"Model architecture is vision: {is_vision_model(model_name)}") + logger.info(f"Audio type: {self._audio_type}, is_audio: {self.is_audio}, is_audio_vlm: {self.is_audio_vlm}") logger.info(f"Dataset has images: {is_dataset_image}, audio: {is_dataset_audio}") - logger.info(f"Using VLM path: {self.is_vlm}, Audio VLM: {self.is_audio_vlm}") + logger.info(f"Using VLM path: {self.is_vlm}") # Reset training state for new run self._update_progress( @@ -947,6 +939,11 @@ class UnslothTrainer: processor = AutoProcessor.from_pretrained(self.model_name) + # Strip pad_to_multiple_of from tokenizer init_kwargs — fine-tuned models + # (e.g. keanteng/sesame-csm-elise) save it in tokenizer_config.json, and + # _merge_kwargs leaks it into audio_kwargs where EncodecFeatureExtractor rejects it. + processor.tokenizer.init_kwargs.pop('pad_to_multiple_of', None) + # Resolve columns from user mapping or hardcoded fallback resolved = self._resolve_audio_columns(dataset, custom_format_mapping) audio_col = resolved["audio_col"] @@ -966,15 +963,27 @@ class UnslothTrainer: dataset = dataset.cast_column(audio_col, Audio(sampling_rate=24000)) - def preprocess_example(example): - conversation = [{ - "role": str(example[speaker_key]), - "content": [ - {"type": "text", "text": example.get(text_col, "")}, - {"type": "audio", "path": example[audio_col]["array"]}, - ], - }] + required_keys = ["input_ids", "attention_mask", "labels", "input_values", "input_values_cutoffs"] + + self._update_progress(status_message="Preprocessing CSM dataset...") + processed_examples = [] + skipped = 0 + for idx in range(len(dataset)): + if self.should_stop: + print("Stopped during CSM preprocessing\n") + break + + example = dataset[idx] try: + conversation = [{ + "role": str(example[speaker_key]), + "content": [ + {"type": "text", "text": example.get(text_col, "")}, + {"type": "audio", "path": example[audio_col]["array"]}, + ], + }] + # NOTE: pad_to_multiple_of intentionally omitted from text_kwargs — + # CsmProcessor._merge_kwargs leaks it to EncodecFeatureExtractor which rejects it. model_inputs = processor.apply_chat_template( conversation, tokenize=True, @@ -983,7 +992,6 @@ class UnslothTrainer: text_kwargs={ "padding": "max_length", "max_length": 256, - "pad_to_multiple_of": 8, "padding_side": "right", }, audio_kwargs={ @@ -993,29 +1001,38 @@ class UnslothTrainer: }, common_kwargs={"return_tensors": "pt"}, ) + + out = {} + for k in required_keys: + if k not in model_inputs: + raise KeyError(f"Missing required key '{k}' in model outputs") + out[k] = model_inputs[k][0] + + if not all(isinstance(out[k], torch.Tensor) for k in out): + skipped += 1 + continue + + processed_examples.append(out) + except Exception as e: - logger.warning(f"Error processing CSM example: {e}") - return None + logger.warning(f"Error processing CSM example {idx}: {e}") + skipped += 1 + continue - required = ["input_ids", "attention_mask", "labels", "input_values", "input_values_cutoffs"] - out = {} - for k in required: - if k not in model_inputs: - return None - out[k] = model_inputs[k][0] + if (idx + 1) % 100 == 0: + self._update_progress( + status_message=f"Preprocessing CSM... {idx + 1}/{len(dataset)}" + ) - if not all(isinstance(out[k], torch.Tensor) for k in out): - return None - return out + if not processed_examples: + raise ValueError( + f"No valid examples after CSM preprocessing (skipped {skipped})" + ) - self._update_progress(status_message="Preprocessing CSM dataset...") - processed = dataset.map( - preprocess_example, - remove_columns=dataset.column_names, - desc="Preprocessing CSM dataset", - ) - print(f"CSM preprocessing complete: {len(processed)} examples\n") - return processed + result_dataset = Dataset.from_list(processed_examples) + print(f"CSM preprocessing complete: {len(result_dataset)} examples " + f"({skipped} skipped)\n") + return result_dataset def _format_audio_vlm_dataset(self, dataset, custom_format_mapping=None): """Format dataset as audio chat messages for multimodal models (e.g. Gemma 3N). @@ -1100,7 +1117,11 @@ class UnslothTrainer: f"SNAC dataset needs 'audio' and 'text' columns, got: {dataset.column_names}" ) - # Get dataset sample rate from first example + # Cast audio column so datasets 4.x AudioDecoder objects are decoded to dicts + from datasets import Audio + dataset = dataset.cast_column(audio_col, Audio(sampling_rate=SNAC_SAMPLE_RATE)) + + # Get dataset sample rate from first example (after cast, always SNAC_SAMPLE_RATE) first_audio = dataset[0][audio_col] ds_sample_rate = first_audio.get("sampling_rate", SNAC_SAMPLE_RATE) if isinstance(first_audio, dict) else SNAC_SAMPLE_RATE @@ -1272,6 +1293,11 @@ class UnslothTrainer: f"BiCodec dataset needs 'audio' and 'text' columns, got: {dataset.column_names}" ) + # Cast audio column so datasets 4.x AudioDecoder objects are decoded to dicts. + # Don't resample here — BiCodec's target_sr may differ; the loop handles resampling. + from datasets import Audio + dataset = dataset.cast_column(audio_col, Audio()) + # Load BiCodec tokenizer self._update_progress(status_message="Loading BiCodec tokenizer...") print("Loading BiCodec tokenizer...\n") @@ -1810,9 +1836,6 @@ class UnslothTrainer: 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: formatted = self._format_audio_vlm_dataset(dataset, custom_format_mapping) return (formatted, None) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 8072c7e1f3..d3a93a5d75 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -26,7 +26,7 @@ try: list_gguf_variants, ModelConfig, ) - from utils.models.model_config import _pick_best_gguf, _extract_quant_label + from utils.models.model_config import _pick_best_gguf, _extract_quant_label, is_audio_input_type from core.inference import get_inference_backend except ImportError: # Fallback: try to import from parent directory @@ -43,7 +43,7 @@ except ImportError: list_gguf_variants, ModelConfig, ) - from utils.models.model_config import _pick_best_gguf, _extract_quant_label + from utils.models.model_config import _pick_best_gguf, _extract_quant_label, is_audio_input_type from core.inference import get_inference_backend from models import ( @@ -277,32 +277,34 @@ async def get_model_config( """ try: logger.info(f"Getting model config for: {model_name}") + from utils.models.model_config import detect_audio_type # Load model defaults from backend config_dict = load_model_defaults(model_name) - - # Check if it's a vision model + + # Detect model capabilities is_vision = is_vision_model(model_name) - + audio_type = detect_audio_type(model_name) + # Check if it's a LoRA adapter is_lora = False base_model = None - - # Try to create ModelConfig to get more info try: model_config = ModelConfig.from_identifier(model_name) is_lora = model_config.is_lora base_model = model_config.base_model if is_lora else None except Exception: - # If ModelConfig creation fails, use defaults pass - - logger.info(f"Model config result for {model_name}: is_vision={is_vision}, is_lora={is_lora}, base_model={base_model}") + + logger.info(f"Model config result for {model_name}: is_vision={is_vision}, audio_type={audio_type}, is_lora={is_lora}") return ModelDetails( id=model_name, model_name=model_name, config=config_dict, is_vision=is_vision, is_lora=is_lora, + is_audio=audio_type is not None, + audio_type=audio_type, + has_audio_input=is_audio_input_type(audio_type), base_model=base_model, ) diff --git a/studio/backend/utils/models/__init__.py b/studio/backend/utils/models/__init__.py index 92e65cf67c..11d5f54539 100644 --- a/studio/backend/utils/models/__init__.py +++ b/studio/backend/utils/models/__init__.py @@ -5,6 +5,9 @@ from .model_config import ( ModelConfig, GgufVariantInfo, is_vision_model, + detect_audio_type, + is_audio_input_type, + VALID_AUDIO_TYPES, scan_trained_loras, scan_exported_models, load_model_defaults, @@ -20,6 +23,9 @@ __all__ = [ 'ModelConfig', 'GgufVariantInfo', 'is_vision_model', + 'detect_audio_type', + 'is_audio_input_type', + 'VALID_AUDIO_TYPES', 'scan_trained_loras', 'scan_exported_models', 'load_model_defaults', diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index 995ff084cb..222017ce00 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -387,6 +387,13 @@ def is_vision_model(model_name: str, hf_token: Optional[str] = None) -> bool: try: config = load_model_config(model_name, use_auth=True, token=hf_token) + # Exclude audio-only models that share ForConditionalGeneration suffix + # (e.g. CsmForConditionalGeneration, WhisperForConditionalGeneration) + _audio_only_model_types = {'csm', 'whisper'} + model_type = getattr(config, 'model_type', None) + if model_type in _audio_only_model_types: + return False + # Check 1: Architecture class name patterns if hasattr(config, 'architectures'): is_vlm = any( @@ -407,10 +414,6 @@ 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") @@ -434,36 +437,115 @@ 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. +VALID_AUDIO_TYPES = ('snac', 'csm', 'bicodec', 'dac', 'whisper', 'audio_vlm') - Returns the audio_type string ('snac', 'csm', 'bicodec', 'dac') or None. +# Cache detection results per session to avoid repeated API calls +_audio_detection_cache: Dict[str, Optional[str]] = {} + +# Tokenizer token patterns → audio_type (all 6 types detected from tokenizer_config.json) +_AUDIO_TOKEN_PATTERNS = { + 'csm': lambda tokens: '<|AUDIO|>' in tokens and '<|audio_eos|>' in tokens, + 'whisper': lambda tokens: '<|startoftranscript|>' in tokens, + 'audio_vlm': lambda tokens: '' in tokens, + 'bicodec': lambda tokens: any(t.startswith('<|bicodec_') for t in tokens), + 'dac': lambda tokens: '<|audio_start|>' in tokens and '<|audio_end|>' in tokens, + 'snac': lambda tokens: sum(1 for t in tokens if t.startswith(' 10000, +} + + +def detect_audio_type(model_name: str, hf_token: Optional[str] = None) -> Optional[str]: """ + Dynamically detect if a model is an audio model and return its type. + + Fully dynamic — works for any model, not just known ones. + Uses tokenizer_config.json special tokens to detect all 6 audio types. + + Returns: audio_type string ('snac', 'csm', 'bicodec', 'dac', 'whisper', 'audio_vlm') or None. + """ + if model_name in _audio_detection_cache: + return _audio_detection_cache[model_name] + + result = _detect_audio_from_tokenizer(model_name, hf_token) + + _audio_detection_cache[model_name] = result + if result: + logger.info(f"Model {model_name} detected as audio model: audio_type={result}") + return result + + +def _detect_audio_from_tokenizer(model_name: str, hf_token: Optional[str] = None) -> Optional[str]: + """Detect audio type from tokenizer special tokens (for LLM-based audio models). + + First checks local HF cache, then fetches tokenizer_config.json from HuggingFace. + Checks added_tokens_decoder for distinctive patterns. + """ + def _check_token_patterns(tok_config: dict) -> Optional[str]: + added = tok_config.get('added_tokens_decoder', {}) + if not added: + return None + token_contents = [v.get('content', '') for v in added.values()] + for audio_type, check_fn in _AUDIO_TOKEN_PATTERNS.items(): + if check_fn(token_contents): + return audio_type + return None + + # 1) Check local HF cache first (works for gated/offline models) try: - 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', 'whisper'): - logger.info(f"Model {model_name} detected as audio model: audio_type={audio_type}") - return audio_type + from huggingface_hub.constants import HF_HUB_CACHE + cache_dir = Path(HF_HUB_CACHE) + repo_dir_name = f"models--{model_name.replace('/', '--')}" + repo_dir = cache_dir / repo_dir_name + if repo_dir.exists(): + snapshots_dir = repo_dir / "snapshots" + if snapshots_dir.exists(): + for snapshot in snapshots_dir.iterdir(): + for tok_path in ['tokenizer_config.json', 'LLM/tokenizer_config.json']: + tok_file = snapshot / tok_path + if tok_file.exists(): + tok_config = json.loads(tok_file.read_text()) + result = _check_token_patterns(tok_config) + if result: + return result + except Exception as e: + logger.debug(f"Could not check local cache for {model_name}: {e}") + + # 2) Fall back to HuggingFace API + try: + import requests + import os + + paths_to_try = ['tokenizer_config.json', 'LLM/tokenizer_config.json'] + # Use provided token, or fall back to env + token = hf_token or os.environ.get('HF_TOKEN') + headers = {} + if token: + headers['Authorization'] = f'Bearer {token}' + + for tok_path in paths_to_try: + url = f"https://huggingface.co/{model_name}/resolve/main/{tok_path}" + resp = requests.get(url, headers=headers, timeout=15) + if not resp.ok: + continue + + tok_config = resp.json() + result = _check_token_patterns(tok_config) + if result: + return result + return None except Exception as e: - logger.debug(f"Could not determine if {model_name} is audio model: {e}") + logger.debug(f"Could not detect audio type from tokenizer for {model_name}: {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. +def is_audio_input_type(audio_type: Optional[str]) -> bool: + """Check if an audio_type accepts audio input (ASR/speech understanding). - Returns True if the model has 'audio_input: true' in its defaults. + Whisper (ASR) and audio_vlm (Gemma3n) accept audio input. """ - 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 + return audio_type in ('whisper', 'audio_vlm') + + def _is_mmproj(filename: str) -> bool: """Check if a GGUF filename is a vision projection (mmproj) file.""" return "mmproj" in filename.lower() @@ -1040,7 +1122,7 @@ class ModelConfig: is_vision = is_vision_model(base_model, hf_token=hf_token) # Check if base model is audio - audio_type = is_audio_model(base_model) + audio_type = detect_audio_type(base_model, hf_token=hf_token) display_name = lora_path_obj.name identifier = lora_path # Use path as identifier for local LoRAs @@ -1055,6 +1137,7 @@ class ModelConfig: is_lora=True, is_audio=audio_type is not None, audio_type=audio_type, + has_audio_input=is_audio_input_type(audio_type), base_model=base_model, ) @@ -1232,22 +1315,16 @@ class ModelConfig: if not base_model: logger.warning(f"Could not determine base model for LoRA '{path}'") return None - vision = is_vision_model(base_model, hf_token=hf_token) - audio_type_val = is_audio_model(base_model) - has_audio_in = has_audio_input_model(base_model) + check_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) + check_model = identifier + + vision = is_vision_model(check_model, hf_token=hf_token) + audio_type_val = detect_audio_type(check_model, hf_token=hf_token) + has_audio_in = is_audio_input_type(audio_type_val) display_name = Path(path).name if is_local else identifier.split("/")[-1] - # Audio models are never vision models (e.g. WhisperForConditionalGeneration - # and CsmForConditionalGeneration match the ForConditionalGeneration suffix - # but are not VLMs). - if audio_type_val is not None: - vision = False - return cls( identifier=identifier, display_name=display_name, @@ -1338,4 +1415,3 @@ class ModelConfig: is_lora=is_lora, base_model=base_model, # This will be None for base models, and populated for LoRAs ) - pass