From 4d06258e9348979ed0e3762c85cf075a880f1bc2 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 22 Feb 2026 18:29:40 +0000 Subject: [PATCH 01/83] Auto-switch transformers version (5.1.0/4.57.1) for Ministral-3, GLM-4.7-Flash, Qwen3-30B-A3B models with LoRA adapter resolution --- studio/backend/routes/export.py | 4 + studio/backend/routes/inference.py | 4 + studio/backend/routes/training.py | 5 + studio/backend/utils/transformers_version.py | 184 +++++++++++++++++++ 4 files changed, 197 insertions(+) create mode 100644 studio/backend/utils/transformers_version.py diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index 6616c9fbd8..11d6377480 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -61,6 +61,10 @@ async def load_checkpoint( Wraps ExportBackend.load_checkpoint. """ try: + # Ensure correct transformers version for this model architecture + from utils.transformers_version import ensure_transformers_version + ensure_transformers_version(request.checkpoint_path) + backend = get_export_backend() success, message = backend.load_checkpoint( checkpoint_path=request.checkpoint_path, diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index c30a1638d4..c238b4019c 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -72,6 +72,10 @@ async def load_model(request: LoadRequest): from the model's YAML config, falling back to default.yaml for missing values. """ try: + # Ensure correct transformers version for this model architecture + from utils.transformers_version import ensure_transformers_version + ensure_transformers_version(request.model_path) + backend = get_inference_backend() # Create config using clean factory method diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index f8de2f639f..6917bed39e 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -85,6 +85,11 @@ async def start_training( """ try: logger.info(f"Starting training job with model: {request.model_name}") + + # Ensure correct transformers version for this model architecture + from utils.transformers_version import ensure_transformers_version + ensure_transformers_version(request.model_name) + backend = get_training_backend() # Generate job ID and attach to backend for later status/progress calls diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py new file mode 100644 index 0000000000..384a408d30 --- /dev/null +++ b/studio/backend/utils/transformers_version.py @@ -0,0 +1,184 @@ +""" +Automatic transformers version switching. + +Some newer model architectures (Ministral-3, GLM-4.7-Flash, Qwen3-30B-A3B MoE, +tiny_qwen3_moe) require transformers>=5.1.0, while everything else needs the +default 4.57.x that ships with Unsloth. + +When loading a LoRA adapter with a custom name, we resolve the base model from +``adapter_config.json`` and check *that* against the model list. + +This module detects the model being loaded and ensures the correct transformers +version is installed before proceeding. The pip install is done *without* +``--force-reinstall`` to avoid rebuilding unrelated dependencies. +""" + +import importlib +import importlib.metadata +import json +import logging +import os +import subprocess +import sys +from pathlib import Path + +logger = logging.getLogger(__name__) + +# --------------------------------------------------------------------------- +# Detection +# --------------------------------------------------------------------------- + +# Lowercase substrings — if ANY appears anywhere in the lowered model name, +# we need transformers 5.x. +TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = ( + "ministral-3-", # Ministral-3-{3,8,14}B-{Instruct,Reasoning,Base}-2512 + "glm-4.7-flash", # GLM-4.7-Flash + "qwen3-30b-a3b", # Qwen3-30B-A3B-Instruct-2507 and variants + "tiny_qwen3_moe", # imdatta0/tiny_qwen3_moe_2.8B_0.7B +) + +# Versions +TRANSFORMERS_5_VERSION = "5.1.0" +TRANSFORMERS_DEFAULT_VERSION = "4.57.1" + + +def _resolve_base_model(model_name: str) -> str: + """If *model_name* points to a LoRA adapter, return its base model. + + Checks for ``adapter_config.json`` locally first (covers the common case of + local output directories with custom names). Falls back to the existing + ``get_base_model_from_lora`` utility which also handles remote HF LoRAs. + + Returns the original *model_name* unchanged if it is not a LoRA adapter. + """ + # --- Fast local check --------------------------------------------------- + adapter_cfg_path = Path(model_name) / "adapter_config.json" + if adapter_cfg_path.is_file(): + try: + with open(adapter_cfg_path) as f: + cfg = json.load(f) + base = cfg.get("base_model_name_or_path") + if base: + logger.info( + "Resolved LoRA adapter '%s' → base model '%s'", model_name, base, + ) + return base + except Exception as exc: + logger.debug("Could not read %s: %s", adapter_cfg_path, exc) + + # --- Fallback: use the project's existing helper (handles HF repos too) -- + try: + from utils.models import get_base_model_from_lora + base = get_base_model_from_lora(model_name) + if base: + logger.info( + "Resolved LoRA adapter '%s' → base model '%s' (via get_base_model_from_lora)", + model_name, base, + ) + return base + except Exception as exc: + logger.debug("get_base_model_from_lora failed for '%s': %s", model_name, exc) + + return model_name + + +def needs_transformers_5(model_name: str) -> bool: + """Return True if *model_name* belongs to an architecture that requires + ``transformers>=5.1.0``.""" + lowered = model_name.lower() + return any(sub in lowered for sub in TRANSFORMERS_5_MODEL_SUBSTRINGS) + + +# --------------------------------------------------------------------------- +# Version switching +# --------------------------------------------------------------------------- + +def _installed_transformers_version() -> str | None: + """Return the currently installed transformers version, or None.""" + try: + return importlib.metadata.version("transformers") + except importlib.metadata.PackageNotFoundError: + return None + + +def _pip_install(package_spec: str) -> None: + """Run ``pip install `` using the running interpreter's pip. + + Mirrors the approach used in ``unsloth_zoo.llama_cpp.check_pip`` — + ``sys.executable -m pip`` is always the safest choice. + """ + cmd = [sys.executable, "-m", "pip", "install", package_spec] + logger.info("Running: %s", " ".join(cmd)) + result = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + if result.returncode != 0: + logger.error("pip install failed:\n%s", result.stdout) + raise RuntimeError( + f"Failed to install {package_spec}. " + f"pip output:\n{result.stdout}" + ) + logger.info("pip install succeeded for %s", package_spec) + + +def _reload_transformers() -> None: + """Invalidate importlib caches and force-reload transformers so the new + version is visible in the current process.""" + # Clear the metadata cache so importlib.metadata.version() returns the + # freshly installed version. + importlib.invalidate_caches() + + # Remove cached transformers modules so the next import picks up the new + # package on disk. + to_remove = [k for k in sys.modules if k == "transformers" or k.startswith("transformers.")] + for key in to_remove: + del sys.modules[key] + + +def ensure_transformers_version(model_name: str) -> None: + """Ensure the correct ``transformers`` version is installed for *model_name*. + + * If the model needs 5.x and the installed version is already 5.x → no-op. + * If the model needs 5.x but 4.x is installed → ``pip install transformers==5.1.0``. + * If the model does NOT need 5.x but 5.x is installed → downgrade to 4.57.1. + * Otherwise → no-op. + + For LoRA adapters with custom names, the base model is resolved from + ``adapter_config.json`` before checking. + + Call this at the top of every model-loading code path (training ``/start``, + inference ``/load``). + """ + # Resolve LoRA adapters to their base model for accurate detection + resolved = _resolve_base_model(model_name) + want_5 = needs_transformers_5(resolved) + current = _installed_transformers_version() + + if current is None: + logger.warning("transformers is not installed — skipping version check") + return + + current_major = int(current.split(".")[0]) + target_version = TRANSFORMERS_5_VERSION if want_5 else TRANSFORMERS_DEFAULT_VERSION + target_major = int(target_version.split(".")[0]) + + if current_major == target_major: + logger.debug( + "transformers %s already satisfies requirement (need major=%d) for model '%s'", + current, target_major, model_name, + ) + return + + logger.info( + "Model '%s' requires transformers %s but %s is installed — switching…", + model_name, target_version, current, + ) + + _pip_install(f"transformers=={target_version}") + _reload_transformers() + + new_version = _installed_transformers_version() + logger.info("Transformers version is now %s", new_version) From 1c2653fcc2cad4a188d8ce873b5b8aba330b9968 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 22 Feb 2026 18:52:06 +0000 Subject: [PATCH 02/83] aggressive reload_transformers --- studio/backend/utils/transformers_version.py | 39 ++++++++++++++++---- 1 file changed, 32 insertions(+), 7 deletions(-) diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 384a408d30..abbc97bc61 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -125,18 +125,43 @@ def _pip_install(package_spec: str) -> None: def _reload_transformers() -> None: - """Invalidate importlib caches and force-reload transformers so the new - version is visible in the current process.""" - # Clear the metadata cache so importlib.metadata.version() returns the - # freshly installed version. + """Purge transformers AND all packages that cache transformers internals + from ``sys.modules``, then force a fresh re-import. + + Simply clearing ``transformers.*`` is not enough because libraries like + ``unsloth``, ``peft``, ``trl``, and ``accelerate`` bind transformers + classes/functions into their own module-level state. We must evict them + all so the next ``import`` picks up the freshly pip-installed version. + """ importlib.invalidate_caches() - # Remove cached transformers modules so the next import picks up the new - # package on disk. - to_remove = [k for k in sys.modules if k == "transformers" or k.startswith("transformers.")] + # All top-level prefixes that hold references to transformers internals. + _PREFIXES = ( + "transformers", + "unsloth", + "unsloth_zoo", + "peft", + "trl", + "accelerate", + "auto_gptq", + "bitsandbytes", + ) + + to_remove = [ + k for k in list(sys.modules.keys()) + if any(k == p or k.startswith(p + ".") for p in _PREFIXES) + ] for key in to_remove: del sys.modules[key] + logger.info("Purged %d cached modules (%s)", len(to_remove), + ", ".join(_PREFIXES)) + + # Force a fresh import so the new version is loaded into the process. + import transformers # noqa: F811 + logger.info("Re-imported transformers — version is now %s", + transformers.__version__) + def ensure_transformers_version(model_name: str) -> None: """Ensure the correct ``transformers`` version is installed for *model_name*. From 0050e78aa35edaa01104e292393e4e38c0e12e53 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 22 Feb 2026 19:08:18 +0000 Subject: [PATCH 03/83] Fix in-memory transformers version detection and aggressive module purge for 5.1.0/4.57.1 switching --- studio/backend/utils/transformers_version.py | 121 +++++++++++++------ 1 file changed, 83 insertions(+), 38 deletions(-) diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index abbc97bc61..f1dd17e3bc 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -46,13 +46,15 @@ def _resolve_base_model(model_name: str) -> str: """If *model_name* points to a LoRA adapter, return its base model. Checks for ``adapter_config.json`` locally first (covers the common case of - local output directories with custom names). Falls back to the existing - ``get_base_model_from_lora`` utility which also handles remote HF LoRAs. + local output directories with custom names). For HF repo IDs that look + like they might be LoRA adapters, falls back to + ``get_base_model_from_lora``. Returns the original *model_name* unchanged if it is not a LoRA adapter. """ # --- Fast local check --------------------------------------------------- - adapter_cfg_path = Path(model_name) / "adapter_config.json" + local_path = Path(model_name) + adapter_cfg_path = local_path / "adapter_config.json" if adapter_cfg_path.is_file(): try: with open(adapter_cfg_path) as f: @@ -66,18 +68,24 @@ def _resolve_base_model(model_name: str) -> str: except Exception as exc: logger.debug("Could not read %s: %s", adapter_cfg_path, exc) - # --- Fallback: use the project's existing helper (handles HF repos too) -- - try: - from utils.models import get_base_model_from_lora - base = get_base_model_from_lora(model_name) - if base: - logger.info( - "Resolved LoRA adapter '%s' → base model '%s' (via get_base_model_from_lora)", - model_name, base, + # --- Only try the heavier fallback for paths that look like local dirs --- + # (Avoids triggering noisy warnings for plain HF model IDs like + # "unsloth/GLM-4.7-Flash" which are obviously not LoRA adapters.) + if local_path.is_dir(): + try: + from utils.models import get_base_model_from_lora + base = get_base_model_from_lora(model_name) + if base: + logger.info( + "Resolved LoRA adapter '%s' → base model '%s' " + "(via get_base_model_from_lora)", + model_name, base, + ) + return base + except Exception as exc: + logger.debug( + "get_base_model_from_lora failed for '%s': %s", model_name, exc, ) - return base - except Exception as exc: - logger.debug("get_base_model_from_lora failed for '%s': %s", model_name, exc) return model_name @@ -93,8 +101,17 @@ def needs_transformers_5(model_name: str) -> bool: # Version switching # --------------------------------------------------------------------------- -def _installed_transformers_version() -> str | None: - """Return the currently installed transformers version, or None.""" +def _get_in_memory_version() -> str | None: + """Return the transformers version currently loaded in this process, + or None if transformers hasn't been imported yet.""" + tf = sys.modules.get("transformers") + if tf is not None: + return getattr(tf, "__version__", None) + return None + + +def _get_on_disk_version() -> str | None: + """Return the transformers version installed on disk (pip metadata).""" try: return importlib.metadata.version("transformers") except importlib.metadata.PackageNotFoundError: @@ -166,44 +183,72 @@ def _reload_transformers() -> None: def ensure_transformers_version(model_name: str) -> None: """Ensure the correct ``transformers`` version is installed for *model_name*. - * If the model needs 5.x and the installed version is already 5.x → no-op. - * If the model needs 5.x but 4.x is installed → ``pip install transformers==5.1.0``. - * If the model does NOT need 5.x but 5.x is installed → downgrade to 4.57.1. + Checks BOTH the on-disk version (pip metadata) and the in-memory version + (``transformers.__version__``) because they can diverge after a previous + pip install in the same process. + + * If the model needs 5.x and the loaded version is already 5.x → no-op. + * If the model needs 5.x but 4.x is loaded → pip install 5.1.0 + reload. + * If the model does NOT need 5.x but 5.x is loaded → downgrade + reload. * Otherwise → no-op. For LoRA adapters with custom names, the base model is resolved from ``adapter_config.json`` before checking. Call this at the top of every model-loading code path (training ``/start``, - inference ``/load``). + inference ``/load``, export ``/load-checkpoint``). """ # Resolve LoRA adapters to their base model for accurate detection resolved = _resolve_base_model(model_name) want_5 = needs_transformers_5(resolved) - current = _installed_transformers_version() - - if current is None: - logger.warning("transformers is not installed — skipping version check") - return - - current_major = int(current.split(".")[0]) target_version = TRANSFORMERS_5_VERSION if want_5 else TRANSFORMERS_DEFAULT_VERSION target_major = int(target_version.split(".")[0]) - if current_major == target_major: - logger.debug( - "transformers %s already satisfies requirement (need major=%d) for model '%s'", - current, target_major, model_name, - ) - return + # --- Check what's actually loaded in memory first ----------------------- + in_memory = _get_in_memory_version() + on_disk = _get_on_disk_version() logger.info( - "Model '%s' requires transformers %s but %s is installed — switching…", - model_name, target_version, current, + "Version check for '%s' (resolved: '%s'): need=%s, " + "in_memory=%s, on_disk=%s", + model_name, resolved, target_version, in_memory, on_disk, ) - _pip_install(f"transformers=={target_version}") + if in_memory is not None: + in_memory_major = int(in_memory.split(".")[0]) + if in_memory_major == target_major: + logger.info( + "transformers %s in memory — already correct for '%s'", + in_memory, model_name, + ) + return + # Wrong major in memory — need to switch + logger.info( + "transformers %s loaded in memory but need %s — switching…", + in_memory, target_version, + ) + elif on_disk is not None: + on_disk_major = int(on_disk.split(".")[0]) + if on_disk_major == target_major: + logger.info( + "transformers %s on disk (not yet imported) — correct for '%s'", + on_disk, model_name, + ) + return + logger.info( + "transformers %s on disk but need %s — switching…", + on_disk, target_version, + ) + else: + logger.warning("transformers is not installed — skipping version check") + return + + # --- pip install the target version if needed --------------------------- + if on_disk is None or int(on_disk.split(".")[0]) != target_major: + _pip_install(f"transformers=={target_version}") + + # --- Purge and reload --------------------------------------------------- _reload_transformers() - new_version = _installed_transformers_version() - logger.info("Transformers version is now %s", new_version) + final = _get_in_memory_version() + logger.info("Transformers version is now %s (in memory)", final) From 15cb9b0f37becf73d7c87b5961838307c34f45ef Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 22 Feb 2026 19:19:12 +0000 Subject: [PATCH 04/83] Use sys.path overlay to switch transformers versions in-process instead of modifying site-packages --- studio/backend/utils/transformers_version.py | 236 ++++++++++--------- 1 file changed, 126 insertions(+), 110 deletions(-) diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index f1dd17e3bc..9bc9e1c94b 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -8,9 +8,13 @@ default 4.57.x that ships with Unsloth. When loading a LoRA adapter with a custom name, we resolve the base model from ``adapter_config.json`` and check *that* against the model list. -This module detects the model being loaded and ensures the correct transformers -version is installed before proceeding. The pip install is done *without* -``--force-reinstall`` to avoid rebuilding unrelated dependencies. +Strategy (sys.path overlay): + • The default transformers (4.57.x) always lives in site-packages. + • When 5.x is needed, we ``pip install --target --no-deps`` to a + separate directory and **prepend** it to ``sys.path``. + • To revert, we simply **remove** that directory from ``sys.path``. + • After either change we purge cached modules so the next import picks + up the correct version. """ import importlib @@ -18,12 +22,24 @@ import importlib.metadata import json import logging import os +import shutil import subprocess import sys +import tempfile from pathlib import Path logger = logging.getLogger(__name__) +# Ensure our logger is visible even if root logger isn't configured for INFO. +if not logger.handlers: + _handler = logging.StreamHandler() + _handler.setLevel(logging.INFO) + _handler.setFormatter( + logging.Formatter("[%(name)s|%(levelname)s]%(message)s") + ) + logger.addHandler(_handler) + logger.setLevel(logging.INFO) + # --------------------------------------------------------------------------- # Detection # --------------------------------------------------------------------------- @@ -41,14 +57,16 @@ TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = ( TRANSFORMERS_5_VERSION = "5.1.0" TRANSFORMERS_DEFAULT_VERSION = "4.57.1" +# Persistent directory for the transformers 5.x overlay +_OVERLAY_DIR = os.path.join(tempfile.gettempdir(), "transformers_5_overlay") + def _resolve_base_model(model_name: str) -> str: """If *model_name* points to a LoRA adapter, return its base model. - Checks for ``adapter_config.json`` locally first (covers the common case of - local output directories with custom names). For HF repo IDs that look - like they might be LoRA adapters, falls back to - ``get_base_model_from_lora``. + Checks for ``adapter_config.json`` locally first. Only calls the heavier + ``get_base_model_from_lora`` for paths that are actual local directories + (avoids noisy warnings for plain HF model IDs). Returns the original *model_name* unchanged if it is not a LoRA adapter. """ @@ -62,15 +80,14 @@ def _resolve_base_model(model_name: str) -> str: base = cfg.get("base_model_name_or_path") if base: logger.info( - "Resolved LoRA adapter '%s' → base model '%s'", model_name, base, + "Resolved LoRA adapter '%s' → base model '%s'", + model_name, base, ) return base except Exception as exc: logger.debug("Could not read %s: %s", adapter_cfg_path, exc) - # --- Only try the heavier fallback for paths that look like local dirs --- - # (Avoids triggering noisy warnings for plain HF model IDs like - # "unsloth/GLM-4.7-Flash" which are obviously not LoRA adapters.) + # --- Only try the heavier fallback for local directories ---------------- if local_path.is_dir(): try: from utils.models import get_base_model_from_lora @@ -84,7 +101,8 @@ def _resolve_base_model(model_name: str) -> str: return base except Exception as exc: logger.debug( - "get_base_model_from_lora failed for '%s': %s", model_name, exc, + "get_base_model_from_lora failed for '%s': %s", + model_name, exc, ) return model_name @@ -102,95 +120,111 @@ def needs_transformers_5(model_name: str) -> bool: # --------------------------------------------------------------------------- def _get_in_memory_version() -> str | None: - """Return the transformers version currently loaded in this process, - or None if transformers hasn't been imported yet.""" + """Return the transformers version currently loaded in this process.""" tf = sys.modules.get("transformers") if tf is not None: return getattr(tf, "__version__", None) return None -def _get_on_disk_version() -> str | None: - """Return the transformers version installed on disk (pip metadata).""" - try: - return importlib.metadata.version("transformers") - except importlib.metadata.PackageNotFoundError: - return None +# All top-level prefixes that hold references to transformers internals. +_PURGE_PREFIXES = ( + "transformers", + "unsloth", + "unsloth_zoo", + "peft", + "trl", + "accelerate", + "auto_gptq", + "bitsandbytes", +) -def _pip_install(package_spec: str) -> None: - """Run ``pip install `` using the running interpreter's pip. +def _purge_modules() -> int: + """Remove all cached modules for transformers and its dependents. - Mirrors the approach used in ``unsloth_zoo.llama_cpp.check_pip`` — - ``sys.executable -m pip`` is always the safest choice. - """ - cmd = [sys.executable, "-m", "pip", "install", package_spec] - logger.info("Running: %s", " ".join(cmd)) - result = subprocess.run( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - ) - if result.returncode != 0: - logger.error("pip install failed:\n%s", result.stdout) - raise RuntimeError( - f"Failed to install {package_spec}. " - f"pip output:\n{result.stdout}" - ) - logger.info("pip install succeeded for %s", package_spec) - - -def _reload_transformers() -> None: - """Purge transformers AND all packages that cache transformers internals - from ``sys.modules``, then force a fresh re-import. - - Simply clearing ``transformers.*`` is not enough because libraries like - ``unsloth``, ``peft``, ``trl``, and ``accelerate`` bind transformers - classes/functions into their own module-level state. We must evict them - all so the next ``import`` picks up the freshly pip-installed version. + Returns the number of modules purged. """ importlib.invalidate_caches() - - # All top-level prefixes that hold references to transformers internals. - _PREFIXES = ( - "transformers", - "unsloth", - "unsloth_zoo", - "peft", - "trl", - "accelerate", - "auto_gptq", - "bitsandbytes", - ) - to_remove = [ k for k in list(sys.modules.keys()) - if any(k == p or k.startswith(p + ".") for p in _PREFIXES) + if any(k == p or k.startswith(p + ".") for p in _PURGE_PREFIXES) ] for key in to_remove: del sys.modules[key] + return len(to_remove) - logger.info("Purged %d cached modules (%s)", len(to_remove), - ", ".join(_PREFIXES)) - # Force a fresh import so the new version is loaded into the process. - import transformers # noqa: F811 - logger.info("Re-imported transformers — version is now %s", +def _install_overlay() -> None: + """Install transformers 5.x into the overlay directory and prepend + it to ``sys.path`` so it shadows the default site-packages version.""" + # Install if the overlay doesn't already exist + if not os.path.isdir(_OVERLAY_DIR) or not os.listdir(_OVERLAY_DIR): + os.makedirs(_OVERLAY_DIR, exist_ok=True) + cmd = [ + sys.executable, "-m", "pip", "install", + "--target", _OVERLAY_DIR, + "--no-deps", + f"transformers=={TRANSFORMERS_5_VERSION}", + ] + logger.info("Installing transformers %s to overlay: %s", + TRANSFORMERS_5_VERSION, " ".join(cmd)) + result = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + if result.returncode != 0: + logger.error("pip install failed:\n%s", result.stdout) + raise RuntimeError( + f"Failed to install transformers=={TRANSFORMERS_5_VERSION} " + f"to {_OVERLAY_DIR}.\npip output:\n{result.stdout}" + ) + logger.info("Overlay install succeeded") + else: + logger.info("Overlay directory already exists at %s", _OVERLAY_DIR) + + # Prepend to sys.path (if not already there) + if _OVERLAY_DIR not in sys.path: + sys.path.insert(0, _OVERLAY_DIR) + logger.info("Prepended %s to sys.path", _OVERLAY_DIR) + + # Purge old modules and force fresh import + count = _purge_modules() + logger.info("Purged %d cached modules", count) + + import transformers + logger.info("Loaded transformers %s from overlay", transformers.__version__) + + +def _remove_overlay() -> None: + """Remove the overlay directory from ``sys.path`` so the default + site-packages version (4.57.x) takes effect again.""" + # Remove from sys.path + changed = False + while _OVERLAY_DIR in sys.path: + sys.path.remove(_OVERLAY_DIR) + changed = True + + if changed: + logger.info("Removed %s from sys.path", _OVERLAY_DIR) + + # Purge old modules and force fresh import + count = _purge_modules() + logger.info("Purged %d cached modules", count) + + import transformers + logger.info("Reverted to transformers %s from site-packages", transformers.__version__) def ensure_transformers_version(model_name: str) -> None: - """Ensure the correct ``transformers`` version is installed for *model_name*. + """Ensure the correct ``transformers`` version is active for *model_name*. - Checks BOTH the on-disk version (pip metadata) and the in-memory version - (``transformers.__version__``) because they can diverge after a previous - pip install in the same process. - - * If the model needs 5.x and the loaded version is already 5.x → no-op. - * If the model needs 5.x but 4.x is loaded → pip install 5.1.0 + reload. - * If the model does NOT need 5.x but 5.x is loaded → downgrade + reload. - * Otherwise → no-op. + Uses sys.path overlay: + • Need 5.x → install to separate dir, prepend sys.path, purge modules. + • Need 4.x → remove overlay from sys.path, purge modules. For LoRA adapters with custom names, the base model is resolved from ``adapter_config.json`` before checking. @@ -204,51 +238,33 @@ def ensure_transformers_version(model_name: str) -> None: target_version = TRANSFORMERS_5_VERSION if want_5 else TRANSFORMERS_DEFAULT_VERSION target_major = int(target_version.split(".")[0]) - # --- Check what's actually loaded in memory first ----------------------- + # Check what's actually loaded in memory in_memory = _get_in_memory_version() - on_disk = _get_on_disk_version() + overlay_active = _OVERLAY_DIR in sys.path logger.info( "Version check for '%s' (resolved: '%s'): need=%s, " - "in_memory=%s, on_disk=%s", - model_name, resolved, target_version, in_memory, on_disk, + "in_memory=%s, overlay_active=%s", + model_name, resolved, target_version, in_memory, overlay_active, ) + # --- Already correct? --------------------------------------------------- if in_memory is not None: in_memory_major = int(in_memory.split(".")[0]) if in_memory_major == target_major: logger.info( - "transformers %s in memory — already correct for '%s'", + "transformers %s already loaded — correct for '%s'", in_memory, model_name, ) return - # Wrong major in memory — need to switch - logger.info( - "transformers %s loaded in memory but need %s — switching…", - in_memory, target_version, - ) - elif on_disk is not None: - on_disk_major = int(on_disk.split(".")[0]) - if on_disk_major == target_major: - logger.info( - "transformers %s on disk (not yet imported) — correct for '%s'", - on_disk, model_name, - ) - return - logger.info( - "transformers %s on disk but need %s — switching…", - on_disk, target_version, - ) + + # --- Switch version ----------------------------------------------------- + if want_5: + logger.info("Activating transformers %s overlay…", TRANSFORMERS_5_VERSION) + _install_overlay() else: - logger.warning("transformers is not installed — skipping version check") - return - - # --- pip install the target version if needed --------------------------- - if on_disk is None or int(on_disk.split(".")[0]) != target_major: - _pip_install(f"transformers=={target_version}") - - # --- Purge and reload --------------------------------------------------- - _reload_transformers() + logger.info("Reverting to default transformers %s…", TRANSFORMERS_DEFAULT_VERSION) + _remove_overlay() final = _get_in_memory_version() - logger.info("Transformers version is now %s (in memory)", final) + logger.info("✓ transformers version is now %s", final) From 7cde5201765e951cde4a904e4bf2160e79cd498c Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 22 Feb 2026 19:34:46 +0000 Subject: [PATCH 05/83] Move transformers overlay to local .venv_overlay/, add huggingface-hub to overlay install --- .gitignore | 1 + studio/backend/utils/transformers_version.py | 71 +++++++++++++------- 2 files changed, 46 insertions(+), 26 deletions(-) diff --git a/.gitignore b/.gitignore index f9be87df48..e93b39259b 100755 --- a/.gitignore +++ b/.gitignore @@ -10,6 +10,7 @@ __pycache__/ # Virtual environments .venv/ +.venv_overlay/ venv/ env/ environment.yaml diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 9bc9e1c94b..76841985d0 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -25,7 +25,6 @@ import os import shutil import subprocess import sys -import tempfile from pathlib import Path logger = logging.getLogger(__name__) @@ -57,8 +56,9 @@ TRANSFORMERS_5_MODEL_SUBSTRINGS: tuple[str, ...] = ( TRANSFORMERS_5_VERSION = "5.1.0" TRANSFORMERS_DEFAULT_VERSION = "4.57.1" -# Persistent directory for the transformers 5.x overlay -_OVERLAY_DIR = os.path.join(tempfile.gettempdir(), "transformers_5_overlay") +# Persistent directory for the transformers 5.x overlay — lives next to .venv/ +_PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent.parent # studio/backend/utils/ → project root +_OVERLAY_DIR = str(_PROJECT_ROOT / ".venv_overlay") def _resolve_base_model(model_name: str) -> str: @@ -130,6 +130,7 @@ def _get_in_memory_version() -> str | None: # All top-level prefixes that hold references to transformers internals. _PURGE_PREFIXES = ( "transformers", + "huggingface_hub", "unsloth", "unsloth_zoo", "peft", @@ -155,32 +156,50 @@ def _purge_modules() -> int: return len(to_remove) +# Packages to install into the overlay (each with --no-deps). +_OVERLAY_PACKAGES = ( + f"transformers=={TRANSFORMERS_5_VERSION}", + "huggingface_hub>=1.3.0,<2.0", +) + + def _install_overlay() -> None: - """Install transformers 5.x into the overlay directory and prepend - it to ``sys.path`` so it shadows the default site-packages version.""" - # Install if the overlay doesn't already exist - if not os.path.isdir(_OVERLAY_DIR) or not os.listdir(_OVERLAY_DIR): + """Install transformers 5.x and its critical deps into the overlay + directory and prepend it to ``sys.path``.""" + # Check if overlay needs (re)creation. + # If the overlay exists but is missing huggingface_hub, it's stale. + needs_install = ( + not os.path.isdir(_OVERLAY_DIR) + or not os.listdir(_OVERLAY_DIR) + or not os.path.isdir(os.path.join(_OVERLAY_DIR, "huggingface_hub")) + ) + + if needs_install: + # Clean up stale overlay if present + if os.path.isdir(_OVERLAY_DIR): + shutil.rmtree(_OVERLAY_DIR) os.makedirs(_OVERLAY_DIR, exist_ok=True) - cmd = [ - sys.executable, "-m", "pip", "install", - "--target", _OVERLAY_DIR, - "--no-deps", - f"transformers=={TRANSFORMERS_5_VERSION}", - ] - logger.info("Installing transformers %s to overlay: %s", - TRANSFORMERS_5_VERSION, " ".join(cmd)) - result = subprocess.run( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - ) - if result.returncode != 0: - logger.error("pip install failed:\n%s", result.stdout) - raise RuntimeError( - f"Failed to install transformers=={TRANSFORMERS_5_VERSION} " - f"to {_OVERLAY_DIR}.\npip output:\n{result.stdout}" + + for pkg in _OVERLAY_PACKAGES: + cmd = [ + sys.executable, "-m", "pip", "install", + "--target", _OVERLAY_DIR, + "--no-deps", + pkg, + ] + logger.info("Installing %s to overlay: %s", pkg, " ".join(cmd)) + result = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, ) + if result.returncode != 0: + logger.error("pip install failed:\n%s", result.stdout) + raise RuntimeError( + f"Failed to install {pkg} to {_OVERLAY_DIR}.\n" + f"pip output:\n{result.stdout}" + ) logger.info("Overlay install succeeded") else: logger.info("Overlay directory already exists at %s", _OVERLAY_DIR) From 60997a75eb48cb44859c8181050d737b2b804970 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 22 Feb 2026 19:44:12 +0000 Subject: [PATCH 06/83] Install transformers into both site-packages and overlay to fix sub-package resolution during version switch --- studio/backend/utils/transformers_version.py | 90 ++++++++++++-------- 1 file changed, 54 insertions(+), 36 deletions(-) diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index 76841985d0..e244525548 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -164,22 +164,41 @@ _OVERLAY_PACKAGES = ( def _install_overlay() -> None: - """Install transformers 5.x and its critical deps into the overlay - directory and prepend it to ``sys.path``.""" - # Check if overlay needs (re)creation. - # If the overlay exists but is missing huggingface_hub, it's stale. - needs_install = ( + """Install transformers 5.x into BOTH site-packages and the overlay + directory, then prepend the overlay to ``sys.path``. + + We install into site-packages so that: + - ``importlib.metadata.version()`` returns the correct version + - sub-package resolution (``transformers.models.*``) uses 5.x code + The overlay is kept as a safety net for sys.path-based resolution. + """ + # --- 1. Install into site-packages (updates code + metadata) ----------- + for pkg in _OVERLAY_PACKAGES: + cmd = [sys.executable, "-m", "pip", "install", pkg] + logger.info("Installing %s into site-packages: %s", pkg, " ".join(cmd)) + result = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + if result.returncode != 0: + logger.error("pip install failed:\n%s", result.stdout) + raise RuntimeError( + f"Failed to install {pkg}.\npip output:\n{result.stdout}" + ) + logger.info("Site-packages install succeeded") + + # --- 2. Install into overlay (safety net for sys.path resolution) ------ + needs_overlay = ( not os.path.isdir(_OVERLAY_DIR) or not os.listdir(_OVERLAY_DIR) or not os.path.isdir(os.path.join(_OVERLAY_DIR, "huggingface_hub")) ) - - if needs_install: - # Clean up stale overlay if present + if needs_overlay: if os.path.isdir(_OVERLAY_DIR): shutil.rmtree(_OVERLAY_DIR) os.makedirs(_OVERLAY_DIR, exist_ok=True) - for pkg in _OVERLAY_PACKAGES: cmd = [ sys.executable, "-m", "pip", "install", @@ -188,54 +207,53 @@ def _install_overlay() -> None: pkg, ] logger.info("Installing %s to overlay: %s", pkg, " ".join(cmd)) - result = subprocess.run( - cmd, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - text=True, - ) - if result.returncode != 0: - logger.error("pip install failed:\n%s", result.stdout) - raise RuntimeError( - f"Failed to install {pkg} to {_OVERLAY_DIR}.\n" - f"pip output:\n{result.stdout}" - ) + subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True) logger.info("Overlay install succeeded") - else: - logger.info("Overlay directory already exists at %s", _OVERLAY_DIR) # Prepend to sys.path (if not already there) if _OVERLAY_DIR not in sys.path: sys.path.insert(0, _OVERLAY_DIR) logger.info("Prepended %s to sys.path", _OVERLAY_DIR) - # Purge old modules and force fresh import + # --- 3. Purge old modules and force fresh import ----------------------- count = _purge_modules() logger.info("Purged %d cached modules", count) import transformers - logger.info("Loaded transformers %s from overlay", transformers.__version__) + logger.info("Loaded transformers %s", transformers.__version__) def _remove_overlay() -> None: - """Remove the overlay directory from ``sys.path`` so the default - site-packages version (4.57.x) takes effect again.""" - # Remove from sys.path - changed = False + """Revert to the default transformers (4.57.x) by restoring site-packages + and removing the overlay from ``sys.path``.""" + # --- 1. Restore default versions in site-packages ---------------------- + default_packages = ( + f"transformers=={TRANSFORMERS_DEFAULT_VERSION}", + "huggingface_hub==0.36.0", + ) + for pkg in default_packages: + cmd = [sys.executable, "-m", "pip", "install", pkg] + logger.info("Restoring %s in site-packages: %s", pkg, " ".join(cmd)) + result = subprocess.run( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + ) + if result.returncode != 0: + logger.error("pip install failed:\n%s", result.stdout) + + # --- 2. Remove overlay from sys.path ----------------------------------- while _OVERLAY_DIR in sys.path: sys.path.remove(_OVERLAY_DIR) - changed = True + logger.info("Removed %s from sys.path", _OVERLAY_DIR) - if changed: - logger.info("Removed %s from sys.path", _OVERLAY_DIR) - - # Purge old modules and force fresh import + # --- 3. Purge old modules and force fresh import ----------------------- count = _purge_modules() logger.info("Purged %d cached modules", count) import transformers - logger.info("Reverted to transformers %s from site-packages", - transformers.__version__) + logger.info("Reverted to transformers %s", transformers.__version__) def ensure_transformers_version(model_name: str) -> None: From c12d75c472932da476bf68b89254013a9e61aa6f Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 22 Feb 2026 19:56:12 +0000 Subject: [PATCH 07/83] Add transformers version switch to model config and vision check endpoints for dropdown selection --- studio/backend/routes/models.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 761c04d3e7..8747aefffa 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -248,6 +248,10 @@ async def get_model_config( This endpoint wraps the backend load_model_defaults function. """ try: + # Ensure correct transformers version for this model architecture + from utils.transformers_version import ensure_transformers_version + ensure_transformers_version(model_name) + logger.info(f"Getting model config for: {model_name}") # Load model defaults from backend config_dict = load_model_defaults(model_name) @@ -371,6 +375,10 @@ async def check_vision_model( This endpoint wraps the backend is_vision_model function. """ try: + # Ensure correct transformers version for this model architecture + from utils.transformers_version import ensure_transformers_version + ensure_transformers_version(model_name) + logger.info(f"Checking if vision model: {model_name}") is_vision = is_vision_model(model_name) From 5de624614218670b0c15b4d9b75d60902a322b58 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 22 Feb 2026 20:04:35 +0000 Subject: [PATCH 08/83] Purge own utils/core modules and use lazy imports so is_vision_model picks up fresh AutoConfig after version switch --- studio/backend/routes/models.py | 9 ++++++--- studio/backend/utils/transformers_version.py | 6 ++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 8747aefffa..cd2aa8b625 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -256,8 +256,9 @@ async def get_model_config( # Load model defaults from backend config_dict = load_model_defaults(model_name) - # Check if it's a vision model - is_vision = is_vision_model(model_name) + # Lazy import to pick up fresh transformers after version switch + from utils.models import is_vision_model as _is_vision_model + is_vision = _is_vision_model(model_name) # Check if it's a LoRA adapter is_lora = False @@ -380,7 +381,9 @@ async def check_vision_model( ensure_transformers_version(model_name) logger.info(f"Checking if vision model: {model_name}") - is_vision = is_vision_model(model_name) + # Lazy import to pick up fresh transformers after version switch + from utils.models import is_vision_model as _is_vision_model + is_vision = _is_vision_model(model_name) logger.info(f"Vision check result for {model_name}: is_vision={is_vision}") return VisionCheckResponse( diff --git a/studio/backend/utils/transformers_version.py b/studio/backend/utils/transformers_version.py index e244525548..15ac70f31e 100644 --- a/studio/backend/utils/transformers_version.py +++ b/studio/backend/utils/transformers_version.py @@ -138,6 +138,12 @@ _PURGE_PREFIXES = ( "accelerate", "auto_gptq", "bitsandbytes", + # Our own modules that import from transformers at module level + # (e.g. model_config.py: `from transformers import AutoConfig`) + "utils.models", + "core.training", + "core.inference", + "core.export", ) From e3fb4f53df35811cd3c33034b4cd8b2e9558862c Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 22 Feb 2026 20:27:52 +0000 Subject: [PATCH 09/83] Patch adapter_config.json with unsloth_training_method and auto-detect load_in_4bit for LoRA inference --- studio/backend/core/training/trainer.py | 34 ++++++++++++++++++++ studio/backend/routes/inference.py | 41 ++++++++++++++++++++++++- 2 files changed, 74 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index d78c429e5a..04da0d2db9 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -68,6 +68,7 @@ class UnslothTrainer: self.is_training = False self.should_stop = False self.save_on_stop = True + self.load_in_4bit = True # Track quantization mode for metadata # Model state tracking self.is_vlm = False @@ -114,6 +115,7 @@ class UnslothTrainer: hf_token: Optional[str] = None, is_dataset_multimodal: bool = False) -> bool: """Load model for training (supports both text and vision models)""" + self.load_in_4bit = load_in_4bit # Store for training_meta.json try: if self.model is not None: del self.model @@ -988,6 +990,7 @@ class UnslothTrainer: # Stopped by user — save model at current checkpoint self.trainer.save_model() self.tokenizer.save_pretrained(output_dir) + self._patch_adapter_config(output_dir) print(f"\nTraining stopped. Model saved to {output_dir}\n") self._update_progress( is_training=False, @@ -1004,6 +1007,7 @@ class UnslothTrainer: # Normal completion self.trainer.save_model() self.tokenizer.save_pretrained(output_dir) + self._patch_adapter_config(output_dir) print(f"\nTraining completed! Model saved to {output_dir}\n") self._update_progress( is_training=False, @@ -1018,6 +1022,36 @@ class UnslothTrainer: finally: self.is_training = False + def _patch_adapter_config(self, output_dir: str) -> None: + """Patch adapter_config.json with unsloth_training_method. + + Values: 'qlora', 'lora', 'FT', 'CPT', 'DPO', 'GRPO', etc. + For LoRA/QLoRA, the distinction comes from load_in_4bit. + """ + config_path = os.path.join(output_dir, "adapter_config.json") + if not os.path.exists(config_path): + logger.info("No adapter_config.json found — skipping training method patch") + return + + try: + with open(config_path, "r") as f: + config = json.load(f) + + # Determine the training method + if self.load_in_4bit: + method = "qlora" + else: + method = "lora" + + config["unsloth_training_method"] = method + logger.info(f"Patching adapter_config.json with unsloth_training_method='{method}'") + + with open(config_path, "w") as f: + json.dump(config, f, indent=2) + + except Exception as e: + logger.warning(f"Failed to patch adapter_config.json: {e}") + def stop_training(self, save: bool = True): """Stop ongoing training""" print(f"\nStopping training (save={save})...") diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index c238b4019c..63fe246008 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -91,11 +91,50 @@ async def load_model(request: LoadRequest): detail=f"Invalid model identifier: {request.model_path}" ) + # Auto-detect quantization for LoRA adapters from adapter_config.json + # The training pipeline patches this file with "unsloth_training_method" + # which is 'qlora' or 'lora'. Only LoRA (16-bit) needs load_in_4bit=False. + load_in_4bit = request.load_in_4bit + if config.is_lora and config.path: + import json + from pathlib import Path + adapter_cfg_path = Path(config.path) / "adapter_config.json" + if adapter_cfg_path.exists(): + try: + with open(adapter_cfg_path) as f: + adapter_cfg = json.load(f) + training_method = adapter_cfg.get("unsloth_training_method") + if training_method == "lora" and load_in_4bit: + logger.info( + f"adapter_config.json says unsloth_training_method='lora' — " + f"setting load_in_4bit=False to match 16-bit training" + ) + load_in_4bit = False + elif training_method == "qlora" and not load_in_4bit: + logger.info( + f"adapter_config.json says unsloth_training_method='qlora' — " + f"setting load_in_4bit=True to match QLoRA training" + ) + load_in_4bit = True + elif training_method: + logger.info(f"Training method: {training_method}, load_in_4bit={load_in_4bit}") + else: + # No unsloth_training_method — fallback to base model name + if config.base_model and "-bnb-4bit" not in config.base_model.lower() and load_in_4bit: + logger.info( + f"No unsloth_training_method in adapter_config.json. " + f"Base model '{config.base_model}' has no -bnb-4bit suffix — " + f"setting load_in_4bit=False" + ) + load_in_4bit = False + except Exception as e: + logger.warning(f"Could not read adapter_config.json: {e}") + # Load the model success = backend.load_model( config=config, max_seq_length=request.max_seq_length, - load_in_4bit=request.load_in_4bit, + load_in_4bit=load_in_4bit, hf_token=request.hf_token, ) From 6995aaf0776b20d27bec5f2651e5c16d3d744834 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Sun, 22 Feb 2026 20:30:00 +0000 Subject: [PATCH 10/83] Clean up stale .venv_overlay directory during setup --- setup.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/setup.sh b/setup.sh index b54f782ded..4c0e167f6a 100755 --- a/setup.sh +++ b/setup.sh @@ -170,6 +170,7 @@ if [ "$IS_COLAB" = true ]; then else # Local: create venv (always start fresh to preserve correct install order) rm -rf .venv + rm -rf .venv_overlay # Clean up stale transformers version overlay "$BEST_PY" -m venv .venv source .venv/bin/activate run_quiet "pip upgrade" pip install --upgrade pip From 3fe85d36ccfa45ec8ba5f39e895ff17492a11b08 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Mon, 23 Feb 2026 05:02:45 +0000 Subject: [PATCH 11/83] Remove stale .venv_overlay on server startup to prevent transformers version conflicts --- studio/backend/main.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/studio/backend/main.py b/studio/backend/main.py index e7a33750ff..a765fc9030 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -25,6 +25,12 @@ UNSLOTH_CACHE_DIR = Path(__file__).parent / "unsloth_compiled_cache" @asynccontextmanager async def lifespan(app: FastAPI): """Startup: detect hardware, print setup token if needed. Shutdown: clean up compiled cache.""" + # Remove stale .venv_overlay from previous sessions — it will be + # rebuilt at runtime if a model needs transformers 5.x + overlay_dir = Path(__file__).resolve().parent.parent.parent / ".venv_overlay" + if overlay_dir.is_dir(): + shutil.rmtree(overlay_dir, ignore_errors=True) + # Detect hardware first — sets DEVICE global used everywhere detect_hardware() From 5416bdd4e69d14ea7997de7ffe9126850c8b5bae Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Mon, 23 Feb 2026 07:44:59 +0000 Subject: [PATCH 12/83] added shutil import to main.py --- studio/backend/main.py | 1 + 1 file changed, 1 insertion(+) diff --git a/studio/backend/main.py b/studio/backend/main.py index b29fdb3908..2d00f985e5 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -3,6 +3,7 @@ Main FastAPI application for Unsloth UI Backend """ import os import secrets +import shutil from contextlib import asynccontextmanager from fastapi import FastAPI From ac27edde35072d03aff0f08c342ccf538c6e94d0 Mon Sep 17 00:00:00 2001 From: Manan17 Date: Thu, 26 Feb 2026 08:01:18 +0000 Subject: [PATCH 13/83] merging with nightly --- .../other/OuteAI_Llama-OuteTTS-1.0-1B.yaml | 3 + .../other/Spark-TTS-0.5B_LLM.yaml | 3 + .../model_defaults/other/sesame_csm-1b.yaml | 3 + .../other/unsloth_orpheus-3b-0.1-ft.yaml | 3 + .../other/unsloth_whisper-large-v3.yaml | 3 + studio/backend/core/training/trainer.py | 1682 ++++++++++++++++- studio/backend/core/training/training.py | 9 +- studio/backend/models/training.py | 1 + studio/backend/routes/training.py | 1 + studio/backend/utils/datasets/__init__.py | 2 + .../backend/utils/datasets/data_collators.py | 27 + studio/backend/utils/models/model_config.py | 3 + .../studio/sections/training-section.tsx | 4 +- .../src/features/training/api/mappers.ts | 1 + .../src/features/training/api/models-api.ts | 1 + .../features/training/lib/model-defaults.ts | 4 + .../training/stores/training-config-store.ts | 4 + .../src/features/training/types/api.ts | 1 + .../src/features/training/types/config.ts | 1 + 19 files changed, 1736 insertions(+), 20 deletions(-) diff --git a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml index 9c65107699..72b5b018e1 100644 --- a/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml +++ b/studio/backend/assets/configs/model_defaults/other/OuteAI_Llama-OuteTTS-1.0-1B.yaml @@ -3,7 +3,10 @@ # Also applies to: OuteAI/Llama-OuteTTS-1.0-1B # added inference parameters from unsloth notebook +audio_type: dac + training: + eval_steps: 0 max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 diff --git a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml index 84cf750262..d20751b0c7 100644 --- a/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml +++ b/studio/backend/assets/configs/model_defaults/other/Spark-TTS-0.5B_LLM.yaml @@ -3,7 +3,10 @@ # Also applies to: Spark-TTS-0.5B/LLM # added inference parameters from unsloth notebook +audio_type: bicodec + training: + eval_steps: 0 max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 diff --git a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml index 294da47e10..f5f49fe1e6 100644 --- a/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml +++ b/studio/backend/assets/configs/model_defaults/other/sesame_csm-1b.yaml @@ -2,7 +2,10 @@ # Based on Sesame_CSM_(1B)-TTS.ipynb # Also applies to: sesame/csm-1b +audio_type: csm + training: + eval_steps: 0 max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml index 1bbbcdf66c..5a3c4abb48 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_orpheus-3b-0.1-ft.yaml @@ -3,7 +3,10 @@ # Also applies to: unsloth/orpheus-3b-0.1-ft-unsloth-bnb-4bit, canopylabs/orpheus-3b-0.1-ft, unsloth/orpheus-3b-0.1-ft-bnb-4bit # added inference parameters from unsloth notebook +audio_type: snac + training: + eval_steps: 0 max_seq_length: 2048 # num_epochs: 4 num_epochs: 0 diff --git a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml index d41c1c65fb..b7cb5830f4 100644 --- a/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml +++ b/studio/backend/assets/configs/model_defaults/other/unsloth_whisper-large-v3.yaml @@ -2,7 +2,10 @@ # Based on Whisper.ipynb # Also applies to: unsloth/whisper-large-v3, openai/whisper-large-v3 +audio_type: whisper + training: + eval_steps: 5 max_seq_length: 448 # num_epochs: 4 num_epochs: 0 diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index e4e7a474be..75b6aac0ef 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -71,6 +71,10 @@ class UnslothTrainer: # Model state tracking self.is_vlm = False + self.is_audio = False + self.is_audio_vlm = False # Multimodal model (e.g. Gemma 3N) trained on audio data + self._audio_type = None # 'csm', 'whisper', 'snac', 'xcodec2', 'bicodec', 'dac' + self._spark_tts_repo_dir = None # Path to downloaded Spark-TTS repo (for BiCodecTokenizer) self.model_name = None # Training metrics tracking @@ -107,12 +111,25 @@ class UnslothTrainer: except Exception as e: logger.error(f"Error in progress callback: {e}") + 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, max_seq_length: int = 2048, load_in_4bit: bool = True, hf_token: Optional[str] = None, - is_dataset_multimodal: bool = False) -> bool: + is_dataset_multimodal: bool = False, + is_dataset_audio: bool = False) -> bool: """Load model for training (supports both text and vision models)""" try: if self.model is not None: @@ -129,15 +146,23 @@ 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 if this is a vision model AND dataset is multimodal # A vision-capable model with a text-only dataset should use FastLanguageModel - self.is_vlm = is_vision_model(model_name) and is_dataset_multimodal + self.is_vlm = not self.is_audio and is_vision_model(model_name) and is_dataset_multimodal + # Audio VLM: multimodal model (e.g. Gemma 3N) trained on audio data + # Uses FastModel + SFTTrainer with audio collator (same pattern as VLM) + self.is_audio_vlm = not self.is_audio and is_vision_model(model_name) and is_dataset_audio self.model_name = model_name - logger.info(f"Model architecture is vision: {is_vision_model(model_name)}") - logger.info(f"Dataset is multimodal: {is_dataset_multimodal}") - logger.info(f"Using VLM path: {self.is_vlm}") + 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"Dataset is multimodal: {is_dataset_multimodal}, audio: {is_dataset_audio}") + logger.info(f"Using VLM path: {self.is_vlm}, Audio VLM: {self.is_audio_vlm}") # Reset training state for new run self._update_progress( @@ -151,11 +176,12 @@ class UnslothTrainer: # Update UI immediately with loading message model_display = model_name.split('/')[-1] if '/' in model_name else model_name + model_type_label = 'audio' if self.is_audio else ('vision' if self.is_vlm else 'text') self._update_progress( - status_message=f"Loading {'vision' if self.is_vlm else 'text'} model... {model_display}" + status_message=f"Loading {model_type_label} model... {model_display}" ) - print(f"\nLoading {'vision' if self.is_vlm else 'text'} model: {model_name}") + print(f"\nLoading {model_type_label} model: {model_name}") # Set HF token if provided if hf_token: @@ -163,7 +189,102 @@ class UnslothTrainer: # Branch based on model type - if self.is_vlm: + if self._audio_type == 'csm': + # CSM: FastModel + auto_model=CsmForConditionalGeneration + load_in_4bit=False + from unsloth import FastModel + from transformers import CsmForConditionalGeneration + self.model, self.tokenizer = FastModel.from_pretrained( + model_name=model_name, + max_seq_length=max_seq_length, + dtype=None, + auto_model=CsmForConditionalGeneration, + load_in_4bit=False, + token=hf_token, + ) + logger.info("Loaded CSM audio model") + + elif self._audio_type == 'whisper': + # Whisper: FastModel + auto_model=WhisperForConditionalGeneration + load_in_4bit=False + from unsloth import FastModel + from transformers import WhisperForConditionalGeneration + self.model, self.tokenizer = FastModel.from_pretrained( + model_name=model_name, + dtype=None, + load_in_4bit=False, + auto_model=WhisperForConditionalGeneration, + whisper_language="English", + whisper_task="transcribe", + token=hf_token, + ) + # Configure generation settings (notebook lines 100-105) + self.model.generation_config.language = "<|en|>" + self.model.generation_config.task = "transcribe" + self.model.config.suppress_tokens = [] + self.model.generation_config.forced_decoder_ids = None + logger.info("Loaded Whisper audio model (FastModel)") + + elif self._audio_type == 'snac': + # Orpheus: language model with audio codec tokens + self.model, self.tokenizer = FastLanguageModel.from_pretrained( + model_name=model_name, + max_seq_length=max_seq_length, + dtype=None, + load_in_4bit=load_in_4bit, + token=hf_token, + ) + logger.info(f"Loaded {self._audio_type} audio model (FastLanguageModel)") + + elif self._audio_type == 'bicodec': + # Spark-TTS: download full repo (contains sparktts package + BiCodec weights), + # then load only the LLM subfolder with FastModel. + # model_name may be: + # "Spark-TTS-0.5B/LLM" (local-style, from YAML mapping) + # "unsloth/Spark-TTS-0.5B" (HF repo ID) + from unsloth import FastModel + from huggingface_hub import snapshot_download + + if model_name.endswith("/LLM"): + # "Spark-TTS-0.5B/LLM" → parent="Spark-TTS-0.5B" + local_dir = model_name.rsplit("/", 1)[0] + hf_repo = f"unsloth/{local_dir}" + llm_path = model_name + else: + # "unsloth/Spark-TTS-0.5B" → local_dir="Spark-TTS-0.5B" + hf_repo = model_name + local_dir = model_name.split("/")[-1] + llm_path = f"{local_dir}/LLM" + + repo_path = snapshot_download(hf_repo, local_dir=local_dir) + self._spark_tts_repo_dir = os.path.abspath(repo_path) # Absolute path for sys.path + llm_path = os.path.join(self._spark_tts_repo_dir, "LLM") + + self.model, self.tokenizer = FastModel.from_pretrained( + model_name=llm_path, + max_seq_length=max_seq_length, + dtype=torch.float32, # Spark-TTS requires float32 + load_in_4bit=False, + token=hf_token, + ) + logger.info("Loaded Spark-TTS (bicodec) model") + + elif self._audio_type == 'dac': + # Phase 2: OuteTTS + raise NotImplementedError(f"Audio model type '{self._audio_type}' not yet implemented") + + elif self.is_audio_vlm: + # Audio VLM: multimodal model trained on audio (e.g. Gemma 3N) + # Uses FastModel (general loader) — returns (model, processor) + from unsloth import FastModel + self.model, self.tokenizer = FastModel.from_pretrained( + model_name=model_name, + max_seq_length=max_seq_length, + dtype=None, + load_in_4bit=load_in_4bit, + token=hf_token, + ) + logger.info("Loaded audio VLM model (FastModel)") + + elif self.is_vlm: # Load vision model - returns (model, tokenizer) self.model, self.tokenizer = FastVisionModel.from_pretrained( model_name=model_name, @@ -281,8 +402,81 @@ class UnslothTrainer: print(f"Configuring LoRA adapters (r={lora_r}, alpha={lora_alpha})...\n") print(f"Gradient checkpointing: {use_gradient_checkpointing} (type: {type(use_gradient_checkpointing).__name__})\n") - # Branch based on vision vs text - if self.is_vlm: + # Branch based on model type: audio, audio_vlm, vision, or text + if self._audio_type in ('csm', 'bicodec', 'dac') or self.is_audio_vlm: + # Models using FastModel.get_peft_model (codec audio + audio VLM) + from unsloth import FastModel + label = self._audio_type or 'audio_vlm' + print(f"{label} LoRA configuration:") + print(f" - Target modules: {target_modules}") + if self.is_audio_vlm: + print(f" - Finetune vision layers: {finetune_vision_layers}") + print(f" - Finetune language layers: {finetune_language_layers}") + print(f" - Finetune attention modules: {finetune_attention_modules}") + print(f" - Finetune MLP modules: {finetune_mlp_modules}") + print() + + peft_kwargs = dict( + r=lora_r, + target_modules=target_modules, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + bias="none", + use_gradient_checkpointing=use_gradient_checkpointing, + random_state=3407, + use_rslora=use_rslora, + loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + ) + # Audio VLM models support VLM-style layer selection + if self.is_audio_vlm: + peft_kwargs.update( + finetune_vision_layers=finetune_vision_layers, + finetune_language_layers=finetune_language_layers, + finetune_attention_modules=finetune_attention_modules, + finetune_mlp_modules=finetune_mlp_modules, + ) + + self.model = FastModel.get_peft_model(self.model, **peft_kwargs) + + elif self._audio_type == 'whisper': + # Phase 2: Whisper uses FastModel.get_peft_model with task_type=None + from unsloth import FastModel + print(f"Audio model (whisper) LoRA configuration:") + print(f" - Target modules: {target_modules}\n") + + self.model = FastModel.get_peft_model( + self.model, + r=lora_r, + target_modules=target_modules, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + bias="none", + use_gradient_checkpointing=use_gradient_checkpointing, + random_state=3407, + use_rslora=use_rslora, + loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + task_type=None, + ) + + elif self._audio_type == 'snac': + # Orpheus uses FastLanguageModel.get_peft_model + print(f"Audio model ({self._audio_type}) LoRA configuration:") + print(f" - Target modules: {target_modules}\n") + + self.model = FastLanguageModel.get_peft_model( + self.model, + r=lora_r, + target_modules=target_modules, + lora_alpha=lora_alpha, + lora_dropout=lora_dropout, + bias="none", + use_gradient_checkpointing=use_gradient_checkpointing, + random_state=3407, + use_rslora=use_rslora, + loftq_config={"loftq_bits": 4, "loftq_iter": 1} if use_loftq else None, + ) + + elif self.is_vlm: # Vision model LoRA print(f"Vision model LoRA configuration:") print(f" - Finetune vision layers: {finetune_vision_layers}") @@ -345,6 +539,715 @@ class UnslothTrainer: self._update_progress(error=error_details) return False + def _apply_csm_forward_fix(self): + """Monkey-patch CsmForConditionalGeneration.forward to fix depth decoder kwargs. + + The original transformers forward passes raw **kwargs (num_items_in_batch, + causal_mask, etc.) from the Trainer/PEFT through to the depth decoder, + causing depth_decoder_loss=None and 'Tensor + NoneType' crash. + + We patch at both instance AND class level for maximum reliability, + and strip non-TransformersKwargs params that Unsloth/PEFT inject. + """ + import types + import torch + import torch.nn as nn + from transformers.models.csm.modeling_csm import ( + CsmForConditionalGeneration, + CsmOutputWithPast, + ) + + base_csm = self.model.base_model.model # CsmForConditionalGeneration + + # Save original forward (the @can_return_tuple wrapped version) + _original_forward = CsmForConditionalGeneration.forward + + # Keys that the depth decoder and its sub-layers actually understand + _TRANSFORMERS_KWARGS = { + 'num_items_in_batch', 'output_hidden_states', 'output_attentions', + 'output_router_logits', 'cu_seq_lens_q', 'cu_seq_lens_k', + 'max_length_q', 'max_length_k', + } + + def _fixed_csm_forward( + self, + input_ids=None, input_values=None, attention_mask=None, + input_values_cutoffs=None, position_ids=None, past_key_values=None, + inputs_embeds=None, labels=None, use_cache=None, + cache_position=None, logits_to_keep=0, **kwargs, + ): + # Strip non-standard kwargs injected by Unsloth/PEFT (causal_mask, + # num_logits_to_keep, task_ids, return_dict, etc.) + output_attentions = kwargs.pop('output_attentions', None) + output_hidden_states = kwargs.pop('output_hidden_states', None) + kwargs.pop('return_dict', None) + kwargs.pop('causal_mask', None) + kwargs.pop('num_logits_to_keep', None) + kwargs.pop('task_ids', None) + + # Only keep recognized TransformersKwargs + clean_kwargs = {k: v for k, v in kwargs.items() if k in _TRANSFORMERS_KWARGS} + + if input_ids is not None and input_ids.ndim == 2: + merged = self._merge_input_ids_with_input_values( + input_ids, input_values, input_values_cutoffs, labels + ) + inputs_embeds = merged["inputs_embeds"] + labels = merged["labels"] + input_ids = None + + backbone_outputs = self.backbone_model( + input_ids=input_ids, attention_mask=attention_mask, + position_ids=position_ids, past_key_values=past_key_values, + inputs_embeds=inputs_embeds, use_cache=use_cache, + cache_position=cache_position, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + **clean_kwargs, + ) + + backbone_hidden_states = backbone_outputs[0] + slice_indices = ( + slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) + else logits_to_keep + ) + backbone_logits = self.lm_head(backbone_hidden_states[:, slice_indices, :]) + + loss = None + backbone_loss = None + depth_decoder_loss = None + depth_decoder_outputs = None + if labels is not None: + backbone_labels = labels[:, :, 0] + backbone_loss = self.loss_function( + logits=backbone_logits, labels=backbone_labels, + vocab_size=self.config.vocab_size, **clean_kwargs, + ) + + train_mask = ~(labels[:, :, 1:] == -100).all(dim=-1) + depth_decoder_input_ids = labels[train_mask][..., :self.config.num_codebooks - 1] + depth_decoder_input_ids = nn.functional.pad( + depth_decoder_input_ids, (1, 0), value=0 + ) + + train_idxs = train_mask.nonzero(as_tuple=True) + backbone_last_hidden_states = backbone_hidden_states[ + train_idxs[0], train_idxs[1] - 1, : + ] + depth_decoder_labels = labels[train_mask] + + # Build clean kwargs for depth decoder + dd_kwargs = clean_kwargs.copy() + # Scale num_items_in_batch for depth decoder (31 codebooks) + if 'num_items_in_batch' in dd_kwargs: + dd_kwargs['num_items_in_batch'] = ( + dd_kwargs['num_items_in_batch'] * (self.config.num_codebooks - 1) + ) + + depth_decoder_outputs = self.depth_decoder( + input_ids=depth_decoder_input_ids, + backbone_last_hidden_state=backbone_last_hidden_states, + use_cache=False, return_dict=True, + labels=depth_decoder_labels, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + **dd_kwargs, + ) + + depth_decoder_loss = depth_decoder_outputs.loss + if depth_decoder_loss is None: + logger.warning( + "CSM depth_decoder_loss is None! " + f"labels shape={depth_decoder_labels.shape}, " + f"train_mask sum={train_mask.sum().item()}" + ) + # Fallback: use only backbone loss instead of crashing + loss = backbone_loss + else: + loss = backbone_loss + depth_decoder_loss + + return CsmOutputWithPast( + loss=loss, backbone_loss=backbone_loss, + depth_decoder_loss=depth_decoder_loss, logits=backbone_logits, + past_key_values=backbone_outputs.past_key_values, + hidden_states=backbone_outputs.hidden_states, + attentions=backbone_outputs.attentions, + depth_decoder_logits=( + depth_decoder_outputs.logits if depth_decoder_outputs else None + ), + depth_decoder_past_key_values=( + depth_decoder_outputs.past_key_values if depth_decoder_outputs else None + ), + depth_decoder_hidden_states=( + depth_decoder_outputs.hidden_states if depth_decoder_outputs else None + ), + depth_decoder_attentions=( + depth_decoder_outputs.attentions if depth_decoder_outputs else None + ), + ) + + # Patch at BOTH instance and class level for maximum reliability. + # Instance-level: catches calls via BaseTuner.forward -> self.model.forward() + base_csm.forward = types.MethodType(_fixed_csm_forward, base_csm) + # Class-level: catches any path that resolves through the class dict + CsmForConditionalGeneration.forward = _fixed_csm_forward + print("Applied CSM forward fix (class + instance level)\n") + + def _preprocess_csm_dataset(self, dataset): + """Preprocess dataset for CSM TTS training (exact notebook copy).""" + from transformers import AutoProcessor + from datasets import Audio + import torch + + processor = AutoProcessor.from_pretrained(self.model_name) + + # Resolve speaker key + speaker_key = "source" + if "source" not in dataset.column_names and "speaker_id" not in dataset.column_names: + print("No speaker found, adding default 'source' of 0 for all examples\n") + dataset = dataset.add_column("source", ["0"] * len(dataset)) + elif "source" not in dataset.column_names and "speaker_id" in dataset.column_names: + speaker_key = "speaker_id" + + # Resolve audio and text columns + audio_col = next((c for c in dataset.column_names if c in ("audio", "Audio")), None) + text_col = next((c for c in dataset.column_names if c in ("text", "sentence", "transcript")), None) + + if audio_col is None: + raise ValueError(f"No audio column found in dataset. Columns: {dataset.column_names}") + if text_col is None: + raise ValueError(f"No text column found in dataset. Columns: {dataset.column_names}") + + print(f"CSM preprocessing: audio_col='{audio_col}', text_col='{text_col}', speaker_key='{speaker_key}'\n") + + dataset = dataset.cast_column(audio_col, Audio(sampling_rate=24000)) + + 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"]}, + ], + }] + try: + model_inputs = processor.apply_chat_template( + conversation, + tokenize=True, + return_dict=True, + output_labels=True, + text_kwargs={ + "padding": "max_length", + "max_length": 256, + "pad_to_multiple_of": 8, + "padding_side": "right", + }, + audio_kwargs={ + "sampling_rate": 24_000, + "max_length": 240001, + "padding": "max_length", + }, + common_kwargs={"return_tensors": "pt"}, + ) + except Exception as e: + logger.warning(f"Error processing CSM example: {e}") + return None + + 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 not all(isinstance(out[k], torch.Tensor) for k in out): + return None + return out + + 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 + + def _format_audio_vlm_dataset(self, dataset): + """Format dataset as audio chat messages for multimodal models (e.g. Gemma 3N). + + Expects columns: audio (Audio), text (str). + Produces: messages column with system/user/assistant chat format. + """ + from datasets import Audio + + # Detect audio and text columns + cols = dataset.column_names + audio_col = next((c for c in cols if c.lower() in ("audio", "speech")), None) + text_col = next((c for c in cols if c.lower() in ("text", "sentence", "transcript", "transcription")), None) + if not audio_col or not text_col: + raise ValueError( + f"Audio VLM dataset needs 'audio' and 'text' columns, got: {cols}" + ) + + # Cast audio to 16kHz (standard for speech models) + dataset = dataset.cast_column(audio_col, Audio(sampling_rate=16000)) + + def format_messages(samples): + formatted = {"messages": []} + for idx in range(len(samples[audio_col])): + audio = samples[audio_col][idx]["array"] + label = str(samples[text_col][idx]) + message = [ + {"role": "system", "content": [ + {"type": "text", "text": "You are an assistant that transcribes speech accurately."} + ]}, + {"role": "user", "content": [ + {"type": "audio", "audio": audio}, + {"type": "text", "text": "Please transcribe this audio."} + ]}, + {"role": "assistant", "content": [ + {"type": "text", "text": label} + ]}, + ] + formatted["messages"].append(message) + return formatted + + self._update_progress(status_message="Formatting audio VLM dataset...") + dataset = dataset.map(format_messages, batched=True, batch_size=4, num_proc=4) + print(f"Audio VLM dataset formatted: {len(dataset)} examples\n") + return dataset + + def _preprocess_snac_dataset(self, dataset): + """Preprocess dataset for Orpheus TTS training with SNAC codec. + + Mirrors Orpheus_(3B)-TTS.ipynb: encode audio with SNAC (24kHz, 3 hierarchical + layers), interleave 7 codes per frame, wrap with Orpheus special tokens, + train on full sequence (no label masking). + """ + import torch + import torchaudio.transforms as T + + SNAC_MODEL_NAME = "hubertsiuzdak/snac_24khz" + SNAC_SAMPLE_RATE = 24000 + device = "cuda" if torch.cuda.is_available() else "cpu" + max_length = getattr(self, '_max_seq_length', 2048) or 2048 + tokenizer = self.tokenizer + + # Orpheus special token IDs (hardcoded in tokenizer vocabulary) + START_OF_HUMAN = 128259 + END_OF_HUMAN = 128260 + START_OF_AI = 128261 + END_OF_AI = 128262 + START_OF_SPEECH = 128257 + END_OF_SPEECH = 128258 + END_OF_TEXT = 128009 + AUDIO_OFFSET = 128266 + + # Resolve audio and text columns (reuse CSM pattern) + cols = dataset.column_names + audio_col = next((c for c in cols if c.lower() in ("audio", "speech")), None) + text_col = next((c for c in cols if c.lower() in ("text", "sentence", "transcript", "transcription")), None) + has_source = "source" in cols + if not audio_col or not text_col: + raise ValueError( + f"SNAC dataset needs 'audio' and 'text' columns, got: {cols}" + ) + + # Get dataset sample rate from first example + first_audio = dataset[0][audio_col] + ds_sample_rate = first_audio.get("sampling_rate", SNAC_SAMPLE_RATE) if isinstance(first_audio, dict) else SNAC_SAMPLE_RATE + + # Load SNAC codec model + self._update_progress(status_message="Loading SNAC codec model...") + print("Loading SNAC codec model...\n") + from snac import SNAC + snac_model = SNAC.from_pretrained(SNAC_MODEL_NAME) + snac_model = snac_model.to(device).eval() + + # Resample transform (created once) + resample_transform = T.Resample(orig_freq=ds_sample_rate, new_freq=SNAC_SAMPLE_RATE) if ds_sample_rate != SNAC_SAMPLE_RATE else None + + self._update_progress(status_message="Encoding audio with SNAC...") + print(f"SNAC preprocessing: audio_col='{audio_col}', text_col='{text_col}', " + f"has_source={has_source}, ds_sample_rate={ds_sample_rate}\n") + + processed_examples = [] + skipped = 0 + for idx in range(len(dataset)): + if self.should_stop: + print("Stopped during SNAC preprocessing\n") + break + + example = dataset[idx] + try: + text = example.get(text_col) + if not text: + skipped += 1 + continue + + audio_data = example.get(audio_col) + if audio_data is None or audio_data.get("array") is None: + skipped += 1 + continue + + # --- Encode audio with SNAC (notebook lines 122-142) --- + waveform = torch.from_numpy(audio_data["array"]).unsqueeze(0).to(dtype=torch.float32) + if resample_transform is not None: + waveform = resample_transform(waveform) + + waveform = waveform.unsqueeze(0).to(device) + with torch.inference_mode(): + codes = snac_model.encode(waveform) + + # Interleave 7 codes per frame with layer offsets (notebook lines 134-142) + all_codes = [] + for i in range(codes[0].shape[1]): + all_codes.append(codes[0][0][i].item() + AUDIO_OFFSET) + all_codes.append(codes[1][0][2*i].item() + AUDIO_OFFSET + 4096) + all_codes.append(codes[2][0][4*i].item() + AUDIO_OFFSET + (2*4096)) + all_codes.append(codes[2][0][(4*i)+1].item() + AUDIO_OFFSET + (3*4096)) + all_codes.append(codes[1][0][(2*i)+1].item() + AUDIO_OFFSET + (4*4096)) + all_codes.append(codes[2][0][(4*i)+2].item() + AUDIO_OFFSET + (5*4096)) + all_codes.append(codes[2][0][(4*i)+3].item() + AUDIO_OFFSET + (6*4096)) + + if len(all_codes) == 0: + skipped += 1 + continue + + # Deduplicate consecutive frames with same first code (notebook lines 185-207) + deduped = all_codes[:7] + for i in range(7, len(all_codes), 7): + if all_codes[i] != deduped[-7]: + deduped.extend(all_codes[i:i+7]) + all_codes = deduped + + # --- Build text tokens (notebook lines 217-224) --- + text_prompt = f"{example['source']}: {text}" if has_source and example.get("source") else text + text_ids = tokenizer.encode(text_prompt, add_special_tokens=True) + text_ids.append(END_OF_TEXT) + + # --- Build full input_ids (notebook lines 225-234) --- + input_ids = ( + [START_OF_HUMAN] + + text_ids + + [END_OF_HUMAN] + + [START_OF_AI] + + [START_OF_SPEECH] + + all_codes + + [END_OF_SPEECH] + + [END_OF_AI] + ) + + # Truncate to max_length + input_ids = input_ids[:max_length] + + # Labels = input_ids (no masking — Orpheus trains on full sequence) + labels = list(input_ids) + attention_mask = [1] * len(input_ids) + + processed_examples.append({ + "input_ids": input_ids, + "labels": labels, + "attention_mask": attention_mask, + }) + + except Exception as e: + logger.warning(f"Error processing SNAC example {idx}: {e}") + skipped += 1 + continue + + # Progress update every 100 examples + if (idx + 1) % 100 == 0: + self._update_progress( + status_message=f"Encoding audio... {idx + 1}/{len(dataset)}" + ) + + # Free SNAC model from GPU + print("Freeing SNAC codec model from GPU...\n") + snac_model.to("cpu") + del snac_model + torch.cuda.empty_cache() + + if not processed_examples: + raise ValueError( + f"No valid examples after SNAC preprocessing (skipped {skipped})" + ) + + result_dataset = Dataset.from_list(processed_examples) + print(f"SNAC preprocessing complete: {len(result_dataset)} examples " + f"({skipped} skipped)\n") + return result_dataset + + def _preprocess_bicodec_dataset(self, dataset): + """Preprocess dataset for Spark-TTS training with BiCodec tokenizer. + + Mirrors Spark_TTS_(0_5B).ipynb: encode audio with BiCodec (semantic + global tokens), + format as special-token text strings for SFTTrainer with dataset_text_field="text". + """ + import sys + import torch + import numpy as np + import torchaudio.transforms as T + + import subprocess + + device = "cuda" if torch.cuda.is_available() else "cpu" + + # The sparktts Python package lives in the SparkAudio/Spark-TTS GitHub repo, + # NOT in the unsloth/Spark-TTS-0.5B HF model repo. Clone it if needed. + spark_code_dir = os.path.join(os.path.dirname(self._spark_tts_repo_dir), "Spark-TTS") + sparktts_pkg = os.path.join(spark_code_dir, "sparktts") + if not os.path.isdir(sparktts_pkg): + self._update_progress(status_message="Cloning Spark-TTS code repo...") + print(f"Cloning SparkAudio/Spark-TTS to {spark_code_dir}...\n") + subprocess.run( + ["git", "clone", "--depth", "1", "https://github.com/SparkAudio/Spark-TTS", spark_code_dir], + check=True, + ) + + if spark_code_dir not in sys.path: + sys.path.insert(0, spark_code_dir) + + from sparktts.models.audio_tokenizer import BiCodecTokenizer + from sparktts.utils.audio import audio_volume_normalize + + # Resolve audio and text columns + cols = dataset.column_names + audio_col = next((c for c in cols if c.lower() in ("audio", "speech")), None) + text_col = next((c for c in cols if c.lower() in ("text", "sentence", "transcript", "transcription")), None) + has_source = "source" in cols + if not audio_col or not text_col: + raise ValueError( + f"BiCodec dataset needs 'audio' and 'text' columns, got: {cols}" + ) + + # Load BiCodec tokenizer + self._update_progress(status_message="Loading BiCodec tokenizer...") + print("Loading BiCodec tokenizer...\n") + audio_tokenizer = BiCodecTokenizer(self._spark_tts_repo_dir, device) + + target_sr = audio_tokenizer.config['sample_rate'] + + self._update_progress(status_message="Encoding audio with BiCodec...") + print(f"BiCodec preprocessing: audio_col='{audio_col}', text_col='{text_col}', " + f"has_source={has_source}, target_sr={target_sr}\n") + + def extract_wav2vec2_features(wavs: torch.Tensor) -> torch.Tensor: + """Extract wav2vec2 features (average of layers 11, 14, 16).""" + if wavs.shape[0] != 1: + raise ValueError(f"Expected batch size 1, but got shape {wavs.shape}") + wav_np = wavs.squeeze(0).cpu().numpy() + + processed = audio_tokenizer.processor( + wav_np, + sampling_rate=16000, + return_tensors="pt", + padding=True, + ) + input_values = processed.input_values.to(audio_tokenizer.feature_extractor.device) + model_output = audio_tokenizer.feature_extractor(input_values) + + if model_output.hidden_states is None: + raise ValueError("Wav2Vec2Model did not return hidden states.") + + feats_mix = ( + model_output.hidden_states[11] + + model_output.hidden_states[14] + + model_output.hidden_states[16] + ) / 3 + return feats_mix + + processed_examples = [] + skipped = 0 + for idx in range(len(dataset)): + if self.should_stop: + print("Stopped during BiCodec preprocessing\n") + break + + example = dataset[idx] + try: + text = example.get(text_col) + if not text: + skipped += 1 + continue + + audio_data = example.get(audio_col) + if audio_data is None or audio_data.get("array") is None: + skipped += 1 + continue + + audio_array = audio_data["array"] + sampling_rate = audio_data.get("sampling_rate", target_sr) + + # Resample if needed + if sampling_rate != target_sr: + resampler = T.Resample(orig_freq=sampling_rate, new_freq=target_sr) + audio_tensor_temp = torch.from_numpy(audio_array).float() + audio_array = resampler(audio_tensor_temp).numpy() + + # Volume normalize if configured + if audio_tokenizer.config.get("volume_normalize", False): + audio_array = audio_volume_normalize(audio_array) + + # Get reference clip + ref_wav_np = audio_tokenizer.get_ref_clip(audio_array) + + # Prepare tensors + audio_tensor = torch.from_numpy(audio_array).unsqueeze(0).float().to(device) + ref_wav_tensor = torch.from_numpy(ref_wav_np).unsqueeze(0).float().to(device) + + # Extract wav2vec2 features + feat = extract_wav2vec2_features(audio_tensor) + + batch = { + "wav": audio_tensor, + "ref_wav": ref_wav_tensor, + "feat": feat.to(device), + } + + # BiCodec tokenize + semantic_token_ids, global_token_ids = audio_tokenizer.model.tokenize(batch) + + global_tokens = "".join( + [f"<|bicodec_global_{i}|>" for i in global_token_ids.squeeze().cpu().numpy()] + ) + semantic_tokens = "".join( + [f"<|bicodec_semantic_{i}|>" for i in semantic_token_ids.squeeze().cpu().numpy()] + ) + + # Format text with source prefix if available + text_content = f"{example['source']}: {text}" if has_source and example.get("source") else text + + formatted = "".join([ + "<|task_tts|>", + "<|start_content|>", + text_content, + "<|end_content|>", + "<|start_global_token|>", + global_tokens, + "<|end_global_token|>", + "<|start_semantic_token|>", + semantic_tokens, + "<|end_semantic_token|>", + "<|im_end|>", + ]) + + processed_examples.append({"text": formatted}) + + except Exception as e: + logger.warning(f"Error processing BiCodec example {idx}: {e}") + skipped += 1 + continue + + # Progress update every 100 examples + if (idx + 1) % 100 == 0: + self._update_progress( + status_message=f"Encoding audio with BiCodec... {idx + 1}/{len(dataset)}" + ) + + # Free BiCodec model from GPU + print("Freeing BiCodec tokenizer from GPU...\n") + audio_tokenizer.model.cpu() + audio_tokenizer.feature_extractor.cpu() + torch.cuda.empty_cache() + + if not processed_examples: + raise ValueError( + f"No valid examples after BiCodec preprocessing (skipped {skipped})" + ) + + result_dataset = Dataset.from_list(processed_examples) + print(f"BiCodec preprocessing complete: {len(result_dataset)} examples " + f"({skipped} skipped)\n") + # Debug: show first example text (truncated) + sample = result_dataset[0]["text"] + print(f"Sample text (first 200 chars): {sample[:200]}...\n") + print(f"Sample text length: {len(sample)} chars\n") + return result_dataset + + def _preprocess_whisper_dataset(self, dataset, eval_split=None): + """Preprocess dataset for Whisper speech-to-text training. + + Mirrors Whisper.ipynb: extract audio features with Whisper's feature + extractor, tokenize text labels. Returns (train_data, eval_data) where + each is a list of dicts with 'input_features' and 'labels'. + """ + from datasets import Audio + + WHISPER_SAMPLE_RATE = 16000 + + # Resolve audio and text columns + cols = dataset.column_names + audio_col = next((c for c in cols if c.lower() in ("audio", "speech")), None) + text_col = next((c for c in cols if c.lower() in ("text", "sentence", "transcript", "transcription")), None) + if not audio_col or not text_col: + raise ValueError( + f"Whisper dataset needs 'audio' and 'text' columns, got: {cols}" + ) + + # Cast audio to 16kHz (Whisper's expected sample rate) + dataset = dataset.cast_column(audio_col, Audio(sampling_rate=WHISPER_SAMPLE_RATE)) + + # Train/eval split (notebook does dataset.train_test_split) + eval_dataset_raw = None + if eval_split: + splits = dataset.train_test_split(test_size=0.06, seed=42) + dataset = splits["train"] + eval_dataset_raw = splits["test"] + + self._update_progress(status_message="Processing audio for Whisper...") + print(f"Whisper preprocessing: audio_col='{audio_col}', text_col='{text_col}', " + f"samples={len(dataset)}\n") + + def process_split(ds, split_name="train"): + processed = [] + skipped = 0 + for idx in range(len(ds)): + if self.should_stop: + print(f"Stopped during Whisper {split_name} preprocessing\n") + break + + example = ds[idx] + try: + audio_data = example.get(audio_col) + text = example.get(text_col) + if audio_data is None or audio_data.get("array") is None or not text: + skipped += 1 + continue + + # Extract audio features (notebook line 112-115) + features = self.tokenizer.feature_extractor( + audio_data["array"], sampling_rate=audio_data["sampling_rate"] + ) + # Tokenize text (notebook line 116) + tokenized_text = self.tokenizer.tokenizer(text) + + processed.append({ + "input_features": features.input_features[0], + "labels": tokenized_text.input_ids, + }) + except Exception as e: + logger.warning(f"Error processing Whisper {split_name} example {idx}: {e}") + skipped += 1 + continue + + if (idx + 1) % 100 == 0: + self._update_progress( + status_message=f"Processing {split_name} audio... {idx + 1}/{len(ds)}" + ) + + print(f"Whisper {split_name} preprocessing: {len(processed)} examples ({skipped} skipped)\n") + return processed + + train_data = process_split(dataset, "train") + eval_data = process_split(eval_dataset_raw, "eval") if eval_dataset_raw else None + + if not train_data: + raise ValueError("No valid examples after Whisper preprocessing") + + return (train_data, eval_data) + def load_and_format_dataset(self, dataset_source: str, format_type: str = "auto", @@ -450,6 +1353,33 @@ class UnslothTrainer: print("Stopped before applying chat template\n") return None + # ========== AUDIO MODELS: custom preprocessing ========== + if self._audio_type == 'csm': + processed = self._preprocess_csm_dataset(dataset) + # CSM returns a ready-to-train Dataset (not a dict) with no eval + return (processed, None) + + elif self._audio_type == 'whisper': + train_data, eval_data = self._preprocess_whisper_dataset(dataset, eval_split=eval_split) + return (train_data, eval_data) + + elif self._audio_type == 'snac': + processed = self._preprocess_snac_dataset(dataset) + return (processed, None) + + elif self._audio_type == 'bicodec': + processed = self._preprocess_bicodec_dataset(dataset) + return (processed, None) + + elif self._audio_type in ('xcodec2', 'dac'): + # Phase 2: remaining codec-to-text models + raise NotImplementedError(f"Audio dataset preprocessing for '{self._audio_type}' not yet implemented") + + elif self.is_audio_vlm: + # Audio VLM (e.g. Gemma 3N): format as chat messages with audio content + formatted = self._format_audio_vlm_dataset(dataset) + return (formatted, None) + # ========== FORMAT FIRST ========== print(f"Formatting dataset with format_type='{format_type}'...\n") @@ -650,6 +1580,675 @@ class UnslothTrainer: output_dir = training_args.get('output_dir', './outputs') os.makedirs(output_dir, exist_ok=True) + # ========== AUDIO TRAINER BRANCH ========== + if self._audio_type == 'csm': + # CSM uses plain HF Trainer 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. + 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 + "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), + ) + print("CSM Trainer initialized\n") + + # 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 + ) + + 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.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 + + 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 + + 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), + ) + 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.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 + + elif self._audio_type == 'whisper': + # Whisper: Seq2SeqTrainer with custom speech collator + from transformers import Seq2SeqTrainer, Seq2SeqTrainingArguments, TrainerCallback + 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 + if eval_dataset: + whisper_training_args["eval_strategy"] = "steps" + whisper_training_args["eval_steps"] = eval_steps_val + + # 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") + + trainer_kwargs = { + "model": self.model, + "train_dataset": train_ds, + "data_collator": data_collator, + "tokenizer": self.tokenizer.feature_extractor, + "args": Seq2SeqTrainingArguments(**whisper_training_args), + } + 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 + + 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), + ) + 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.trainer.train() + + # 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 + raise NotImplementedError(f"Audio training for '{self._audio_type}' not yet implemented") + # ========== DATA COLLATOR SELECTION ========== # Detect special model types model_name_lower = self.model_name.lower() @@ -703,6 +2302,39 @@ class UnslothTrainer: data_collator = UnslothVisionDataCollator(self.model, self.tokenizer) print("Vision data collator configured\n") + elif self.is_audio_vlm: + # Audio VLM collator (e.g. Gemma 3N with audio data) + # Mirrors the collate_fn from Gemma3N_(4B)-Audio notebook + print("Configuring audio VLM data collator...\n") + processor = self.tokenizer # FastModel returns processor as tokenizer + + def audio_vlm_collate_fn(examples): + texts = [] + audios = [] + for example in examples: + text = processor.apply_chat_template( + example["messages"], tokenize=False, add_generation_prompt=False + ).strip() + texts.append(text) + audios.append(example["audio"]["array"]) + + batch = processor( + text=texts, audio=audios, return_tensors="pt", padding=True + ) + + # Labels = input_ids with special tokens masked + labels = batch["input_ids"].clone() + labels[labels == processor.tokenizer.pad_token_id] = -100 + for attr in ('audio_token_id', 'image_token_id', 'boi_token_id', 'eoi_token_id'): + token_id = getattr(processor.tokenizer, attr, None) + if token_id is not None: + labels[labels == token_id] = -100 + batch["labels"] = labels + return batch + + data_collator = audio_vlm_collate_fn + print("Audio VLM data collator configured\n") + # ========== TRAINING CONFIGURATION ========== # Handle epochs vs max_steps properly max_steps_val = training_args.get('max_steps', 0) @@ -775,9 +2407,10 @@ class UnslothTrainer: optim_value = training_args.get('optim', "adamw_8bit") lr_scheduler_type_value = training_args.get('lr_scheduler_type', "linear") - if self.is_vlm: - # Vision-specific config - print("Configuring vision model training parameters\n") + if self.is_vlm or self.is_audio_vlm: + # Vision / audio VLM config (both need skip_prepare_dataset + remove_unused_columns) + label = "audio VLM" if self.is_audio_vlm else "vision" + print(f"Configuring {label} model training parameters\n") # Use provided values or defaults for vision models optim_value = training_args.get('optim', "adamw_torch_fused") lr_scheduler_type_value = training_args.get('lr_scheduler_type', "cosine") @@ -786,7 +2419,7 @@ class UnslothTrainer: "lr_scheduler_type": lr_scheduler_type_value, "gradient_checkpointing": True, "gradient_checkpointing_kwargs": {"use_reentrant": False}, - "max_grad_norm": 0.3, # Recommended for vision models + "max_grad_norm": 0.3, "remove_unused_columns": False, "dataset_text_field": "", "dataset_kwargs": {"skip_prepare_dataset": True}, @@ -810,11 +2443,19 @@ class UnslothTrainer: print("Training configuration prepared\n") # ========== TRAINER INITIALIZATION ========== - if self.is_vlm: + if self.is_vlm or self.is_audio_vlm: + # VLM: dataset is dict wrapper from format_and_template_dataset + # Audio VLM: dataset is raw Dataset from _format_audio_vlm_dataset + train_dataset = dataset['dataset'] if isinstance(dataset, dict) else dataset trainer_kwargs = { "model": self.model, +<<<<<<< HEAD "train_dataset": dataset['dataset'], "processing_class": self.tokenizer, +======= + "train_dataset": train_dataset, + "processing_class": self.tokenizer.tokenizer, +>>>>>>> 0a7e75e (Adding support for audio llms) "data_collator": data_collator, "args": SFTConfig(**config_args), } @@ -852,7 +2493,8 @@ class UnslothTrainer: train_on_responses_enabled = training_args.get('train_on_completions', False) # DeepSeek OCR handles this internally in its collator, so skip - if train_on_responses_enabled and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'): + # Audio VLM handles label masking in its collator, so skip + if train_on_responses_enabled and not self.is_audio_vlm and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'): try: print("Configuring train on responses only...\n") @@ -881,7 +2523,7 @@ class UnslothTrainer: train_on_responses_enabled = False # Apply train on responses only if we have valid parts - if train_on_responses_enabled and instruction_part and response_part and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'): + if train_on_responses_enabled and instruction_part and response_part and not self.is_audio_vlm and not (is_deepseek_ocr or dataset["final_format"].lower() == 'alpaca'): try: from unsloth.chat_templates import train_on_responses_only @@ -1008,7 +2650,11 @@ class UnslothTrainer: progress_callback = ProgressCallback(self) self.trainer.add_callback(progress_callback) +<<<<<<< HEAD num_samples = len(self.trainer.train_dataset) +======= + num_samples = len(dataset['dataset'] if isinstance(dataset, dict) else dataset) +>>>>>>> 0a7e75e (Adding support for audio llms) 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) @@ -1069,7 +2715,9 @@ class UnslothTrainer: ) except Exception as e: + import traceback logger.error(f"Training error: {e}") + logger.error(f"Full traceback:\n{traceback.format_exc()}") self._update_progress(is_training=False, error=str(e)) finally: diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 9123d36b39..a2020cf68c 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -116,7 +116,8 @@ class TrainingBackend: train_split: str = "train", eval_split: str = None, eval_steps: float = 0.00, - is_dataset_multimodal: bool = False) -> bool: + is_dataset_multimodal: bool = False, + is_dataset_audio: bool = False) -> bool: """ Start training. @@ -145,6 +146,11 @@ class TrainingBackend: import torch as _torch if _torch.cuda.is_available(): _torch.cuda.synchronize() + # Reset torch dynamo/compiler caches — Unsloth's compiled SFTTrainer + # and model.for_training() set class-level state that persists between + # runs (e.g. BiCodec text trainer pollutes subsequent VLM runs). + _torch._dynamo.reset() + _torch.compiler.reset() import gc gc.collect() clear_gpu_cache() @@ -177,6 +183,7 @@ class TrainingBackend: load_in_4bit=load_in_4bit if use_lora_actual else False, # Only 4bit for LoRA hf_token=hf_token if hf_token.strip() else None, is_dataset_multimodal=is_dataset_multimodal, + is_dataset_audio=is_dataset_audio, ) if not success or self.trainer.should_stop: diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 54de974100..72d8b256ac 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -66,6 +66,7 @@ class TrainingStartRequest(BaseModel): finetune_attention_modules: bool = Field(False, description="Finetune attention modules") finetune_mlp_modules: bool = Field(False, description="Finetune MLP modules") is_dataset_multimodal: bool = Field(False, description="Whether the dataset contains multimodal (image) data") + is_dataset_audio: bool = Field(False, description="Whether the dataset contains audio data") # Logging parameters enable_wandb: bool = Field(False, description="Enable Weights & Biases logging") diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index f8de2f639f..1fb49f493f 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -179,6 +179,7 @@ async def start_training( "finetune_attention_modules": request.finetune_attention_modules, "finetune_mlp_modules": request.finetune_mlp_modules, "is_dataset_multimodal": request.is_dataset_multimodal, + "is_dataset_audio": request.is_dataset_audio, "enable_wandb": request.enable_wandb, "wandb_token": request.wandb_token or "", "wandb_project": request.wandb_project or "", diff --git a/studio/backend/utils/datasets/__init__.py b/studio/backend/utils/datasets/__init__.py index b47db737c3..2e78057237 100644 --- a/studio/backend/utils/datasets/__init__.py +++ b/studio/backend/utils/datasets/__init__.py @@ -45,6 +45,7 @@ from .vlm_processing import ( # Data collators from .data_collators import ( + DataCollatorSpeechSeq2SeqWithPadding, DeepSeekOCRDataCollator, VLMDataCollator, ) @@ -85,6 +86,7 @@ __all__ = [ # VLM "generate_smart_vlm_instruction", # Collators + "DataCollatorSpeechSeq2SeqWithPadding", "DeepSeekOCRDataCollator", "VLMDataCollator", # Mappings diff --git a/studio/backend/utils/datasets/data_collators.py b/studio/backend/utils/datasets/data_collators.py index f453eaea1b..41062f6a6f 100644 --- a/studio/backend/utils/datasets/data_collators.py +++ b/studio/backend/utils/datasets/data_collators.py @@ -10,6 +10,33 @@ from dataclasses import dataclass from typing import Any, List, Optional, Union +@dataclass +class DataCollatorSpeechSeq2SeqWithPadding: + """ + Data collator for Whisper speech-to-text training. + + Pads input features (audio) and label sequences (text) separately, + masks padding in labels with -100, and strips leading BOS token. + Mirrors the collator from the Whisper.ipynb notebook. + """ + processor: Any + + def __call__(self, features: List[dict]) -> dict: + input_features = [{"input_features": feature["input_features"]} for feature in features] + batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt") + + label_features = [{"input_ids": feature["labels"]} for feature in features] + labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt") + + labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100) + + if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item(): + labels = labels[:, 1:] + + batch["labels"] = labels + return batch + + @dataclass class DeepSeekOCRDataCollator: """ diff --git a/studio/backend/utils/models/model_config.py b/studio/backend/utils/models/model_config.py index b84a14d226..3dde2b3a63 100644 --- a/studio/backend/utils/models/model_config.py +++ b/studio/backend/utils/models/model_config.py @@ -213,6 +213,7 @@ MODEL_NAME_MAPPING = { "unsloth/Nemotron-3-Nano-30B-A3B", ], "unsloth_orpheus-3b-0.1-ft.yaml": [ + "unsloth/orpheus-3b-0.1-ft", "unsloth/orpheus-3b-0.1-ft-unsloth-bnb-4bit", "canopylabs/orpheus-3b-0.1-ft", "unsloth/orpheus-3b-0.1-ft-bnb-4bit", @@ -317,9 +318,11 @@ MODEL_NAME_MAPPING = { ], "sesame_csm-1b.yaml": [ "sesame/csm-1b", + "unsloth/csm-1b", ], "Spark-TTS-0.5B_LLM.yaml": [ "Spark-TTS-0.5B/LLM", + "unsloth/Spark-TTS-0.5B", ], "unsloth_tinyllama-bnb-4bit.yaml": [ "unsloth/tinyllama", diff --git a/studio/frontend/src/features/studio/sections/training-section.tsx b/studio/frontend/src/features/studio/sections/training-section.tsx index af6a16fc44..160a0808c9 100644 --- a/studio/frontend/src/features/studio/sections/training-section.tsx +++ b/studio/frontend/src/features/studio/sections/training-section.tsx @@ -42,8 +42,8 @@ export function TrainingSection() { const store = useTrainingConfigStore(); const { isStarting, startError, startTrainingRun } = useTrainingActions(); const isIncompatible = - !store.isVisionModel && store.isDatasetMultimodal === true; - const fileInputRef = useRef(null); + !store.isVisionModel && !store.isDatasetAudio && store.isDatasetMultimodal === true; + const fileInputRef = useRef(null); const handleFileUpload = (e: React.ChangeEvent) => { const file = e.target.files?.[0]; diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index 1adfbd8b6d..923306861d 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -58,6 +58,7 @@ export function buildTrainingStartPayload( finetune_attention_modules: config.finetuneAttentionModules, finetune_mlp_modules: config.finetuneMLPModules, is_dataset_multimodal: !!config.isDatasetMultimodal, + is_dataset_audio: config.isDatasetAudio, enable_wandb: config.enableWandb, wandb_token: config.enableWandb ? config.wandbToken.trim() || null : null, wandb_project: config.enableWandb diff --git a/studio/frontend/src/features/training/api/models-api.ts b/studio/frontend/src/features/training/api/models-api.ts index e22de5f581..0536eddf4c 100644 --- a/studio/frontend/src/features/training/api/models-api.ts +++ b/studio/frontend/src/features/training/api/models-api.ts @@ -46,6 +46,7 @@ interface BackendLoggingDefaults { } export interface BackendModelConfig { + audio_type?: string | null; training?: BackendTrainingDefaults; lora?: BackendLoraDefaults; logging?: BackendLoggingDefaults; diff --git a/studio/frontend/src/features/training/lib/model-defaults.ts b/studio/frontend/src/features/training/lib/model-defaults.ts index 35ce562dbf..e68516e7de 100644 --- a/studio/frontend/src/features/training/lib/model-defaults.ts +++ b/studio/frontend/src/features/training/lib/model-defaults.ts @@ -4,6 +4,7 @@ import type { TrainingConfigState } from "../types/config"; type ModelDefaultsPatch = Partial< Pick< TrainingConfigState, + | "isDatasetAudio" | "epochs" | "contextLength" | "learningRate" @@ -79,6 +80,9 @@ export function mapBackendModelConfigToTrainingPatch( const lora = config.lora; const logging = config.logging; + // Audio models: set isDatasetAudio based on audio_type from YAML + patch.isDatasetAudio = typeof config.audio_type === "string" && config.audio_type.length > 0; + const maxSeqLength = toNumber(training?.max_seq_length); if (maxSeqLength !== undefined) patch.contextLength = maxSeqLength; diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index b2d1858716..e72acef5d2 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -36,6 +36,7 @@ const initialState: TrainingConfigState = { modelDefaultsAppliedFor: null, isCheckingDataset: false, isDatasetMultimodal: null, + isDatasetAudio: false, ...DEFAULT_HYPERPARAMS, }; @@ -57,6 +58,7 @@ const NON_PERSISTED_STATE_KEYS: ReadonlySet = new Set "modelDefaultsAppliedFor", "isCheckingDataset", "isDatasetMultimodal", + "isDatasetAudio", "trainOnCompletions", ]); @@ -205,6 +207,7 @@ export const useTrainingConfigStore = create()( selectedModel: null, isCheckingVision: false, isVisionModel: false, + isDatasetAudio: false, isLoadingModelDefaults: false, modelDefaultsError: null, modelDefaultsAppliedFor: null, @@ -220,6 +223,7 @@ export const useTrainingConfigStore = create()( set({ isCheckingVision: false, isVisionModel: false, + isDatasetAudio: false, isLoadingModelDefaults: false, modelDefaultsError: null, modelDefaultsAppliedFor: null, diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index 22f02e4331..9f4c849200 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/studio/frontend/src/features/training/types/api.ts @@ -39,6 +39,7 @@ export interface TrainingStartRequest { finetune_attention_modules: boolean; finetune_mlp_modules: boolean; is_dataset_multimodal: boolean; + is_dataset_audio: boolean; enable_wandb: boolean; wandb_token: string | null; wandb_project: string | null; diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index 6c2feec172..0cf2129f3a 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -60,6 +60,7 @@ export interface TrainingConfigState { modelDefaultsAppliedFor: string | null; isCheckingDataset: boolean; isDatasetMultimodal: boolean | null; + isDatasetAudio: boolean; finetuneVisionLayers: boolean; finetuneLanguageLayers: boolean; finetuneAttentionModules: boolean; From ab2ac39017f6b6cf187a553b7627a8cf49321f6e Mon Sep 17 00:00:00 2001 From: Manan17 Date: Thu, 26 Feb 2026 22:17:25 +0000 Subject: [PATCH 14/83] Changes with audio training --- studio/backend/core/training/trainer.py | 207 +++++++++++------- studio/backend/models/datasets.py | 3 + studio/backend/routes/datasets.py | 3 + .../backend/utils/datasets/dataset_utils.py | 112 ++++++---- .../utils/datasets/format_detection.py | 121 ++++++++-- .../dataset-preview-dialog-mapping.tsx | 34 ++- .../sections/dataset-preview-dialog.tsx | 10 +- .../training/hooks/use-training-actions.ts | 30 ++- .../features/training/lib/model-defaults.ts | 4 - .../training/stores/training-config-store.ts | 2 + .../src/features/training/types/datasets.ts | 3 + 11 files changed, 371 insertions(+), 158 deletions(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 75b6aac0ef..d77effc0d4 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -111,6 +111,41 @@ class UnslothTrainer: except Exception as e: logger.error(f"Error in progress callback: {e}") + def _resolve_audio_columns(self, dataset, custom_format_mapping: dict = None): + """Resolve audio, text, and speaker columns from user mapping or hardcoded fallback. + + Returns: + dict with keys: audio_col, text_col, speaker_col (speaker_col may be None) + """ + cols = dataset.column_names + + if custom_format_mapping: + audio_col = None + text_col = None + speaker_col = None + for col, role in custom_format_mapping.items(): + if role == "audio": + audio_col = col + elif role == "text": + text_col = col + elif role == "speaker_id": + speaker_col = col + # Use mapping if both required columns exist in the dataset + if audio_col and audio_col in cols and text_col and text_col in cols: + return {"audio_col": audio_col, "text_col": text_col, "speaker_col": speaker_col} + + # Hardcoded fallback (existing behavior) + audio_col = next((c for c in cols if c.lower() in ("audio", "speech")), None) + text_col = next((c for c in cols if c.lower() in ("text", "sentence", "transcript", "transcription")), None) + + speaker_col = None + if "source" in cols: + speaker_col = "source" + elif "speaker_id" in cols: + speaker_col = "speaker_id" + + return {"audio_col": audio_col, "text_col": text_col, "speaker_col": speaker_col} + def _resolve_audio_type(self, model_name: str) -> Optional[str]: """Resolve audio_type from YAML model config. Returns None for non-audio models.""" try: @@ -150,12 +185,11 @@ class UnslothTrainer: self._audio_type = self._resolve_audio_type(model_name) self.is_audio = self._audio_type is not None - # Detect if this is a vision model AND dataset is multimodal - # A vision-capable model with a text-only dataset should use FastLanguageModel - self.is_vlm = not self.is_audio and is_vision_model(model_name) and is_dataset_multimodal # Audio VLM: multimodal model (e.g. Gemma 3N) trained on audio data - # Uses FastModel + SFTTrainer with audio collator (same pattern as VLM) + # 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_multimodal self.model_name = model_name logger.info(f"Audio type: {self._audio_type}") @@ -693,7 +727,7 @@ class UnslothTrainer: CsmForConditionalGeneration.forward = _fixed_csm_forward print("Applied CSM forward fix (class + instance level)\n") - def _preprocess_csm_dataset(self, dataset): + def _preprocess_csm_dataset(self, dataset, custom_format_mapping=None): """Preprocess dataset for CSM TTS training (exact notebook copy).""" from transformers import AutoProcessor from datasets import Audio @@ -701,22 +735,20 @@ class UnslothTrainer: processor = AutoProcessor.from_pretrained(self.model_name) - # Resolve speaker key - speaker_key = "source" - if "source" not in dataset.column_names and "speaker_id" not in dataset.column_names: - print("No speaker found, adding default 'source' of 0 for all examples\n") - dataset = dataset.add_column("source", ["0"] * len(dataset)) - elif "source" not in dataset.column_names and "speaker_id" in dataset.column_names: - speaker_key = "speaker_id" - - # Resolve audio and text columns - audio_col = next((c for c in dataset.column_names if c in ("audio", "Audio")), None) - text_col = next((c for c in dataset.column_names if c in ("text", "sentence", "transcript")), None) + # Resolve columns from user mapping or hardcoded fallback + resolved = self._resolve_audio_columns(dataset, custom_format_mapping) + audio_col = resolved["audio_col"] + text_col = resolved["text_col"] + speaker_key = resolved["speaker_col"] if audio_col is None: raise ValueError(f"No audio column found in dataset. Columns: {dataset.column_names}") if text_col is None: raise ValueError(f"No text column found in dataset. Columns: {dataset.column_names}") + if speaker_key is None: + print("No speaker found, adding default 'source' of 0 for all examples\n") + dataset = dataset.add_column("source", ["0"] * len(dataset)) + speaker_key = "source" print(f"CSM preprocessing: audio_col='{audio_col}', text_col='{text_col}', speaker_key='{speaker_key}'\n") @@ -773,7 +805,7 @@ class UnslothTrainer: print(f"CSM preprocessing complete: {len(processed)} examples\n") return processed - def _format_audio_vlm_dataset(self, dataset): + def _format_audio_vlm_dataset(self, dataset, custom_format_mapping=None): """Format dataset as audio chat messages for multimodal models (e.g. Gemma 3N). Expects columns: audio (Audio), text (str). @@ -781,15 +813,17 @@ class UnslothTrainer: """ from datasets import Audio - # Detect audio and text columns - cols = dataset.column_names - audio_col = next((c for c in cols if c.lower() in ("audio", "speech")), None) - text_col = next((c for c in cols if c.lower() in ("text", "sentence", "transcript", "transcription")), None) + resolved = self._resolve_audio_columns(dataset, custom_format_mapping) + audio_col = resolved["audio_col"] + text_col = resolved["text_col"] if not audio_col or not text_col: raise ValueError( - f"Audio VLM dataset needs 'audio' and 'text' columns, got: {cols}" + f"Audio VLM dataset needs 'audio' and 'text' columns, got: {dataset.column_names}" ) + # Store resolved audio column name for the collator closure + self._audio_vlm_audio_col = audio_col + # Cast audio to 16kHz (standard for speech models) dataset = dataset.cast_column(audio_col, Audio(sampling_rate=16000)) @@ -818,7 +852,7 @@ class UnslothTrainer: print(f"Audio VLM dataset formatted: {len(dataset)} examples\n") return dataset - def _preprocess_snac_dataset(self, dataset): + def _preprocess_snac_dataset(self, dataset, custom_format_mapping=None): """Preprocess dataset for Orpheus TTS training with SNAC codec. Mirrors Orpheus_(3B)-TTS.ipynb: encode audio with SNAC (24kHz, 3 hierarchical @@ -844,14 +878,14 @@ class UnslothTrainer: END_OF_TEXT = 128009 AUDIO_OFFSET = 128266 - # Resolve audio and text columns (reuse CSM pattern) - cols = dataset.column_names - audio_col = next((c for c in cols if c.lower() in ("audio", "speech")), None) - text_col = next((c for c in cols if c.lower() in ("text", "sentence", "transcript", "transcription")), None) - has_source = "source" in cols + resolved = self._resolve_audio_columns(dataset, custom_format_mapping) + audio_col = resolved["audio_col"] + text_col = resolved["text_col"] + speaker_col = resolved["speaker_col"] + has_source = speaker_col is not None if not audio_col or not text_col: raise ValueError( - f"SNAC dataset needs 'audio' and 'text' columns, got: {cols}" + f"SNAC dataset needs 'audio' and 'text' columns, got: {dataset.column_names}" ) # Get dataset sample rate from first example @@ -923,7 +957,7 @@ class UnslothTrainer: all_codes = deduped # --- Build text tokens (notebook lines 217-224) --- - text_prompt = f"{example['source']}: {text}" if has_source and example.get("source") else text + text_prompt = f"{example[speaker_col]}: {text}" if has_source and example.get(speaker_col) else text text_ids = tokenizer.encode(text_prompt, add_special_tokens=True) text_ids.append(END_OF_TEXT) @@ -979,7 +1013,7 @@ class UnslothTrainer: f"({skipped} skipped)\n") return result_dataset - def _preprocess_bicodec_dataset(self, dataset): + def _preprocess_bicodec_dataset(self, dataset, custom_format_mapping=None): """Preprocess dataset for Spark-TTS training with BiCodec tokenizer. Mirrors Spark_TTS_(0_5B).ipynb: encode audio with BiCodec (semantic + global tokens), @@ -1013,13 +1047,14 @@ class UnslothTrainer: from sparktts.utils.audio import audio_volume_normalize # Resolve audio and text columns - cols = dataset.column_names - audio_col = next((c for c in cols if c.lower() in ("audio", "speech")), None) - text_col = next((c for c in cols if c.lower() in ("text", "sentence", "transcript", "transcription")), None) - has_source = "source" in cols + resolved = self._resolve_audio_columns(dataset, custom_format_mapping) + audio_col = resolved["audio_col"] + text_col = resolved["text_col"] + speaker_col = resolved["speaker_col"] + has_source = speaker_col is not None if not audio_col or not text_col: raise ValueError( - f"BiCodec dataset needs 'audio' and 'text' columns, got: {cols}" + f"BiCodec dataset needs 'audio' and 'text' columns, got: {dataset.column_names}" ) # Load BiCodec tokenizer @@ -1117,7 +1152,7 @@ class UnslothTrainer: ) # Format text with source prefix if available - text_content = f"{example['source']}: {text}" if has_source and example.get("source") else text + text_content = f"{example[speaker_col]}: {text}" if has_source and example.get(speaker_col) else text formatted = "".join([ "<|task_tts|>", @@ -1166,7 +1201,7 @@ class UnslothTrainer: print(f"Sample text length: {len(sample)} chars\n") return result_dataset - def _preprocess_whisper_dataset(self, dataset, eval_split=None): + def _preprocess_whisper_dataset(self, dataset, eval_split=None, custom_format_mapping=None): """Preprocess dataset for Whisper speech-to-text training. Mirrors Whisper.ipynb: extract audio features with Whisper's feature @@ -1177,13 +1212,12 @@ class UnslothTrainer: WHISPER_SAMPLE_RATE = 16000 - # Resolve audio and text columns - cols = dataset.column_names - audio_col = next((c for c in cols if c.lower() in ("audio", "speech")), None) - text_col = next((c for c in cols if c.lower() in ("text", "sentence", "transcript", "transcription")), None) + resolved = self._resolve_audio_columns(dataset, custom_format_mapping) + audio_col = resolved["audio_col"] + text_col = resolved["text_col"] if not audio_col or not text_col: raise ValueError( - f"Whisper dataset needs 'audio' and 'text' columns, got: {cols}" + f"Whisper dataset needs 'audio' and 'text' columns, got: {dataset.column_names}" ) # Cast audio to 16kHz (Whisper's expected sample rate) @@ -1355,20 +1389,21 @@ class UnslothTrainer: # ========== AUDIO MODELS: custom preprocessing ========== if self._audio_type == 'csm': - processed = self._preprocess_csm_dataset(dataset) - # CSM returns a ready-to-train Dataset (not a dict) with no eval + processed = self._preprocess_csm_dataset(dataset, custom_format_mapping) return (processed, None) elif self._audio_type == 'whisper': - train_data, eval_data = self._preprocess_whisper_dataset(dataset, eval_split=eval_split) + train_data, eval_data = self._preprocess_whisper_dataset( + dataset, eval_split=eval_split, custom_format_mapping=custom_format_mapping + ) return (train_data, eval_data) elif self._audio_type == 'snac': - processed = self._preprocess_snac_dataset(dataset) + processed = self._preprocess_snac_dataset(dataset, custom_format_mapping) return (processed, None) elif self._audio_type == 'bicodec': - processed = self._preprocess_bicodec_dataset(dataset) + processed = self._preprocess_bicodec_dataset(dataset, custom_format_mapping) return (processed, None) elif self._audio_type in ('xcodec2', 'dac'): @@ -1376,8 +1411,7 @@ class UnslothTrainer: raise NotImplementedError(f"Audio dataset preprocessing for '{self._audio_type}' not yet implemented") elif self.is_audio_vlm: - # Audio VLM (e.g. Gemma 3N): format as chat messages with audio content - formatted = self._format_audio_vlm_dataset(dataset) + formatted = self._format_audio_vlm_dataset(dataset, custom_format_mapping) return (formatted, None) # ========== FORMAT FIRST ========== @@ -1438,7 +1472,7 @@ class UnslothTrainer: from datasets import get_dataset_split_names load_kwargs = {"path": dataset_source} if subset: - load_kwargs["name"] = subset + load_kwargs["config_name"] = subset available_splits = get_dataset_split_names(**load_kwargs) print(f"Available splits: {available_splits}\n") @@ -1518,6 +1552,22 @@ class UnslothTrainer: self._update_progress(error="Model not loaded") return False + # Pre-import heavy transformers modules on the main thread. + # Unsloth's patched_import hook (deepseek_v3_moe.py) is not thread-safe + # with Python's importlib cache, causing KeyError: 'size' if these are + # first imported inside the worker thread. + import transformers # noqa: F401 – ensures submodules are cached + from transformers import ( # noqa: F401 + Trainer as _HFTrainer, + TrainingArguments as _TrainingArguments, + TrainerCallback as _TrainerCallback, + ) + if self._audio_type == 'whisper': + from transformers import ( # noqa: F401 + Seq2SeqTrainer as _Seq2SeqTrainer, + Seq2SeqTrainingArguments as _Seq2SeqTrainingArguments, + ) + # Start training in separate thread self.training_thread = threading.Thread( target=self._train_worker, @@ -1970,7 +2020,7 @@ class UnslothTrainer: "model": self.model, "train_dataset": train_ds, "data_collator": data_collator, - "tokenizer": self.tokenizer.feature_extractor, + "processing_class": self.tokenizer.feature_extractor, "args": Seq2SeqTrainingArguments(**whisper_training_args), } if eval_dataset: @@ -2293,21 +2343,14 @@ class UnslothTrainer: self._update_progress(error=error_msg, is_training=False) return - elif self.is_vlm: - # Standard VLM collator - print("Using UnslothVisionDataCollator for vision model\n") - from unsloth.trainer import UnslothVisionDataCollator - - FastVisionModel.for_training(self.model) - data_collator = UnslothVisionDataCollator(self.model, self.tokenizer) - print("Vision data collator configured\n") - elif self.is_audio_vlm: # Audio VLM collator (e.g. Gemma 3N with audio data) # Mirrors the collate_fn from Gemma3N_(4B)-Audio notebook print("Configuring audio VLM data collator...\n") processor = self.tokenizer # FastModel returns processor as tokenizer + audio_col_name = getattr(self, '_audio_vlm_audio_col', 'audio') + def audio_vlm_collate_fn(examples): texts = [] audios = [] @@ -2316,7 +2359,7 @@ class UnslothTrainer: example["messages"], tokenize=False, add_generation_prompt=False ).strip() texts.append(text) - audios.append(example["audio"]["array"]) + audios.append(example[audio_col_name]["array"]) batch = processor( text=texts, audio=audios, return_tensors="pt", padding=True @@ -2335,6 +2378,15 @@ class UnslothTrainer: data_collator = audio_vlm_collate_fn print("Audio VLM data collator configured\n") + elif self.is_vlm: + # Standard VLM collator (images) + print("Using UnslothVisionDataCollator for vision model\n") + from unsloth.trainer import UnslothVisionDataCollator + + FastVisionModel.for_training(self.model) + data_collator = UnslothVisionDataCollator(self.model, self.tokenizer) + print("Vision data collator configured\n") + # ========== TRAINING CONFIGURATION ========== # Handle epochs vs max_steps properly max_steps_val = training_args.get('max_steps', 0) @@ -2443,19 +2495,28 @@ class UnslothTrainer: print("Training configuration prepared\n") # ========== TRAINER INITIALIZATION ========== - if self.is_vlm or self.is_audio_vlm: - # VLM: dataset is dict wrapper from format_and_template_dataset - # Audio VLM: dataset is raw Dataset from _format_audio_vlm_dataset + if self.is_audio_vlm: + # Audio VLM (e.g. Gemma 3N + audio): raw Dataset from _format_audio_vlm_dataset + # Notebook uses processing_class=processor.tokenizer (text tokenizer only) + train_dataset = dataset if isinstance(dataset, Dataset) else dataset['dataset'] + processing_class = self.tokenizer.tokenizer if hasattr(self.tokenizer, 'tokenizer') else self.tokenizer + trainer_kwargs = { + "model": self.model, + "train_dataset": train_dataset, + "processing_class": processing_class, + "data_collator": data_collator, + "args": SFTConfig(**config_args), + } + if eval_dataset is not None: + trainer_kwargs["eval_dataset"] = eval_dataset + self.trainer = SFTTrainer(**trainer_kwargs) + elif self.is_vlm: + # Image VLM: dataset is dict wrapper from format_and_template_dataset train_dataset = dataset['dataset'] if isinstance(dataset, dict) else dataset trainer_kwargs = { "model": self.model, -<<<<<<< HEAD - "train_dataset": dataset['dataset'], - "processing_class": self.tokenizer, -======= "train_dataset": train_dataset, - "processing_class": self.tokenizer.tokenizer, ->>>>>>> 0a7e75e (Adding support for audio llms) + "processing_class": self.tokenizer, "data_collator": data_collator, "args": SFTConfig(**config_args), } @@ -2650,11 +2711,7 @@ class UnslothTrainer: progress_callback = ProgressCallback(self) self.trainer.add_callback(progress_callback) -<<<<<<< HEAD - num_samples = len(self.trainer.train_dataset) -======= num_samples = len(dataset['dataset'] if isinstance(dataset, dict) else dataset) ->>>>>>> 0a7e75e (Adding support for audio llms) 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) diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py index 81adef7577..a144cf65cb 100644 --- a/studio/backend/models/datasets.py +++ b/studio/backend/models/datasets.py @@ -28,9 +28,12 @@ class CheckFormatResponse(BaseModel): detected_format: str columns: List[str] is_multimodal: bool = False + is_audio: bool = False multimodal_columns: Optional[List[str]] = None suggested_mapping: Optional[Dict[str, str]] = None detected_image_column: Optional[str] = None + detected_audio_column: Optional[str] = None detected_text_column: Optional[str] = None + detected_speaker_column: Optional[str] = None preview_samples: Optional[List[Dict]] = None total_rows: Optional[int] = None diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index cb1ea75e33..af50b2b72e 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -212,10 +212,13 @@ def check_format(request: CheckFormatRequest): detected_format=result["detected_format"], columns=result["columns"], is_multimodal=result.get("is_multimodal", False), + is_audio=result.get("is_audio", False), multimodal_columns=result.get("multimodal_columns"), suggested_mapping=result.get("suggested_mapping"), detected_image_column=result.get("detected_image_column"), + detected_audio_column=result.get("detected_audio_column"), detected_text_column=result.get("detected_text_column"), + detected_speaker_column=result.get("detected_speaker_column"), preview_samples=preview_samples, total_rows=total_rows, ) diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 3a4d54f93f..6ddd2c753c 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -65,13 +65,22 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: # Auto-detect multimodal data regardless of is_vlm flag multimodal_info = detect_multimodal_dataset(dataset) - if multimodal_info["is_multimodal"]: - is_vlm = True # Route to VLM detection automatically - + is_audio = multimodal_info.get("is_audio", False) + + if multimodal_info["is_multimodal"] and not is_audio: + is_vlm = True # Route to VLM detection for image datasets only + + # Common audio fields for all return paths + audio_fields = { + "is_audio": is_audio, + "detected_audio_column": multimodal_info.get("detected_audio_column"), + "detected_speaker_column": multimodal_info.get("detected_speaker_column"), + } + if is_vlm: vlm_structure = detect_vlm_dataset_structure(dataset) requires_mapping = vlm_structure["format"] == "unknown" - + return { "requires_manual_mapping": requires_mapping, "detected_format": vlm_structure["format"], @@ -81,51 +90,70 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: "detected_text_column": vlm_structure.get("text_column"), "is_multimodal": multimodal_info["is_multimodal"], "multimodal_columns": multimodal_info.get("multimodal_columns"), + **audio_fields, } - else: - # LLM flow - detected = detect_dataset_format(dataset) - - # If format is unknown, try heuristic detection - if detected["format"] == "unknown": - heuristic_mapping = detect_custom_format_heuristic(dataset) - if heuristic_mapping: - # Heuristic succeeded - no manual mapping needed - return { - "requires_manual_mapping": False, - "detected_format": "custom_heuristic", - "columns": columns, - "suggested_mapping": heuristic_mapping, - "detected_image_column": None, - "detected_text_column": None, - "is_multimodal": False, - "multimodal_columns": None, - } - else: - # Both detection and heuristic failed - return { - "requires_manual_mapping": True, - "detected_format": "unknown", - "columns": columns, - "suggested_mapping": None, - "detected_image_column": None, - "detected_text_column": None, - "is_multimodal": False, - "multimodal_columns": None, - } - - # Known format detected + + if is_audio: + # Audio dataset — require manual mapping only when columns can't be auto-detected + detected_audio = multimodal_info.get("detected_audio_column") + detected_text = multimodal_info.get("detected_text_column") + needs_mapping = not detected_audio or not detected_text return { - "requires_manual_mapping": False, - "detected_format": detected["format"], + "requires_manual_mapping": needs_mapping, + "detected_format": "audio", "columns": columns, "suggested_mapping": None, "detected_image_column": None, - "detected_text_column": None, - "is_multimodal": False, - "multimodal_columns": None, + "detected_text_column": multimodal_info.get("detected_text_column"), + "is_multimodal": True, + "multimodal_columns": multimodal_info.get("audio_columns"), + **audio_fields, } + # LLM flow + detected = detect_dataset_format(dataset) + + # If format is unknown, try heuristic detection + if detected["format"] == "unknown": + heuristic_mapping = detect_custom_format_heuristic(dataset) + if heuristic_mapping: + return { + "requires_manual_mapping": False, + "detected_format": "custom_heuristic", + "columns": columns, + "suggested_mapping": heuristic_mapping, + "detected_image_column": None, + "detected_text_column": None, + "is_multimodal": False, + "multimodal_columns": None, + **audio_fields, + } + else: + return { + "requires_manual_mapping": True, + "detected_format": "unknown", + "columns": columns, + "suggested_mapping": None, + "detected_image_column": None, + "detected_text_column": None, + "is_multimodal": False, + "multimodal_columns": None, + **audio_fields, + } + + # Known format detected + return { + "requires_manual_mapping": False, + "detected_format": detected["format"], + "columns": columns, + "suggested_mapping": None, + "detected_image_column": None, + "detected_text_column": None, + "is_multimodal": False, + "multimodal_columns": None, + **audio_fields, + } + # Normalise any format-specific role to canonical chatml (user/assistant/system) _TO_CHATML = { "user": "user", "human": "user", "instruction": "user", diff --git a/studio/backend/utils/datasets/format_detection.py b/studio/backend/utils/datasets/format_detection.py index 9283ea5d55..1a557b5df6 100644 --- a/studio/backend/utils/datasets/format_detection.py +++ b/studio/backend/utils/datasets/format_detection.py @@ -326,45 +326,51 @@ def detect_custom_format_heuristic(dataset): def detect_multimodal_dataset(dataset): """ - Detects if dataset contains multimodal data (images/vision). + Detects if dataset contains multimodal data (images and/or audio). - Two-pass approach: - 1. Column-name heuristic (fast): checks for keywords like 'image', 'img', 'pixel'. - 2. Value-type inspection (reliable): checks if actual values are PIL Images, - bytes with image headers, or HF Image-feature dicts. + Two-pass approach for each modality: + 1. Column-name heuristic (fast): checks for keywords. + 2. Value-type inspection (reliable): checks actual sample values. Returns: dict: { "is_multimodal": bool, "multimodal_columns": list of column names containing image data, - "modality_types": list of detected types (e.g., ["image", "pixel"]) + "modality_types": list of detected types (e.g., ["image", "audio"]), + "is_audio": bool, + "audio_columns": list of column names containing audio data, + "detected_audio_column": str or None, + "detected_text_column": str or None, } """ sample = next(iter(dataset)) column_names = list(sample.keys()) - # Keywords that indicate multimodal/image data - multimodal_keywords = [ + # Keywords that indicate image data + image_keywords = [ 'image', 'img', 'pixel', 'jpg', 'jpeg', 'png', 'webp', 'bmp', 'gif', 'tiff', 'svg', 'photo', 'pic', 'picture', 'visual', ] + # Keywords that indicate audio data + audio_keywords = ['audio', 'speech', 'wav', 'waveform', 'sound'] + multimodal_columns = [] + audio_columns = [] modality_types = set() - # ── Pass 1: column-name heuristic ─────────────────────── + # ── Image detection ───────────────────────────────────── + # Pass 1: column-name heuristic for col_name in column_names: col_lower = col_name.lower() - - for keyword in multimodal_keywords: + for keyword in image_keywords: if keyword in col_lower: multimodal_columns.append(col_name) modality_types.add(keyword) - break # Don't check other keywords for this column + break - # ── Pass 2: inspect actual values ─────────────────────── - # Catches columns with non-obvious names (e.g. "jpg", "photo", "pic") + # Pass 2: inspect actual values already_detected = set(multimodal_columns) for col_name in column_names: if col_name in already_detected: @@ -374,10 +380,61 @@ def detect_multimodal_dataset(dataset): multimodal_columns.append(col_name) modality_types.add("image") + # ── Audio detection ───────────────────────────────────── + # Pass 1: column-name heuristic + for col_name in column_names: + col_lower = col_name.lower() + for keyword in audio_keywords: + if keyword in col_lower: + audio_columns.append(col_name) + modality_types.add("audio") + break + + # Pass 2: inspect actual values (catches non-obvious column names) + already_audio = set(audio_columns) + for col_name in column_names: + if col_name in already_audio: + continue + value = sample[col_name] + if _is_audio_value(value): + audio_columns.append(col_name) + modality_types.add("audio") + + # Filter out columns that are actually audio from the image list + # (e.g. a column named "audio" with {"bytes", "path"} could match _is_image_value) + if audio_columns: + audio_set = set(audio_columns) + multimodal_columns = [c for c in multimodal_columns if c not in audio_set] + + # Detect text column for audio datasets + detected_text_col = None + if audio_columns: + text_keywords = ['text', 'sentence', 'transcript', 'transcription', 'label'] + for col_name in column_names: + if col_name.lower() in text_keywords: + detected_text_col = col_name + break + + is_audio = len(audio_columns) > 0 + + # Detect speaker_id column for TTS datasets (CSM, Orpheus, Spark) + detected_speaker_col = None + if audio_columns: + speaker_keywords = ['source', 'speaker', 'speaker_id'] + for col_name in column_names: + if col_name.lower() in speaker_keywords: + detected_speaker_col = col_name + break + return { - "is_multimodal": len(multimodal_columns) > 0, + "is_multimodal": len(multimodal_columns) > 0 or is_audio, "multimodal_columns": multimodal_columns, - "modality_types": list(modality_types) + "modality_types": list(modality_types), + "is_audio": is_audio, + "audio_columns": audio_columns, + "detected_audio_column": audio_columns[0] if audio_columns else None, + "detected_text_column": detected_text_col, + "detected_speaker_column": detected_speaker_col, } @@ -395,9 +452,16 @@ def _is_image_value(value) -> bool: pass # HF datasets Image feature stores decoded images as PIL or dicts with - # {"bytes": b"...", "path": "..."} when not yet decoded + # {"bytes": b"...", "path": "..."} when not yet decoded. + # Exclude audio dicts (decoded audio has "array" + "sampling_rate"). if isinstance(value, dict): + if "array" in value and "sampling_rate" in value: + return False # This is audio, not image if "bytes" in value and "path" in value: + # Check path extension to exclude audio files + path = value.get("path") or "" + if isinstance(path, str) and any(path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS): + return False return True # Raw bytes with a known image magic header @@ -407,6 +471,29 @@ def _is_image_value(value) -> bool: return False +_AUDIO_EXTENSIONS = ( + ".wav", ".mp3", ".flac", ".ogg", ".opus", ".m4a", ".aac", ".wma", ".webm", +) + + +def _is_audio_value(value) -> bool: + """Check if a single sample value looks like audio data.""" + if value is None: + return False + + # HF datasets Audio feature: decoded → {"array": np.ndarray, "sampling_rate": int} + if isinstance(value, dict): + if "array" in value and "sampling_rate" in value: + return True + # Undecoded/streaming → {"bytes": b"...", "path": "some.wav"} + if "bytes" in value or "path" in value: + path = value.get("path") or "" + if isinstance(path, str) and any(path.lower().endswith(ext) for ext in _AUDIO_EXTENSIONS): + return True + + return False + + def _has_image_header(data: bytes) -> bool: """Quick magic-byte check for common image formats.""" if len(data) < 4: diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx index c1455c20d2..a0938f1338 100644 --- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx @@ -16,6 +16,7 @@ const CHATML_ROLES = ["system", "user", "assistant"] as const; const ALPACA_ROLES = ["instruction", "input", "output"] as const; const SHAREGPT_ROLES = ["system", "human", "gpt"] as const; const VLM_ROLES = ["image", "text"] as const; +const AUDIO_ROLES = ["audio", "text", "speaker_id"] as const; const ROLE_LABELS: Record = { system: "System", @@ -28,9 +29,12 @@ const ROLE_LABELS: Record = { output: "Output", image: "Image", text: "Text", + audio: "Audio", + speaker_id: "Speaker ID", }; -export function getAvailableRoles(isVlm: boolean, format?: string): readonly string[] { +export function getAvailableRoles(isVlm: boolean, format?: string, isAudio?: boolean): readonly string[] { + if (isAudio) return AUDIO_ROLES; if (isVlm) return VLM_ROLES; if (format === "alpaca") return ALPACA_ROLES; if (format === "sharegpt") return SHAREGPT_ROLES; @@ -41,8 +45,10 @@ export function isMappingComplete( mapping: Record, isVlm: boolean, format?: string, + isAudio?: boolean, ): boolean { const roles = new Set(Object.values(mapping)); + if (isAudio) return roles.has("audio") && roles.has("text"); if (isVlm) return roles.has("image") && roles.has("text"); if (format === "alpaca") return roles.has("instruction") && roles.has("output"); if (format === "sharegpt") return roles.has("human") && roles.has("gpt"); @@ -91,16 +97,19 @@ export function DatasetMappingCard({ mappingOk: boolean; autoDetected?: boolean; isVlm?: boolean; + isAudio?: boolean; format?: string; }) { const entries = Object.entries(mapping); - const requiredLabel = isVlm - ? "image and text" - : format === "alpaca" - ? "instruction and output" - : format === "sharegpt" - ? "human and gpt" - : "user and assistant"; + const requiredLabel = isAudio + ? "audio and text" + : isVlm + ? "image and text" + : format === "alpaca" + ? "instruction and output" + : format === "sharegpt" + ? "human and gpt" + : "user and assistant"; return (
= { instruction: "user", input: "system", output: "assistant", human: "user", gpt: "assistant", image: "image", text: "text", + audio: "audio", speaker_id: "speaker_id", }; /** Chatml → format-specific role names (only for formats that differ). */ @@ -257,10 +267,18 @@ export function deriveDefaultMapping( data: CheckFormatResponse, isVlm: boolean, format?: string, + isAudio?: boolean, ): Record { if (data.suggested_mapping) { return remapRolesForFormat({ ...data.suggested_mapping }, format); } + if (isAudio) { + const result: Record = {}; + if (data.detected_audio_column) result[data.detected_audio_column] = "audio"; + if (data.detected_text_column) result[data.detected_text_column] = "text"; + if (data.detected_speaker_column) result[data.detected_speaker_column] = "speaker_id"; + return result; + } if (isVlm) { const result: Record = {}; if (data.detected_image_column) result[data.detected_image_column] = "image"; diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx index 00a738bc64..7537534329 100644 --- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx @@ -64,13 +64,14 @@ export function DatasetPreviewDialog({ // If the backend reports multimodal data, treat as VLM even if the prop // hasn't caught up yet (isDatasetMultimodal may still be null in the store). - const effectiveIsVlm = isVlm || !!data?.is_multimodal; + const effectiveIsAudio = !!data?.is_audio; + const effectiveIsVlm = !effectiveIsAudio && (isVlm || !!data?.is_multimodal); const hasHeuristicMapping = !data?.requires_manual_mapping && !!data?.suggested_mapping; const mappingEnabled = !!data?.requires_manual_mapping || hasHeuristicMapping; const showMappingFooter = mode === "mapping" && mappingEnabled; - const mappingOk = isMappingComplete(manualMapping, effectiveIsVlm, datasetFormat); - const availableRoles = getAvailableRoles(effectiveIsVlm, datasetFormat); + const mappingOk = isMappingComplete(manualMapping, effectiveIsVlm, datasetFormat, effectiveIsAudio); + const availableRoles = getAvailableRoles(effectiveIsVlm, datasetFormat, effectiveIsAudio); const isHfDataset = !!datasetName && datasetName.includes("/"); // When format changes, remap existing mapping roles to the new format's role names @@ -150,7 +151,7 @@ export function DatasetPreviewDialog({ if (!data?.requires_manual_mapping && !data?.suggested_mapping) return; // Don't overwrite if mapping already has entries if (Object.keys(manualMapping).length > 0) return; - const derived = deriveDefaultMapping(data, effectiveIsVlm, datasetFormat); + const derived = deriveDefaultMapping(data, effectiveIsVlm, datasetFormat, effectiveIsAudio); if (Object.keys(derived).length === 0) return; setManualMapping(derived); }, [open, datasetName, data, effectiveIsVlm, datasetFormat, manualMapping, setManualMapping]); @@ -346,6 +347,7 @@ export function DatasetPreviewDialog({ mappingOk={mappingOk} autoDetected={hasHeuristicMapping} isVlm={effectiveIsVlm} + isAudio={effectiveIsAudio} format={datasetFormat} /> )} diff --git a/studio/frontend/src/features/training/hooks/use-training-actions.ts b/studio/frontend/src/features/training/hooks/use-training-actions.ts index bfe035cf4a..10f67493dc 100644 --- a/studio/frontend/src/features/training/hooks/use-training-actions.ts +++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts @@ -47,12 +47,22 @@ export function useTrainingActions() { isVlm, }); - // Backend auto-detects multimodal even if we didn't know yet - if (check.is_multimodal && config.isVisionModel) { + // Backend auto-detects multimodal/audio from dataset content. + // Sync these flags into the store so buildTrainingStartPayload picks them up. + const isAudio = !!check.is_audio; + const isMultimodal = !!check.is_multimodal; + + if (isMultimodal && config.isVisionModel) { isVlm = true; } + if (isMultimodal !== config.isDatasetMultimodal || isAudio !== config.isDatasetAudio) { + useTrainingConfigStore.setState({ + isDatasetMultimodal: isMultimodal, + isDatasetAudio: isAudio, + }); + } - if (check.requires_manual_mapping && !hasManualMapping(config, isVlm)) { + if (check.requires_manual_mapping && !hasManualMapping(config, isVlm, isAudio)) { // Pre-fill from suggested_mapping or VLM detected columns const hint: Record = {}; if (check.suggested_mapping) { @@ -60,6 +70,10 @@ export function useTrainingActions() { for (const [col, role] of Object.entries(check.suggested_mapping)) { hint[col] = table ? (table[role] ?? role) : role; } + } else if (isAudio) { + if (check.detected_audio_column) hint[check.detected_audio_column] = "audio"; + if (check.detected_text_column) hint[check.detected_text_column] = "text"; + if (check.detected_speaker_column) hint[check.detected_speaker_column] = "speaker_id"; } else if (isVlm) { if (check.detected_image_column) hint[check.detected_image_column] = "image"; if (check.detected_text_column) hint[check.detected_text_column] = "text"; @@ -75,7 +89,8 @@ export function useTrainingActions() { } } - const payload = buildTrainingStartPayload(config); + // Re-read config after potential store updates from dataset check + const payload = buildTrainingStartPayload(useTrainingConfigStore.getState()); const response = await startTraining(payload); if (response.status === "error") { @@ -143,12 +158,11 @@ function getDatasetName(config: TrainingConfigState): string | null { : config.uploadedFile; } -function hasManualMapping(config: TrainingConfigState, isVlm = false): boolean { +function hasManualMapping(config: TrainingConfigState, isVlm = false, isAudio = false): boolean { const mapping = config.datasetManualMapping; const roles = new Set(Object.values(mapping)); - if (isVlm) { - return roles.has("image") && roles.has("text"); - } + if (isAudio) return roles.has("audio") && roles.has("text"); + if (isVlm) return roles.has("image") && roles.has("text"); const fmt = config.datasetFormat; if (fmt === "alpaca") return roles.has("instruction") && roles.has("output"); if (fmt === "sharegpt") return roles.has("human") && roles.has("gpt"); diff --git a/studio/frontend/src/features/training/lib/model-defaults.ts b/studio/frontend/src/features/training/lib/model-defaults.ts index e68516e7de..35ce562dbf 100644 --- a/studio/frontend/src/features/training/lib/model-defaults.ts +++ b/studio/frontend/src/features/training/lib/model-defaults.ts @@ -4,7 +4,6 @@ import type { TrainingConfigState } from "../types/config"; type ModelDefaultsPatch = Partial< Pick< TrainingConfigState, - | "isDatasetAudio" | "epochs" | "contextLength" | "learningRate" @@ -80,9 +79,6 @@ export function mapBackendModelConfigToTrainingPatch( const lora = config.lora; const logging = config.logging; - // Audio models: set isDatasetAudio based on audio_type from YAML - patch.isDatasetAudio = typeof config.audio_type === "string" && config.audio_type.length > 0; - const maxSeqLength = toNumber(training?.max_seq_length); if (maxSeqLength !== undefined) patch.contextLength = maxSeqLength; diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index e72acef5d2..4e50ed6ae3 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -175,8 +175,10 @@ export const useTrainingConfigStore = create()( .then((res) => { if (controller.signal.aborted) return; const isMultimodal = !!res.is_multimodal; + const isAudio = !!res.is_audio; const updates: Record = { isDatasetMultimodal: isMultimodal, + isDatasetAudio: isAudio, isCheckingDataset: false, }; if (!_trainOnCompletionsManuallySet) { diff --git a/studio/frontend/src/features/training/types/datasets.ts b/studio/frontend/src/features/training/types/datasets.ts index 96cf699f89..c1880116ca 100644 --- a/studio/frontend/src/features/training/types/datasets.ts +++ b/studio/frontend/src/features/training/types/datasets.ts @@ -4,10 +4,13 @@ export type CheckFormatResponse = { columns: string[]; suggested_mapping?: Record | null; detected_image_column?: string | null; + detected_audio_column?: string | null; detected_text_column?: string | null; + detected_speaker_column?: string | null; preview_samples?: Record[] | null; total_rows?: number | null; is_multimodal?: boolean; + is_audio?: boolean; multimodal_columns?: string[] | null; }; From c48437848d20f81575b77ed5ccedf82ee2883f7f Mon Sep 17 00:00:00 2001 From: Manan17 Date: Sun, 1 Mar 2026 02:30:31 +0000 Subject: [PATCH 15/83] 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} + +
+ )}
)}