Run name
New Finetune
+
+ Engine
+
+ torch
+ mlx
+
+
Base Quant.
@@ -6172,6 +6179,7 @@ Download Audio Selection
const steps = parseInt(document.getElementById('new-ft-steps').value, 10);
const batch = parseInt(document.getElementById('new-ft-batch').value, 10);
const base_model = document.getElementById('new-ft-base-model').value;
+ const engine = (document.getElementById('new-ft-engine') || {}).value || 'torch';
const lora_type = document.getElementById('new-ft-lora-type').value;
const rank = parseInt(document.getElementById('new-ft-rank').value, 10);
const alphaStr = document.getElementById('new-ft-alpha').value.trim();
@@ -6182,7 +6190,7 @@ Download Audio Selection
if (!name || _newFtSelectedGpu === null || isNaN(steps) || isNaN(batch)) return;
closeNewFinetuneModal();
const payload = {name, gpu: _newFtSelectedGpu, max_steps: steps, batch_size: batch,
- base_model, lora_type, rank, checkpoint_every, demo_every};
+ base_model, engine, lora_type, rank, checkpoint_every, demo_every};
if (alphaStr) payload.alpha = Number(alphaStr);
const includeStr = document.getElementById('new-ft-include').value.trim();
const excludeStr = document.getElementById('new-ft-exclude').value.trim();
diff --git a/dashboard/server.py b/dashboard/server.py
index bcc3682..fc50bd4 100644
--- a/dashboard/server.py
+++ b/dashboard/server.py
@@ -1458,12 +1458,16 @@ def launch(self, checkpoint_path, gpu, run_id=None, checkpoint_name=None, title=
checkpoint_path = _bash_path(_app_path(checkpoint_path))
if not title:
title = checkpoint_name or Path(checkpoint_path).name
- # Resolve base model from run record
+ # Resolve base model + engine from run record. engine (torch|mlx)
+ # is threaded onto the run_gradio cmd so MLX-trained LoRAs launch
+ # the MLX gradio, matching how they were trained.
base_model = "sa3-medium"
+ engine = "torch"
if run_id:
run = registry.get_run(run_id)
if run:
base_model = run.get("base_model", "sa3-medium")
+ engine = run.get("engine", "torch")
gmi = _get_model_info(base_model)
# Choose ARC vs RF base checkpoint
use_arc_full = (model_variant == "arc" and gmi.get("arc_type") == "full_model")
@@ -1529,6 +1533,7 @@ def launch(self, checkpoint_path, gpu, run_id=None, checkpoint_name=None, title=
f"--model-config {_bash_quote(config_path_model)} "
f"--ckpt-path {_bash_quote(ckpt_path_model)} "
f"{lora_args} "
+ f"--engine {engine} "
f"--model-half "
f"--title {shlex.quote(title)}"
f"{default_prompt_arg}"
@@ -3716,6 +3721,12 @@ def _handle_new_finetune(self, body):
max_steps = int(body.get("max_steps", 20000))
batch_size = int(body.get("batch_size", 8))
base_model = body.get("base_model", "sa3-medium")
+ # Engine: torch (default, in-process PyTorch loop) or mlx (Apple-Silicon,
+ # delegates to the sibling MLX trainer). Threaded into the launch cmd as
+ # --engine and stored on the run so resume + "use this checkpoint" match.
+ engine = str(body.get("engine", "torch") or "torch").strip().lower()
+ if engine not in ("torch", "mlx"):
+ engine = "torch"
lora_type = body.get("lora_type", "lora") # lora, dora, bora, lora-xs
rank = int(body.get("rank", 16))
checkpoint_every = int(body.get("checkpoint_every", 1000))
@@ -4013,6 +4024,7 @@ def _handle_new_finetune(self, body):
f" --max-steps {max_steps}"
f" --gradient-clip-val 1.0"
f" --logger ''"
+ f" --engine {engine}"
)
gpu_env = f"CUDA_VISIBLE_DEVICES={gpu} "
@@ -4027,10 +4039,13 @@ def _handle_new_finetune(self, body):
# aren't installed — otherwise the trainer subprocess crashes mid-
# startup with a cryptic TypeError when the model config hits a
# conditioner whose kwargs don't match the older backend's API.
- backend_err = _missing_backend_error_for_model(base_model)
- if backend_err:
- self._json_response({"error": backend_err}, status=400)
- return
+ # engine=mlx trains in a separate MLX venv/checkout, so the in-venv
+ # torch backend need not be installed — skip the torch-backend gate.
+ if engine != "mlx":
+ backend_err = _missing_backend_error_for_model(base_model)
+ if backend_err:
+ self._json_response({"error": backend_err}, status=400)
+ return
backend_env = _backend_env_for_model(base_model)
# UNDERFIT_LOG_PATH lets lora_train.py write a sidecar .exit on
# any unhandled exception — robust to buffering issues that can cause
@@ -4073,6 +4088,7 @@ def _handle_new_finetune(self, body):
"gpu": gpu,
"restart_cmd": restart_cmd,
"base_model": base_model,
+ "engine": engine,
"dataset_id": dataset_id,
"dataset_history": [{
"dataset_id": dataset_id,
diff --git a/defaults.ini b/defaults.ini
index 4cbb6d2..6142288 100644
--- a/defaults.ini
+++ b/defaults.ini
@@ -6,6 +6,14 @@
# stable-audio-tools used so backend code that expects these keys keeps
# working.
+# Training/inference engine: torch | mlx
+# torch (default) — the in-process PyTorch loop (MPS/CUDA/CPU), unchanged.
+# mlx — Apple-Silicon-only; shells out to the separate MLX
+# trainer/gradio in the sibling stable-audio-3 checkout
+# (see underfit/backends/mlx_engine.py). Overridable via
+# --engine or the UNDERFIT_ENGINE env var.
+engine = torch
+
# Name of the run (dashboard overrides via --name)
name = underfit-run
diff --git a/lora_train.py b/lora_train.py
index 8e08e10..095a653 100644
--- a/lora_train.py
+++ b/lora_train.py
@@ -100,6 +100,14 @@ def get_all_args(defaults_file="defaults.ini"):
p.add_argument("--wandb-config", default=None)
p.add_argument("--backend", default=None,
help="sat | sa3 (default: env UNDERFIT_BACKEND or auto)")
+ # Pre-register --engine so the generic defaults loop below (which would add
+ # it from defaults.ini) is swallowed by its ArgumentError guard — this
+ # explicit registration keeps default=None so we can apply the
+ # CLI > env > defaults.ini > torch precedence after parsing.
+ p.add_argument("--engine", default=None, choices=["torch", "mlx"],
+ help="torch | mlx (default: env UNDERFIT_ENGINE, defaults.ini, "
+ "else torch). mlx runs the Apple-Silicon MLX trainer in "
+ "the sibling stable-audio-3 checkout.")
for key, value in defaults.items():
arg_name = f"--{key.replace('_', '-')}"
try:
@@ -124,6 +132,13 @@ def get_all_args(defaults_file="defaults.ini"):
setattr(args, key, float(val))
except ValueError:
pass
+ # Engine selection precedence: CLI --engine > env UNDERFIT_ENGINE >
+ # defaults.ini > torch.
+ if getattr(args, "engine", None) is None:
+ args.engine = os.environ.get("UNDERFIT_ENGINE") or defaults.get("engine") or "torch"
+ args.engine = str(args.engine).strip().strip("'\"").lower()
+ if args.engine not in ("torch", "mlx"):
+ raise SystemExit(f"--engine must be 'torch' or 'mlx', got {args.engine!r}")
return args
@@ -131,6 +146,13 @@ def main():
args = get_all_args()
if os.environ.get("SLURM_PROCID") is not None:
args.seed = (args.seed or 0) + int(os.environ["SLURM_PROCID"])
+
+ # --- MLX engine: delegate to the sibling stable-audio-3 MLX trainer. ---
+ # Purely additive; the torch path below is unchanged for engine=torch.
+ if getattr(args, "engine", "torch") == "mlx":
+ from underfit.backends import mlx_engine
+ sys.exit(mlx_engine.run_mlx_training(args))
+
# Pre-warn (and quiet torch's noisy autotune warnings) on pre-Ampere GPUs.
from underfit.utils import check_attention_compute_capability, check_attention_backends
check_attention_compute_capability()
diff --git a/run_gradio.py b/run_gradio.py
index f47a270..9d69b73 100644
--- a/run_gradio.py
+++ b/run_gradio.py
@@ -40,6 +40,16 @@
def main(args):
+ # --- MLX engine: delegate to the sibling stable-audio-3 MLX gradio. ---
+ # Purely additive; the torch path below is unchanged for engine=torch.
+ engine = (getattr(args, "engine", None) or os.environ.get("UNDERFIT_ENGINE") or "torch").strip().lower()
+ if engine == "mlx":
+ from underfit.backends import mlx_engine
+ model = mlx_engine.resolve_dit_model(args.model_config, args.pretrained_name)
+ port_env = os.environ.get("GRADIO_SERVER_PORT")
+ port = int(port_env) if port_env and port_env.isdigit() else None
+ sys.exit(mlx_engine.run_mlx_gradio(model, args.lora_ckpt_path, share=True, port=port))
+
backend = get_backend(args.backend)
print(f"Using backend: {backend.NAME}", flush=True)
@@ -73,6 +83,10 @@ def main(args):
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Run gradio interface (Underfit)")
parser.add_argument("--backend", default=None, help="sat | sa3 (default: env UNDERFIT_BACKEND or auto)")
+ parser.add_argument("--engine", default=None, choices=["torch", "mlx"],
+ help="torch | mlx (default: env UNDERFIT_ENGINE, else torch). "
+ "mlx launches the Apple-Silicon MLX gradio in the sibling "
+ "stable-audio-3 checkout.")
parser.add_argument("--pretrained-name", type=str, required=False)
parser.add_argument("--model-config", type=str, required=False)
parser.add_argument("--ckpt-path", type=str, required=False)
diff --git a/underfit/backends/mlx_engine.py b/underfit/backends/mlx_engine.py
new file mode 100644
index 0000000..4dd4b34
--- /dev/null
+++ b/underfit/backends/mlx_engine.py
@@ -0,0 +1,438 @@
+"""MLX engine adapter for Underfit.
+
+Apple-Silicon-only alternate training/inference engine. Instead of running the
+in-process torch training loop (underfit.training.loop), the MLX engine shells
+out to a *separate* MLX trainer/gradio that lives in a sibling stable-audio-3
+checkout (`/optimized/mlx`). Everything here is a thin launcher: it maps
+Underfit's args/config onto the MLX scripts' locked CLIs, streams the trainer's
+stdout back in the dashboard's expected format, and propagates exit codes.
+
+This module never imports torch or stable_audio_3 — it only reads the run's
+JSON configs and spawns the MLX venv python. It is only reached when
+`--engine mlx` (default is torch, and the torch path is untouched).
+
+Environment overrides:
+ UNDERFIT_MLX_ROOT MLX code root (default /optimized/mlx)
+ UNDERFIT_MLX_PYTHON MLX venv python (default /.venv/bin/python)
+ UNDERFIT_MLX_BASE_WEIGHTS base DiT weights npz (default
+ /models/mlx/dit_-base_f16.npz)
+"""
+import json
+import os
+import re
+import shlex
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+
+
+# Mirror the sibling-checkout resolution in backends/sa3.py without importing
+# it (sa3.py pulls in torch at import time; the MLX path must stay torch-free).
+_HERE = Path(__file__).resolve()
+_APP_ROOT = _HERE.parent.parent.parent # underfit repo root
+_SA3_LOCAL = str(_HERE.parent.parent.parent.parent / "stable-audio-3")
+
+_SUPPORTED_DITS = ("sm-music", "sm-sfx", "medium")
+
+# Underfit's tqdm-postfix step line the MLX trainer emits, per its locked
+# contract: "step {N} train/loss {L:.6f} train/lr {lr:.3e} epoch {E} (…)"
+_STEP_LINE_RE = re.compile(
+ r"^step (\d+)\s+train/loss (\S+)\s+train/lr (\S+)\s+epoch (\d+)"
+)
+
+_STEP_TOKEN_RE = re.compile(r"step=(\d+)")
+_EPOCH_TOKEN_RE = re.compile(r"epoch=(\d+)")
+
+
+# --------------------------------------------------------------------------- #
+# Path / name resolution
+# --------------------------------------------------------------------------- #
+def resolve_mlx_paths():
+ """Return (mlx_root, mlx_python).
+
+ Uses UNDERFIT_MLX_ROOT / UNDERFIT_MLX_PYTHON when set, else the sibling
+ stable-audio-3 checkout defaults. Raises FileNotFoundError (mentioning the
+ env vars) if the trainer script or venv python is missing.
+ """
+ mlx_root = os.environ.get("UNDERFIT_MLX_ROOT") or os.path.join(
+ _SA3_LOCAL, "optimized", "mlx"
+ )
+ mlx_root = os.path.abspath(mlx_root)
+ mlx_python = os.environ.get("UNDERFIT_MLX_PYTHON") or os.path.join(
+ mlx_root, ".venv", "bin", "python"
+ )
+ trainer = os.path.join(mlx_root, "scripts", "lora_train_mlx.py")
+ if not os.path.isfile(trainer):
+ raise FileNotFoundError(
+ f"MLX trainer script not found at {trainer!r}. Point "
+ f"UNDERFIT_MLX_ROOT at the stable-audio-3/optimized/mlx checkout "
+ f"(resolved root: {mlx_root!r})."
+ )
+ if not (os.path.isfile(mlx_python) or shutil.which(mlx_python)):
+ raise FileNotFoundError(
+ f"MLX venv python not found at {mlx_python!r}. Set "
+ f"UNDERFIT_MLX_PYTHON to the MLX venv interpreter (default is "
+ f"/.venv/bin/python)."
+ )
+ return mlx_root, mlx_python
+
+
+def map_model_name(base_model):
+ """Map an Underfit base-model key -> the MLX trainer's --dit value.
+
+ Accepts both the dashboard form ("sa3-medium") and the already-stripped
+ form ("medium"). Raises ValueError for unsupported models.
+ """
+ if not base_model:
+ raise ValueError(
+ "engine=mlx: the model config has no 'base_model' — cannot pick "
+ "the MLX --dit value."
+ )
+ name = str(base_model).strip()
+ if name.startswith("sa3-"):
+ name = name[len("sa3-"):]
+ if name not in _SUPPORTED_DITS:
+ raise ValueError(
+ f"engine=mlx does not support base_model={base_model!r}. "
+ f"Supported: sa3-sm-music, sa3-sm-sfx, sa3-medium."
+ )
+ return name
+
+
+def resolve_base_weights(dit_model):
+ """Path to the MLX BASE-checkpoint DiT weights npz for a --dit value.
+
+ UNDERFIT_MLX_BASE_WEIGHTS overrides; default is
+ /models/mlx/dit_-base_f16.npz.
+ """
+ override = os.environ.get("UNDERFIT_MLX_BASE_WEIGHTS")
+ if override:
+ return os.path.abspath(override)
+ mlx_root, _ = resolve_mlx_paths()
+ return os.path.join(mlx_root, "models", "mlx", f"dit_{dit_model}-base_f16.npz")
+
+
+def resolve_dit_model(model_config_path=None, pretrained_name=None):
+ """Derive the MLX --dit value for gradio from a model-config path or a
+ pretrained name. The dashboard passes the run's _model.json, which carries
+ a 'base_model' key."""
+ if model_config_path:
+ try:
+ mc = _load_json(model_config_path)
+ except (OSError, ValueError):
+ mc = {}
+ base = mc.get("base_model")
+ if base:
+ return map_model_name(base)
+ if pretrained_name:
+ return map_model_name(pretrained_name)
+ raise ValueError(
+ "engine=mlx gradio: could not determine the --dit model. Pass a "
+ "--model-config whose JSON has 'base_model', or --pretrained-name."
+ )
+
+
+# --------------------------------------------------------------------------- #
+# Small helpers
+# --------------------------------------------------------------------------- #
+def _load_json(path):
+ with open(path) as f:
+ return json.load(f)
+
+
+def _resolve_app_relative_path(value):
+ """Match backends/sa3._resolve_app_relative_path so the MLX trainer reads
+ the same latents dir the torch dataloader would."""
+ if not value:
+ return value
+ p = Path(value)
+ if p.is_absolute():
+ return str(p)
+ cwd_path = Path.cwd() / p
+ app_path = _APP_ROOT / p
+ if cwd_path.exists() and not app_path.exists():
+ return str(cwd_path)
+ return str(app_path)
+
+
+def _fmt(v):
+ if isinstance(v, bool):
+ return "true" if v else "false"
+ return str(v)
+
+
+def _add(cmd, flag, value):
+ """Append `flag value` unless value is None or an empty/whitespace string."""
+ if value is None:
+ return
+ if isinstance(value, str) and not value.strip():
+ return
+ cmd.extend([flag, _fmt(value)])
+
+
+def _csv(value):
+ """Normalize an include/exclude field (list or string) to a CSV string, or
+ None when empty."""
+ if value is None:
+ return None
+ if isinstance(value, (list, tuple)):
+ items = [str(x).strip() for x in value if str(x).strip()]
+ return ",".join(items) if items else None
+ s = str(value).strip()
+ return s or None
+
+
+def _as_list(value):
+ if value is None:
+ return []
+ if isinstance(value, (list, tuple)):
+ return list(value)
+ return [value]
+
+
+def _dit_model_from_config(model_config):
+ return map_model_name(model_config.get("base_model"))
+
+
+def _latents_dir_from_dataset(dataset_config):
+ datasets = dataset_config.get("datasets") or []
+ if not datasets:
+ raise ValueError(
+ "engine=mlx: dataset config has no 'datasets' entry to read "
+ "--latents-dir from."
+ )
+ path = datasets[0].get("path")
+ if not path:
+ raise ValueError(
+ "engine=mlx: dataset config datasets[0] has no 'path' for "
+ "--latents-dir."
+ )
+ return _resolve_app_relative_path(path)
+
+
+def _resolve_lr(training, args):
+ lr = None
+ opt = training.get("optimizer_configs")
+ if isinstance(opt, dict):
+ lr = (
+ opt.get("diffusion", {})
+ .get("optimizer", {})
+ .get("config", {})
+ .get("lr")
+ )
+ if lr is None:
+ lr = getattr(args, "lr", None)
+ if lr is None:
+ raise ValueError(
+ "engine=mlx: could not resolve the learning rate. Set "
+ "training.optimizer_configs.diffusion.optimizer.config.lr in the "
+ "model config (the dashboard writes this from the LR field)."
+ )
+ return lr
+
+
+def _dist_shift_from_config(model_config):
+ diffusion = (model_config.get("model", {}) or {}).get("diffusion", {}) or {}
+ opts = diffusion.get("distribution_shift_options") or {}
+ return opts.get("type")
+
+
+def _parse_filename_offsets(path):
+ """(step, epoch) parsed from a checkpoint filename; either may be None.
+ Mirrors underfit.training.loop._parse_filename_offsets."""
+ base = os.path.basename(str(path))
+ s = _STEP_TOKEN_RE.search(base)
+ e = _EPOCH_TOKEN_RE.search(base)
+ return (int(s.group(1)) if s else None,
+ int(e.group(1)) if e else None)
+
+
+def _resolve_offsets(training, resume_path):
+ """(step_offset, epoch_offset) mirroring the torch loop's resolution:
+ explicit config keys win (even at 0), else parse from the resume filename."""
+ has_step = "step_offset" in training
+ has_epoch = "epoch_offset" in training
+ step = int(training.get("step_offset", 0) or 0) if has_step else None
+ epoch = int(training.get("epoch_offset", 0) or 0) if has_epoch else None
+ if (step is None or epoch is None) and resume_path:
+ f_step, f_epoch = _parse_filename_offsets(resume_path)
+ if step is None:
+ step = f_step
+ if epoch is None:
+ epoch = f_epoch
+ return step, epoch
+
+
+def _write_demo_config(demo):
+ """Write the demo block to a temp json and return its path."""
+ fd, path = tempfile.mkstemp(prefix="underfit_mlx_demo_", suffix=".json")
+ with os.fdopen(fd, "w") as f:
+ json.dump(demo, f)
+ return path
+
+
+# --------------------------------------------------------------------------- #
+# Command builders
+# --------------------------------------------------------------------------- #
+def build_trainer_cmd(args, base_weights):
+ """Map Underfit's args + run configs onto the locked lora_train_mlx.py CLI.
+
+ Reads args.model_config / args.dataset_config (paths). base_weights is the
+ resolved --dit-weights npz path (see resolve_base_weights).
+ """
+ mlx_root, mlx_python = resolve_mlx_paths()
+ trainer = os.path.join(mlx_root, "scripts", "lora_train_mlx.py")
+
+ model_config = _load_json(args.model_config)
+ dataset_config = _load_json(args.dataset_config)
+ training = model_config.get("training", {}) or {}
+ lora_cfg = training.get("lora_config", {}) or {}
+
+ dit = _dit_model_from_config(model_config)
+ latents_dir = _latents_dir_from_dataset(dataset_config)
+ lr = _resolve_lr(training, args)
+
+ cmd = [
+ mlx_python, trainer,
+ "--dit", dit,
+ "--dit-weights", str(base_weights),
+ "--latents-dir", latents_dir,
+ "--lr", _fmt(lr),
+ "--name", str(args.name),
+ "--save-dir", str(args.save_dir),
+ ]
+
+ _add(cmd, "--batch-size", getattr(args, "batch_size", None))
+ _add(cmd, "--seed", getattr(args, "seed", None))
+ _add(cmd, "--max-steps", getattr(args, "max_steps", None))
+ _add(cmd, "--checkpoint-every", getattr(args, "checkpoint_every", None))
+
+ _add(cmd, "--adapter-type", lora_cfg.get("adapter_type"))
+ _add(cmd, "--rank", lora_cfg.get("rank"))
+ _add(cmd, "--alpha", lora_cfg.get("alpha"))
+ _add(cmd, "--include", _csv(lora_cfg.get("include")))
+ _add(cmd, "--exclude", _csv(lora_cfg.get("exclude")))
+
+ _add(cmd, "--latent-crop-length", dataset_config.get("latent_crop_length"))
+ _add(cmd, "--timestep-sampler", training.get("timestep_sampler"))
+ _add(cmd, "--dist-shift", _dist_shift_from_config(model_config))
+ _add(cmd, "--cfg-dropout-prob", training.get("cfg_dropout_prob"))
+
+ # --- Resume: CLI arg wins over training_config.lora_ckpt_path (matches
+ # underfit.training.loop) ---
+ resume_path = getattr(args, "lora_ckpt_path", None) or training.get("lora_ckpt_path")
+ if resume_path and str(resume_path).strip():
+ resume_path = str(resume_path).strip()
+ _add(cmd, "--lora-ckpt-path", resume_path)
+ else:
+ resume_path = None
+ step_offset, epoch_offset = _resolve_offsets(training, resume_path)
+ _add(cmd, "--step-offset", step_offset)
+ _add(cmd, "--epoch-offset", epoch_offset)
+
+ # --- Demos: only when the run actually has a demo config ---
+ demo = training.get("demo") or {}
+ demo_every = demo.get("demo_every")
+ if demo.get("demo_cond") and demo_every:
+ _add(cmd, "--demo-config", _write_demo_config(demo))
+ _add(cmd, "--demo-every", demo_every)
+
+ return cmd
+
+
+def build_gradio_cmd(model, lora_ckpt_paths, share=False, port=None):
+ """Map onto the locked sa3_gradio.py CLI:
+ --dit --lora … [--share] [--port N].
+
+ `model` may be an Underfit base-model key ("sa3-medium") or a --dit value
+ ("medium"); it is normalized via map_model_name.
+ """
+ mlx_root, mlx_python = resolve_mlx_paths()
+ gradio_script = os.path.join(mlx_root, "scripts", "sa3_gradio.py")
+ if not os.path.isfile(gradio_script):
+ raise FileNotFoundError(
+ f"MLX gradio script not found at {gradio_script!r}. Point "
+ f"UNDERFIT_MLX_ROOT at the stable-audio-3/optimized/mlx checkout."
+ )
+ cmd = [mlx_python, gradio_script, "--dit", map_model_name(model)]
+ for p in _as_list(lora_ckpt_paths):
+ if p and str(p).strip():
+ cmd.extend(["--lora", str(p)])
+ if share:
+ cmd.append("--share")
+ if port:
+ cmd.extend(["--port", str(int(port))])
+ return cmd
+
+
+# --------------------------------------------------------------------------- #
+# Runners
+# --------------------------------------------------------------------------- #
+def run_mlx_training(args):
+ """Launch the MLX trainer, stream its stdout, and return its exit code.
+
+ cwd is the current process cwd (the dashboard already cd's into the run's
+ demo dir before invoking lora_train.py), so loss_by_timestep.bin and demo
+ files land where the dashboard expects. Every trainer step line is echoed
+ verbatim and ALSO re-emitted in the dashboard's tqdm-progress format so the
+ step/loss/lr parsers pick it up.
+ """
+ model_config = _load_json(args.model_config)
+ dit = _dit_model_from_config(model_config)
+ base_weights = resolve_base_weights(dit)
+ if not os.path.isfile(base_weights):
+ raise FileNotFoundError(
+ f"MLX base weights not found at {base_weights!r}. Set "
+ f"UNDERFIT_MLX_BASE_WEIGHTS or place the npz at "
+ f"/models/mlx/dit_{dit}-base_f16.npz."
+ )
+
+ cmd = build_trainer_cmd(args, base_weights)
+ print("[mlx-engine] launching MLX trainer:", flush=True)
+ print(" " + " ".join(shlex.quote(c) for c in cmd), flush=True)
+
+ env = dict(os.environ)
+ env.setdefault("PYTHONUNBUFFERED", "1")
+ proc = subprocess.Popen(
+ cmd,
+ cwd=os.getcwd(),
+ stdout=subprocess.PIPE,
+ stderr=subprocess.STDOUT,
+ text=True,
+ bufsize=1,
+ env=env,
+ )
+ try:
+ for line in proc.stdout:
+ line = line.rstrip("\n")
+ print(line, flush=True)
+ m = _STEP_LINE_RE.match(line)
+ if m:
+ n, loss, lr, ep = m.group(1), m.group(2), m.group(3), m.group(4)
+ # Format satisfies dashboard/server.py:
+ # _parse_latest_step (Epoch N: … a/b + "Step N," prefix)
+ # _HISTORY_RE / _LR_RE (train/loss= and train/lr=)
+ print(
+ f"Step {n}, Epoch {ep}: 100%|##########| 1/1 "
+ f"train/loss={loss} train/lr={lr}",
+ flush=True,
+ )
+ finally:
+ if proc.stdout:
+ proc.stdout.close()
+ code = proc.wait()
+ print(f"[mlx-engine] MLX trainer exited with code {code}", flush=True)
+ return code
+
+
+def run_mlx_gradio(model, lora_ckpt_paths, share=False, port=None):
+ """Launch the MLX gradio server (inherits stdout/stderr so the dashboard's
+ log redirect captures it) and return its exit code."""
+ cmd = build_gradio_cmd(model, lora_ckpt_paths, share=share, port=port)
+ print("[mlx-engine] launching MLX gradio:", flush=True)
+ print(" " + " ".join(shlex.quote(c) for c in cmd), flush=True)
+ env = dict(os.environ)
+ env.setdefault("PYTHONUNBUFFERED", "1")
+ proc = subprocess.Popen(cmd, cwd=os.getcwd(), env=env)
+ return proc.wait()
From 50b0122057971d06edc589eb55e07418ef7b3705 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 12:11:56 -0400
Subject: [PATCH 02/19] mlx-engine: transform underfit demo block to the MLX
trainer's --demo-config list
The MLX trainer's --demo-config is a flat LIST of {prompt,cfg,seed,steps,...}
entries (its demo_cond equivalent), not underfit's whole training.demo dict.
Map demo_cond -> entries (cfg from demo_cfg_scales[0], steps from demo_steps),
skipping ARC entries (the MLX trainer finetunes the base model; ARC demos need
a weight-swap it doesn't do).
---
underfit/backends/mlx_engine.py | 34 ++++++++++++++++++++++++++++-----
1 file changed, 29 insertions(+), 5 deletions(-)
diff --git a/underfit/backends/mlx_engine.py b/underfit/backends/mlx_engine.py
index 4dd4b34..78c0396 100644
--- a/underfit/backends/mlx_engine.py
+++ b/underfit/backends/mlx_engine.py
@@ -264,11 +264,33 @@ def _resolve_offsets(training, resume_path):
return step, epoch
-def _write_demo_config(demo):
- """Write the demo block to a temp json and return its path."""
+def _build_demo_entries(demo):
+ """Map underfit's `training.demo` block onto the MLX trainer's --demo-config,
+ which is a flat LIST of entries {prompt, cfg, seed, steps, lora_strength?,
+ lora_interval_max?}. ARC entries are skipped: the MLX trainer finetunes the
+ BASE model, and ARC demos need a weight-swap/second-model it doesn't do."""
+ cfg_scales = demo.get("demo_cfg_scales") or [7]
+ default_cfg = cfg_scales[0] if cfg_scales else 7
+ default_steps = demo.get("demo_steps", 50)
+ entries = []
+ for e in demo.get("demo_cond") or []:
+ if e.get("arc"):
+ continue
+ entry = {"prompt": e.get("prompt", ""),
+ "cfg": e.get("cfg", default_cfg),
+ "steps": e.get("steps", default_steps)}
+ for k in ("seed", "lora_strength", "lora_interval_max"):
+ if e.get(k) is not None:
+ entry[k] = e[k]
+ entries.append(entry)
+ return entries
+
+
+def _write_demo_config(entries):
+ """Write the MLX --demo-config entry list to a temp json and return its path."""
fd, path = tempfile.mkstemp(prefix="underfit_mlx_demo_", suffix=".json")
with os.fdopen(fd, "w") as f:
- json.dump(demo, f)
+ json.dump(entries, f)
return path
@@ -335,8 +357,10 @@ def build_trainer_cmd(args, base_weights):
demo = training.get("demo") or {}
demo_every = demo.get("demo_every")
if demo.get("demo_cond") and demo_every:
- _add(cmd, "--demo-config", _write_demo_config(demo))
- _add(cmd, "--demo-every", demo_every)
+ entries = _build_demo_entries(demo)
+ if entries:
+ _add(cmd, "--demo-config", _write_demo_config(entries))
+ _add(cmd, "--demo-every", demo_every)
return cmd
From 21eb9ba8d699e4e444d30f919e3bd98e1bee6fb3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 12:47:05 -0400
Subject: [PATCH 03/19] mlx-engine: map optimizer/scheduler/effective-length +
adapter_type default
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Map the SA3 template's training conventions onto the MLX trainer CLI: AdamW
betas/eps/weight_decay, InverseLR scheduler (+inv_gamma/power/warmup/final_lr),
model.diffusion.use_effective_length_for_schedule → --use-effective-length. Also
default adapter_type to 'lora' when the lora_config omits it (underfit's config-layer
default lora_config.get('adapter_type','lora'); the MLX trainer's own default is
dora-rows, which would otherwise mismatch a template without an explicit type).
---
underfit/backends/mlx_engine.py | 45 ++++++++++++++++++++++++++++++++-
1 file changed, 44 insertions(+), 1 deletion(-)
diff --git a/underfit/backends/mlx_engine.py b/underfit/backends/mlx_engine.py
index 78c0396..b6aa6af 100644
--- a/underfit/backends/mlx_engine.py
+++ b/underfit/backends/mlx_engine.py
@@ -238,6 +238,25 @@ def _dist_shift_from_config(model_config):
return opts.get("type")
+def _use_effective_length(model_config):
+ """model.diffusion.use_effective_length_for_schedule (True in SA3 templates):
+ shift timesteps by ceil(int(seconds_total*sr)/4096) not the crop length."""
+ diffusion = (model_config.get("model", {}) or {}).get("diffusion", {}) or {}
+ return bool(diffusion.get("use_effective_length_for_schedule", False))
+
+
+def _optimizer_config(training):
+ """training.optimizer_configs.diffusion.optimizer.config (AdamW betas/eps/wd)."""
+ opt = training.get("optimizer_configs") or {}
+ return ((opt.get("diffusion", {}) or {}).get("optimizer", {}) or {}).get("config", {}) or {}
+
+
+def _scheduler_config(training):
+ """training.optimizer_configs.diffusion.scheduler (type + config), or None."""
+ opt = training.get("optimizer_configs") or {}
+ return (opt.get("diffusion", {}) or {}).get("scheduler")
+
+
def _parse_filename_offsets(path):
"""(step, epoch) parsed from a checkpoint filename; either may be None.
Mirrors underfit.training.loop._parse_filename_offsets."""
@@ -330,7 +349,10 @@ def build_trainer_cmd(args, base_weights):
_add(cmd, "--max-steps", getattr(args, "max_steps", None))
_add(cmd, "--checkpoint-every", getattr(args, "checkpoint_every", None))
- _add(cmd, "--adapter-type", lora_cfg.get("adapter_type"))
+ # adapter_type: underfit's config-layer default is "lora" when absent
+ # (lora_config.get("adapter_type", "lora")); the MLX trainer defaults to
+ # dora-rows, so pass "lora" explicitly to match a template without one.
+ _add(cmd, "--adapter-type", lora_cfg.get("adapter_type", "lora"))
_add(cmd, "--rank", lora_cfg.get("rank"))
_add(cmd, "--alpha", lora_cfg.get("alpha"))
_add(cmd, "--include", _csv(lora_cfg.get("include")))
@@ -339,8 +361,29 @@ def build_trainer_cmd(args, base_weights):
_add(cmd, "--latent-crop-length", dataset_config.get("latent_crop_length"))
_add(cmd, "--timestep-sampler", training.get("timestep_sampler"))
_add(cmd, "--dist-shift", _dist_shift_from_config(model_config))
+ if _use_effective_length(model_config):
+ cmd.append("--use-effective-length")
_add(cmd, "--cfg-dropout-prob", training.get("cfg_dropout_prob"))
+ # optimizer: AdamW betas/eps/weight_decay (torch decoupled-wd == MLX AdamW)
+ opt_cfg = _optimizer_config(training)
+ betas = opt_cfg.get("betas")
+ if isinstance(betas, (list, tuple)) and len(betas) == 2:
+ _add(cmd, "--beta1", betas[0])
+ _add(cmd, "--beta2", betas[1])
+ _add(cmd, "--eps", opt_cfg.get("eps"))
+ _add(cmd, "--weight-decay", opt_cfg.get("weight_decay"))
+
+ # LR scheduler: only InverseLR is supported (the SA3 templates' scheduler)
+ sched = _scheduler_config(training)
+ if isinstance(sched, dict) and sched.get("type") == "InverseLR":
+ sc = sched.get("config", {}) or {}
+ cmd.extend(["--lr-scheduler", "inverse"])
+ _add(cmd, "--inv-gamma", sc.get("inv_gamma"))
+ _add(cmd, "--lr-power", sc.get("power"))
+ _add(cmd, "--lr-warmup", sc.get("warmup"))
+ _add(cmd, "--lr-final", sc.get("final_lr"))
+
# --- Resume: CLI arg wins over training_config.lora_ckpt_path (matches
# underfit.training.loop) ---
resume_path = getattr(args, "lora_ckpt_path", None) or training.get("lora_ckpt_path")
From e92f97328322d0bc307e36cfa2d77d7c2b955c41 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 14:37:09 -0400
Subject: [PATCH 04/19] mlx-engine: --dit-weights is optional (trainer
auto-downloads the base npz)
The MLX trainer now auto-downloads the base-model npz from HF, so underfit no longer
needs a local base file. Pass --dit-weights only when UNDERFIT_MLX_BASE_WEIGHTS is set
or a local base npz exists; otherwise omit it and let the trainer download. Removes the
hard FileNotFoundError on a missing local base file.
---
underfit/backends/mlx_engine.py | 30 +++++++++++++++++++-----------
1 file changed, 19 insertions(+), 11 deletions(-)
diff --git a/underfit/backends/mlx_engine.py b/underfit/backends/mlx_engine.py
index b6aa6af..d4a600a 100644
--- a/underfit/backends/mlx_engine.py
+++ b/underfit/backends/mlx_engine.py
@@ -316,11 +316,12 @@ def _write_demo_config(entries):
# --------------------------------------------------------------------------- #
# Command builders
# --------------------------------------------------------------------------- #
-def build_trainer_cmd(args, base_weights):
+def build_trainer_cmd(args, base_weights=None):
"""Map Underfit's args + run configs onto the locked lora_train_mlx.py CLI.
- Reads args.model_config / args.dataset_config (paths). base_weights is the
- resolved --dit-weights npz path (see resolve_base_weights).
+ Reads args.model_config / args.dataset_config (paths). base_weights is an
+ optional --dit-weights override; when None the MLX trainer auto-downloads
+ the base-model npz from HF (its default), so no local file is required.
"""
mlx_root, mlx_python = resolve_mlx_paths()
trainer = os.path.join(mlx_root, "scripts", "lora_train_mlx.py")
@@ -337,12 +338,13 @@ def build_trainer_cmd(args, base_weights):
cmd = [
mlx_python, trainer,
"--dit", dit,
- "--dit-weights", str(base_weights),
"--latents-dir", latents_dir,
"--lr", _fmt(lr),
"--name", str(args.name),
"--save-dir", str(args.save_dir),
]
+ if base_weights: # override only; else the trainer downloads the base npz
+ cmd[2:2] = ["--dit-weights", str(base_weights)]
_add(cmd, "--batch-size", getattr(args, "batch_size", None))
_add(cmd, "--seed", getattr(args, "seed", None))
@@ -447,13 +449,19 @@ def run_mlx_training(args):
"""
model_config = _load_json(args.model_config)
dit = _dit_model_from_config(model_config)
- base_weights = resolve_base_weights(dit)
- if not os.path.isfile(base_weights):
- raise FileNotFoundError(
- f"MLX base weights not found at {base_weights!r}. Set "
- f"UNDERFIT_MLX_BASE_WEIGHTS or place the npz at "
- f"/models/mlx/dit_{dit}-base_f16.npz."
- )
+ # Base weights: pass an explicit --dit-weights only for the override env, or
+ # if a local base npz happens to exist. Otherwise leave it off — the MLX
+ # trainer auto-downloads the base-model npz from HF (weights.TRAINING_BASE).
+ override = os.environ.get("UNDERFIT_MLX_BASE_WEIGHTS")
+ if override:
+ base_weights = os.path.abspath(override)
+ if not os.path.isfile(base_weights):
+ raise FileNotFoundError(
+ f"UNDERFIT_MLX_BASE_WEIGHTS={base_weights!r} does not exist."
+ )
+ else:
+ default = resolve_base_weights(dit)
+ base_weights = default if os.path.isfile(default) else None
cmd = build_trainer_cmd(args, base_weights)
print("[mlx-engine] launching MLX trainer:", flush=True)
From 68558cf7fffcabdf5afdae8f3a967a05275dfed7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 16:11:29 -0400
Subject: [PATCH 05/19] dashboard: platform-aware device UI (Apple Silicon
MPS/MLX, not just CUDA)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
The dashboard modeled "devices" as an nvidia-smi CUDA-GPU list and hard-gated on
it — on a Mac it showed "No CUDA GPUs detected ... require an NVIDIA GPU" and left
Start disabled, even though the training loop already resolves cuda>mps>cpu and the
MLX engine exists.
server.py:
- _detect_platform() (torch-free: Darwin+arm64 -> apple, else nvidia count -> cuda, else cpu).
- _apple_gpu_info() — the Metal counterpart of the nvidia-smi query: chip/GPU name +
core count via "system_profiler SPDisplaysDataType", unified-memory total via
"sysctl hw.memsize"; _apple_gpu_mem_used_mb() via "vm_stat".
- /api/gpu returns platform + engines + (on apple) a single synthetic Metal device
{kind:apple, name, cores, used/total/free_mb}. CUDA path unchanged.
- Launch env: _cuda_env_prefix() omits CUDA_VISIBLE_DEVICES on apple/cpu;
_free_gpu_memory no-ops off-CUDA.
index.html:
- Store platform/engines; _noGpuHtml() keeps the exact CUDA message on cuda, softer
"CPU (slow)" otherwise; Apple renders a normal device tile (name + N-core / unified
memory) — no warning.
- Auto-select the single Apple device (Start/Launch enable without a click).
- Engine dropdown filtered/relabelled to the platform: torch -> "MPS (PyTorch)" + "MLX"
on apple (default MLX); mlx hidden on cuda/cpu (default torch).
All Apple/CPU branches guarded on platform -> the Linux/CUDA path is unchanged.
---
dashboard/index.html | 74 ++++++++++++++++---
dashboard/server.py | 166 +++++++++++++++++++++++++++++++++++--------
2 files changed, 201 insertions(+), 39 deletions(-)
diff --git a/dashboard/index.html b/dashboard/index.html
index 76720e5..c0b1a81 100644
--- a/dashboard/index.html
+++ b/dashboard/index.html
@@ -1928,6 +1928,8 @@ Download Audio Selection
let _selectedGpu = null;
let _pendingCkptHighlight = null;
let _gpuData = [];
+let _platform = 'cuda'; // 'apple' | 'cuda' | 'cpu' — from /api/gpu
+let _engines = ['torch']; // available training engines — from /api/gpu
let _gradioVramEstimate = {load_mb: 10000, peak_mb: 12000, n_samples: 0};
let _arcInfo = {}; // {model_key: {arc_type, diffusion_objective}}
@@ -3391,7 +3393,7 @@ Download Audio Selection
gpuGradioCount[g.gpu] = (gpuGradioCount[g.gpu] || 0) + 1;
});
if (_gpuData.length === 0) {
- grid.innerHTML = '⚠ No CUDA GPUs detected — Gradio inference requires an NVIDIA GPU.
';
+ grid.innerHTML = _noGpuHtml('⚠ No CUDA GPUs detected — Gradio inference requires an NVIDIA GPU.', true);
return;
}
const _numGpus = _gpuData.length;
@@ -3418,6 +3420,8 @@ Download Audio Selection
if (!hidden) btn.onclick = () => selectGpu(i);
grid.appendChild(btn);
}
+ // Non-CUDA single accelerator (Apple): default-select it so Launch enables without a click.
+ if (_platform !== 'cuda' && _numGpus === 1) selectGpu(0);
document.getElementById('gpu-modal').classList.add('active');
};
@@ -3782,16 +3786,55 @@ Download Audio Selection
};
+// Platform-aware "no device" message. Keeps the original CUDA wording verbatim
+// (so the cuda path is byte-identical); softer, non-blocking note on cpu.
+function _noGpuHtml(cudaMsg, gridColumn) {
+ const gc = gridColumn ? ';grid-column:1/-1' : '';
+ const style = 'padding:8px 12px;color:var(--orange);font-size:0.85em;border:1px solid var(--orange);border-radius:6px;background:rgba(255,165,0,0.08)' + gc;
+ const msg = (_platform === 'cuda')
+ ? cudaMsg
+ : '⚠ No GPU accelerator — running on CPU (slow).';
+ return '' + msg + '
';
+}
+
+// Filter/relabel the engine dropdown to the platform's available engines.
+// Idempotent — safe to call on every /api/gpu tick. Never clobbers a valid
+// user selection (only falls back when the current option becomes hidden).
+function _applyEngineOptions() {
+ const sel = document.getElementById('new-ft-engine');
+ if (!sel) return;
+ const hasMlx = _engines.includes('mlx');
+ const isApple = _platform === 'apple';
+ Array.from(sel.options).forEach(o => {
+ if (o.value === 'mlx') {
+ o.hidden = !hasMlx;
+ o.disabled = !hasMlx;
+ o.textContent = isApple ? 'MLX' : 'mlx';
+ } else if (o.value === 'torch') {
+ o.textContent = isApple ? 'MPS (PyTorch)' : 'torch';
+ }
+ });
+ const cur = sel.options[sel.selectedIndex];
+ if (!cur || cur.hidden || cur.disabled) {
+ sel.value = hasMlx ? 'mlx' : 'torch';
+ }
+}
+
function _applyGpu(data) {
try {
const gpus = data.gpus || [];
+ _platform = data.platform || 'cuda';
+ _engines = data.engines || ['torch'];
+ window._platform = _platform;
+ window._engines = _engines;
+ _applyEngineOptions();
if (data.gradio_estimate) _gradioVramEstimate = data.gradio_estimate;
if (data.arc_info) _arcInfo = data.arc_info;
_cudaVisibleDevices = data.cuda_visible_devices || null;
_gpuData = gpus;
const row = document.getElementById('gpu-row');
if (gpus.length === 0) {
- row.innerHTML = '⚠ No CUDA GPUs detected. Training and encoding require an NVIDIA GPU with nvidia-smi available on PATH.
';
+ row.innerHTML = _noGpuHtml('⚠ No CUDA GPUs detected. Training and encoding require an NVIDIA GPU with nvidia-smi available on PATH.', false);
return;
}
row.innerHTML = gpus.map(g => {
@@ -3848,11 +3891,16 @@ Download Audio Selection
}).join('')
: (idle ? 'idle
' : '');
const activeClass = (_gpuVramPanelGpu === g.gpu) ? ' gpu-card-active' : '';
+ const isApple = g.kind === 'apple';
+ const idLabel = isApple ? escAttr(g.name || 'Apple GPU') : `GPU ${g.gpu}`;
+ const appleNote = isApple
+ ? `${g.cores ? g.cores + '-core · ' : ''}unified memory
`
+ : '';
return `
-
GPU ${g.gpu}
+
${idLabel}
${(g.used_mb/1024).toFixed(1)} / ${(g.total_mb/1024).toFixed(1)} GB
- ${labelsHtml}
+ ${isApple ? appleNote : labelsHtml}
Download Audio Selection
grid.innerHTML = '';
_buildCudaWarn(gridId);
if (_gpuData.length === 0) {
- grid.innerHTML = '⚠ No CUDA GPUs detected — training requires an NVIDIA GPU.
';
+ grid.innerHTML = _noGpuHtml('⚠ No CUDA GPUs detected — training requires an NVIDIA GPU.', true);
return;
}
const _numGpus = _gpuData.length;
@@ -5271,6 +5319,7 @@ Download Audio Selection
const usedMb = gpuInfo ? gpuInfo.used_mb : 0;
const totalMb = gpuInfo ? gpuInfo.total_mb : 0;
const labels = gpuInfo ? (gpuInfo.labels || []) : [];
+ const isApple = gpuInfo && gpuInfo.kind === 'apple';
const pct = totalMb > 0 ? Math.round(usedMb / totalMb * 100) : 0;
const barColor = pct > 80 ? 'var(--red)' : pct > 50 ? 'var(--orange)' : 'var(--green)';
const labelsHtml = labels.map(l => {
@@ -5278,15 +5327,21 @@ Download Audio Selection
}).join('');
const idle = labels.length === 0 && usedMb < 500;
const idleHtml = idle ? 'idle
' : '';
+ const appleNote = isApple
+ ? `${gpuInfo.cores ? gpuInfo.cores + '-core · ' : ''}unified memory
`
+ : '';
+ const idLabel = isApple ? escAttr(gpuInfo.name || 'Apple GPU') : `GPU ${i}`;
const hidden = !_isGpuVisible(i);
btn.className = 'gpu-btn' + (hidden ? ' cuda-hidden' : '');
- btn.innerHTML = `GPU ${i}
` +
+ btn.innerHTML = `${idLabel}
` +
`` +
`${(usedMb/1024).toFixed(1)} / ${(totalMb/1024).toFixed(1)} GB
` +
- (labelsHtml || idleHtml);
+ (isApple ? appleNote : (labelsHtml || idleHtml));
if (!hidden) btn.onclick = () => selectFn(i);
grid.appendChild(btn);
}
+ // Non-CUDA single accelerator (Apple): default-select it so Start enables without a click.
+ if (_platform !== 'cuda' && _numGpus === 1 && typeof selectFn === 'function') selectFn(0);
}
// LoRA per-layer data — driven by lora_layer_template in /api/models response.
@@ -6018,6 +6073,9 @@ Download Audio Selection
if (_bmOptions.has(_pref)) { _bmEl.value = _pref; break; }
}
document.getElementById('new-ft-base-precision').value = 'fp16';
+ // Engine: default mlx on Apple, torch elsewhere (options filtered/relabeled by _applyEngineOptions).
+ _applyEngineOptions();
+ document.getElementById('new-ft-engine').value = _engines.includes('mlx') ? 'mlx' : 'torch';
document.getElementById('new-ft-batch').value = '1';
document.getElementById('new-ft-steps').value = '40000';
document.getElementById('new-ft-lora-type').value = 'dora-rows';
@@ -7252,7 +7310,7 @@ Download Audio Selection
_buildCudaWarn('ds-gpu-grid');
_dsSelectedGpus = new Set();
if (_gpuData.length === 0) {
- grid.innerHTML = '⚠ No CUDA GPUs detected — dataset encoding requires an NVIDIA GPU.
';
+ grid.innerHTML = _noGpuHtml('⚠ No CUDA GPUs detected — dataset encoding requires an NVIDIA GPU.', true);
return;
}
const _numGpus = _gpuData.length;
diff --git a/dashboard/server.py b/dashboard/server.py
index fc50bd4..6736828 100644
--- a/dashboard/server.py
+++ b/dashboard/server.py
@@ -607,8 +607,9 @@ def _kill_process_group(pid, paused=False):
def _free_gpu_memory(gpu):
- """Run cuda.empty_cache() on the specified GPU to free VRAM."""
- if gpu is None:
+ """Run cuda.empty_cache() on the specified GPU to free VRAM. CUDA-only —
+ a no-op on Apple (MPS/MLX) / CPU, where there's no per-GPU CUDA cache."""
+ if gpu is None or _detect_platform() != "cuda":
return
try:
cmd = (
@@ -1526,7 +1527,7 @@ def launch(self, checkpoint_path, gpu, run_id=None, checkpoint_name=None, title=
# a non-IPython subprocess (matplotlib crashes on import).
cmd = (
f"source {_bash_quote(VENV_ACTIVATE)} && "
- f"{backend_env}CUDA_VISIBLE_DEVICES={gpu} GRADIO_SERVER_PORT={port} "
+ f"{backend_env}{_cuda_env_prefix(gpu)}GRADIO_SERVER_PORT={port} "
f"PYTHONUNBUFFERED=1 MPLBACKEND=Agg "
f"{thread_env}"
f"python {_bash_quote(RUN_GRADIO_SCRIPT)} "
@@ -2214,7 +2215,7 @@ def monitor_loop(self):
if "--gradient-clip-val" not in restart_cmd:
restart_cmd += " --gradient-clip-val 1.0"
# Restart training and append stdout/stderr to the log file.
- gpu_env = f"CUDA_VISIBLE_DEVICES={gpu} " if gpu is not None else ""
+ gpu_env = _cuda_env_prefix(gpu)
backend_env = _backend_env_for_model(fresh_run.get("base_model"))
demo_dir = fresh_run.get("demo_source_dir", str(RUNS_DIR))
os.makedirs(demo_dir, exist_ok=True)
@@ -2631,6 +2632,92 @@ def _update_gradio_vram_estimates():
_gradio_vram_baselines.pop(iid, None)
+_PLATFORM_CACHE = None
+
+
+def _detect_platform():
+ """Accelerator platform: 'apple' (Apple-Silicon / Metal — MPS + MLX),
+ 'cuda' (nvidia-smi GPUs present), or 'cpu'. Torch-free so it works even when
+ the training backend isn't importable in this venv. Cached for the session."""
+ global _PLATFORM_CACHE
+ if _PLATFORM_CACHE is None:
+ import platform as _pf
+ if _pf.system() == "Darwin" and _pf.machine() == "arm64":
+ _PLATFORM_CACHE = "apple"
+ elif _get_gpu_count() > 0:
+ _PLATFORM_CACHE = "cuda"
+ else:
+ _PLATFORM_CACHE = "cpu"
+ return _PLATFORM_CACHE
+
+
+_APPLE_INFO_CACHE = None
+
+
+def _apple_gpu_info():
+ """Static Apple-Silicon accelerator info — the Metal counterpart of the
+ nvidia-smi query. Uses `system_profiler SPDisplaysDataType` for the chip/GPU
+ name (+ core count when present) and `sysctl hw.memsize` for the unified
+ memory total. Cached (hardware is fixed for the session)."""
+ global _APPLE_INFO_CACHE
+ if _APPLE_INFO_CACHE is not None:
+ return _APPLE_INFO_CACHE
+ name, cores = "Apple GPU", None
+ try:
+ out = subprocess.check_output(
+ ["system_profiler", "-json", "SPDisplaysDataType"],
+ timeout=10, stderr=subprocess.DEVNULL,
+ ).decode()
+ disp = json.loads(out).get("SPDisplaysDataType", [])
+ if disp:
+ d0 = disp[0]
+ name = d0.get("sppci_model") or d0.get("_name") or name
+ raw_cores = d0.get("sppci_cores") or d0.get("spdisplays_ndrvs")
+ if raw_cores:
+ try:
+ cores = int(str(raw_cores).split()[0])
+ except (ValueError, IndexError):
+ pass
+ except Exception:
+ pass
+ mem_total_mb = 0
+ try:
+ mem_total_mb = int(subprocess.check_output(
+ ["sysctl", "-n", "hw.memsize"], timeout=5,
+ ).decode().strip()) // (1024 * 1024)
+ except Exception:
+ pass
+ _APPLE_INFO_CACHE = {"name": name, "cores": cores, "mem_total_mb": mem_total_mb}
+ return _APPLE_INFO_CACHE
+
+
+def _apple_gpu_mem_used_mb():
+ """Live unified-memory used (MB) via `vm_stat` — the closest analog to CUDA
+ memory.used on Apple's unified-memory GPU (active + wired + compressed)."""
+ try:
+ out = subprocess.check_output(["vm_stat"], timeout=5).decode()
+ m = re.search(r"page size of (\d+) bytes", out)
+ pgsize = int(m.group(1)) if m else 16384
+
+ def _pages(label):
+ mm = re.search(rf"{label}:\s+(\d+)\.", out)
+ return int(mm.group(1)) if mm else 0
+ used = (_pages("Pages active") + _pages("Pages wired down")
+ + _pages("Pages occupied by compressor"))
+ return used * pgsize // (1024 * 1024)
+ except Exception:
+ return 0
+
+
+def _cuda_env_prefix(gpu):
+ """`CUDA_VISIBLE_DEVICES= ` prefix for a launch command — empty on
+ non-CUDA platforms. Apple (MPS/MLX) uses its single device implicitly and
+ CPU has none, so pinning CUDA_VISIBLE_DEVICES there is meaningless."""
+ if gpu is None or _detect_platform() != "cuda":
+ return ""
+ return f"CUDA_VISIBLE_DEVICES={gpu} "
+
+
def _get_gpu_count(force_refresh=False):
"""Return number of CUDA GPUs from nvidia-smi (cached after the first
*successful* call). Returns 0 if nvidia-smi is missing or fails — the UI
@@ -4027,7 +4114,7 @@ def _handle_new_finetune(self, body):
f" --engine {engine}"
)
- gpu_env = f"CUDA_VISIBLE_DEVICES={gpu} "
+ gpu_env = _cuda_env_prefix(gpu)
# Cap CPU thread pools (same rationale as gradio launches: nproc-sized
# default pools can exhaust ulimit -u when multiple training procs are
# alive). Tied to GRADIO_THREAD_CAP for consistency.
@@ -4389,7 +4476,7 @@ def _handle_resume(self, run_id, body):
# Launch — use GPU from request body if provided, else fall back to run's previous GPU
gpu = body.get("gpu", run.get("gpu"))
- gpu_env = f"CUDA_VISIBLE_DEVICES={gpu} " if gpu is not None else ""
+ gpu_env = _cuda_env_prefix(gpu)
# Refuse the resume up front if the run's declared backends aren't
# installed (see _missing_backend_error_for_model docstring).
backend_err = _missing_backend_error_for_model(run.get("base_model"))
@@ -5017,31 +5104,46 @@ def _get_gpu_info(self):
"""Query nvidia-smi and annotate GPUs with training/gradio labels.
Generous timeout because the first nvidia-smi call on a fresh Colab
VM can take 5–10 s while the driver initializes."""
+ plat = _detect_platform()
gpus = []
- try:
- out = subprocess.check_output(
- ["nvidia-smi", "--query-gpu=index,memory.used,memory.total,memory.free,utilization.gpu",
- "--format=csv,noheader,nounits"],
- timeout=15, stderr=subprocess.DEVNULL,
- ).decode().strip()
- for line in out.split("\n"):
- parts = [p.strip() for p in line.split(",")]
- if len(parts) < 5:
- continue
- gpus.append({
- "gpu": int(parts[0]),
- "used_mb": int(parts[1]),
- "total_mb": int(parts[2]),
- "free_mb": int(parts[3]),
- "util_pct": int(parts[4]),
- "labels": [],
- })
- except Exception:
- return {"gpus": [], "gradio_estimate": _gradio_vram}
-
- caps = _query_gpu_compute_caps()
- for g in gpus:
- g["compute_cap"] = caps.get(g["gpu"])
+ if plat == "apple":
+ # Single unified-memory Metal device — usable via MPS (torch) or MLX.
+ info = _apple_gpu_info()
+ used = _apple_gpu_mem_used_mb()
+ total = info["mem_total_mb"]
+ gpus = [{
+ "gpu": 0, "kind": "apple",
+ "name": info["name"], "cores": info.get("cores"),
+ "used_mb": used, "total_mb": total,
+ "free_mb": max(0, total - used), "util_pct": 0,
+ "labels": [], "compute_cap": None,
+ }]
+ elif plat == "cuda":
+ try:
+ out = subprocess.check_output(
+ ["nvidia-smi", "--query-gpu=index,memory.used,memory.total,memory.free,utilization.gpu",
+ "--format=csv,noheader,nounits"],
+ timeout=15, stderr=subprocess.DEVNULL,
+ ).decode().strip()
+ for line in out.split("\n"):
+ parts = [p.strip() for p in line.split(",")]
+ if len(parts) < 5:
+ continue
+ gpus.append({
+ "gpu": int(parts[0]),
+ "used_mb": int(parts[1]),
+ "total_mb": int(parts[2]),
+ "free_mb": int(parts[3]),
+ "util_pct": int(parts[4]),
+ "labels": [],
+ })
+ except Exception:
+ return {"gpus": [], "platform": "cuda", "engines": ["torch"],
+ "gradio_estimate": _gradio_vram}
+ caps = _query_gpu_compute_caps()
+ for g in gpus:
+ g["compute_cap"] = caps.get(g["gpu"])
+ # plat == "cpu": no accelerator device; gpus stays [].
# Build lookup: gpu -> used_mb for checking occupancy
gpu_mem = {g["gpu"]: g["used_mb"] for g in gpus}
@@ -5131,7 +5233,9 @@ def _get_gpu_info(self):
else:
visible_gpus = None # not set = all visible
- resp = {"gpus": gpus, "gradio_estimate": _gradio_vram, "arc_info": arc_info}
+ resp = {"gpus": gpus, "gradio_estimate": _gradio_vram, "arc_info": arc_info,
+ "platform": plat,
+ "engines": ["torch", "mlx"] if plat == "apple" else ["torch"]}
if visible_gpus is not None:
resp["cuda_visible_devices"] = visible_gpus
return resp
From 90628e86d1508cb4903ae8f4eb6a18a03b473222 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 16:19:30 -0400
Subject: [PATCH 06/19] dashboard: engine-aware base-model gate + shorten MPS
label
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
The "base model not downloaded — run installer again" block in validateNewFt is
the TORCH base checkpoint (modelMeta.registered). Skip it when engine=mlx: the MLX
trainer brings its own base npz (auto-downloaded from HF / local in the MLX
checkout), so the torch download isn't required. Added an onchange on #new-ft-engine
so the error clears/appears when the engine is switched.
Also shorten the Apple engine label "MPS (PyTorch)" -> "MPS" to save UI width.
Dataset ENCODING still gates on the torch ckpt (validateNewDataset / launchEncoding)
— wiring the MLX pre-encode into the dashboard encoding flow is a separate follow-up.
---
dashboard/index.html | 10 +++++++---
1 file changed, 7 insertions(+), 3 deletions(-)
diff --git a/dashboard/index.html b/dashboard/index.html
index c0b1a81..9800ef3 100644
--- a/dashboard/index.html
+++ b/dashboard/index.html
@@ -1283,7 +1283,7 @@ New Finetune
Engine
-
+
torch
mlx
@@ -3811,7 +3811,7 @@ Download Audio Selection
o.disabled = !hasMlx;
o.textContent = isApple ? 'MLX' : 'mlx';
} else if (o.value === 'torch') {
- o.textContent = isApple ? 'MPS (PyTorch)' : 'torch';
+ o.textContent = isApple ? 'MPS' : 'torch';
}
});
const cur = sel.options[sel.selectedIndex];
@@ -6201,9 +6201,13 @@ Download Audio Selection
const btn = document.getElementById('new-ft-next-btn');
const baseModelKey = document.getElementById('new-ft-base-model').value;
const baseModelMeta = _remoteModels[baseModelKey] || {};
+ // The 'registered' check is the torch base checkpoint. The MLX engine brings
+ // its own base npz (auto-downloaded / local in the MLX checkout), so the torch
+ // download isn't required for an mlx run.
+ const _ftEngine = (document.getElementById('new-ft-engine') || {}).value || 'torch';
let error = '';
if (_newFtNoDatasets) error = 'Create a dataset first';
- else if (baseModelKey && baseModelMeta.registered === false)
+ else if (baseModelKey && baseModelMeta.registered === false && _ftEngine !== 'mlx')
error = `Base model "${baseModelMeta.label || baseModelKey}" not downloaded — please run installer again`;
else if (!name) error = '_name_empty';
else if (runsData.some(r => r.display_name === name)) error = `Run "${name}" already exists`;
From 4492c40a9e52a29c76718595d2c545beeb543112 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 16:24:06 -0400
Subject: [PATCH 07/19] dashboard: MLX dataset encoding on Apple (torch-free)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
The Encode Dataset flow launched dataset_processing/pre_encode.py (torch, reads the
torch base checkpoint) and gated on it being downloaded — so on Apple, where the
torch backend isn't installed, dataset encoding was blocked with "base model not
downloaded".
On Apple, encode with the MLX SAME encoder instead:
- mlx_engine.build_encode_cmd(input, output, model) -> pre_encode_mlx.py --audio-dir
--output-dir --codec (small -> same-s, medium -> same-l). Encoder weights are the
MLX SAME npz (auto-downloaded / local in the MLX checkout) — no torch checkpoint.
- server.py /api/datasets/encode: dispatch to the MLX encoder when platform == apple
(no CUDA_VISIBLE_DEVICES); torch path unchanged elsewhere. Output layout (npy + json
+ details.json) matches pre_encode.py, so the existing completion check works as-is.
- index.html: skip the "not downloaded" gate (validateNewDataset + launchEncoding)
when platform == apple.
Known gap: the MLX encoder does not yet honor the per-file exclude list (encodes the
full folder); logged when excludes are set.
---
dashboard/index.html | 6 ++--
dashboard/server.py | 49 ++++++++++++++++++++-------------
underfit/backends/mlx_engine.py | 17 ++++++++++++
3 files changed, 51 insertions(+), 21 deletions(-)
diff --git a/dashboard/index.html b/dashboard/index.html
index 9800ef3..7d36cea 100644
--- a/dashboard/index.html
+++ b/dashboard/index.html
@@ -7670,7 +7670,9 @@ Download Audio Selection
const modelMeta = _remoteModels[modelKey] || {};
const errEl = document.getElementById('ds-encode-error');
const btn = document.getElementById('ds-encode-btn');
- if (modelKey && modelMeta.registered === false) {
+ // On Apple the dashboard encodes with the MLX SAME encoder (own npz), so the
+ // torch base checkpoint ('registered') isn't required.
+ if (modelKey && modelMeta.registered === false && _platform !== 'apple') {
errEl.textContent = `Base model "${modelMeta.label || modelKey}" not downloaded — please run installer again`;
btn.disabled = true;
} else {
@@ -7818,7 +7820,7 @@ Download Audio Selection
const modelMeta = _remoteModels[model] || {};
const gpus = Array.from(_dsSelectedGpus).sort();
const errEl = document.getElementById('ds-encode-error');
- if (model && modelMeta.registered === false) {
+ if (model && modelMeta.registered === false && _platform !== 'apple') {
errEl.textContent = `Base model "${modelMeta.label || model}" not downloaded — please run installer again`;
return;
}
diff --git a/dashboard/server.py b/dashboard/server.py
index 6736828..84be81c 100644
--- a/dashboard/server.py
+++ b/dashboard/server.py
@@ -6285,25 +6285,36 @@ def _handle_datasets_encode(self, body):
exclude_file = output_dir / "exclude.txt"
exclude_file.write_text("\n".join(sorted(exclude_set)) + "\n")
- encode_args = [
- str(VENV_PYTHON),
- str(PRE_DIR / "pre_encode.py"),
- "--input-dir", str(input_path),
- "--model", model,
- "--output-dir", str(output_dir),
- "--num-gpus", str(len(gpus)),
- ]
- if half:
- encode_args.append("--half")
- if exclude_file is not None:
- encode_args.extend(["--exclude-file", str(exclude_file)])
-
- encode_env = os.environ.copy()
- encode_env.update({
- "CUDA_VISIBLE_DEVICES": gpu_str,
- "PYTHONUNBUFFERED": "1",
- "MPLBACKEND": "Agg",
- })
+ if _detect_platform() == "apple":
+ # MLX SAME encoder (torch-free) — no torch base checkpoint, no CUDA.
+ # Writes the same npy+json (+ details.json) layout pre_encode.py does.
+ from underfit.backends import mlx_engine
+ encode_args = mlx_engine.build_encode_cmd(input_path, output_dir, model)
+ if exclude_set:
+ print(f"[encode] mlx encoder does not yet honor the exclude list "
+ f"({len(exclude_set)} file(s)) — encoding the full folder.")
+ encode_env = os.environ.copy()
+ encode_env.update({"PYTHONUNBUFFERED": "1"})
+ else:
+ encode_args = [
+ str(VENV_PYTHON),
+ str(PRE_DIR / "pre_encode.py"),
+ "--input-dir", str(input_path),
+ "--model", model,
+ "--output-dir", str(output_dir),
+ "--num-gpus", str(len(gpus)),
+ ]
+ if half:
+ encode_args.append("--half")
+ if exclude_file is not None:
+ encode_args.extend(["--exclude-file", str(exclude_file)])
+
+ encode_env = os.environ.copy()
+ encode_env.update({
+ "CUDA_VISIBLE_DEVICES": gpu_str,
+ "PYTHONUNBUFFERED": "1",
+ "MPLBACKEND": "Agg",
+ })
log_handle = None
try:
diff --git a/underfit/backends/mlx_engine.py b/underfit/backends/mlx_engine.py
index d4a600a..16416f1 100644
--- a/underfit/backends/mlx_engine.py
+++ b/underfit/backends/mlx_engine.py
@@ -435,6 +435,23 @@ def build_gradio_cmd(model, lora_ckpt_paths, share=False, port=None):
return cmd
+def build_encode_cmd(input_dir, output_dir, base_model):
+ """Map the dashboard's encode request onto pre_encode_mlx.py (the torch-free
+ MLX SAME encoder). base_model → codec: small (sm-music/sm-sfx) → same-s,
+ medium → same-l. The encoder weights (SAME npz) are auto-downloaded / local
+ in the MLX checkout — no torch base checkpoint needed."""
+ mlx_root, mlx_python = resolve_mlx_paths()
+ script = os.path.join(mlx_root, "scripts", "pre_encode_mlx.py")
+ if not os.path.isfile(script):
+ raise FileNotFoundError(
+ f"MLX pre-encode script not found at {script!r}. Point "
+ f"UNDERFIT_MLX_ROOT at the stable-audio-3/optimized/mlx checkout."
+ )
+ codec = "same-l" if map_model_name(base_model) == "medium" else "same-s"
+ return [mlx_python, script, "--audio-dir", str(input_dir),
+ "--output-dir", str(output_dir), "--codec", codec]
+
+
# --------------------------------------------------------------------------- #
# Runners
# --------------------------------------------------------------------------- #
From 25a37fd78850b8c998052792e8ed5fb760baaf35 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 16:37:02 -0400
Subject: [PATCH 08/19] dashboard: MLX encode writes into latent_dir, not
output_dir
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
pre_encode_mlx.py writes flat into --output-dir, but the dashboard expects latents
in output_dir/latents/ (latent_dir) — where the torch pre_encode.py writes and
where the encoding monitor looks for details.json + *.npy. Pointing the MLX encoder at
output_dir left latent_dir empty, so the monitor saw a finished-but-empty encode and
dropped the dataset. Point build_encode_cmd at latent_dir (+ mkdir it).
---
dashboard/server.py | 8 ++++++--
1 file changed, 6 insertions(+), 2 deletions(-)
diff --git a/dashboard/server.py b/dashboard/server.py
index 84be81c..d1bb0ee 100644
--- a/dashboard/server.py
+++ b/dashboard/server.py
@@ -6287,9 +6287,13 @@ def _handle_datasets_encode(self, body):
if _detect_platform() == "apple":
# MLX SAME encoder (torch-free) — no torch base checkpoint, no CUDA.
- # Writes the same npy+json (+ details.json) layout pre_encode.py does.
+ # Write into latent_dir (output_dir/latents/) — the exact dir the
+ # torch pre_encode.py targets and the encoding monitor checks for
+ # details.json + *.npy. (pre_encode_mlx writes flat into --output-dir,
+ # so point it AT latent_dir rather than output_dir.)
from underfit.backends import mlx_engine
- encode_args = mlx_engine.build_encode_cmd(input_path, output_dir, model)
+ latent_dir.mkdir(parents=True, exist_ok=True)
+ encode_args = mlx_engine.build_encode_cmd(input_path, latent_dir, model)
if exclude_set:
print(f"[encode] mlx encoder does not yet honor the exclude list "
f"({len(exclude_set)} file(s)) — encoding the full folder.")
From 08f90f438fa25bf04735a18e50bae39f98e55824 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 16:44:08 -0400
Subject: [PATCH 09/19] dashboard: pass exclude-file to the MLX encoder
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
build_encode_cmd takes exclude_file and appends --exclude-file; the encode dispatch
passes the run's exclude.txt. Removes the earlier 'mlx encoder does not honor the
exclude list' note — per-file exclusions now work on the MLX path.
---
dashboard/server.py | 6 ++----
underfit/backends/mlx_engine.py | 12 ++++++++----
2 files changed, 10 insertions(+), 8 deletions(-)
diff --git a/dashboard/server.py b/dashboard/server.py
index d1bb0ee..d9219f7 100644
--- a/dashboard/server.py
+++ b/dashboard/server.py
@@ -6293,10 +6293,8 @@ def _handle_datasets_encode(self, body):
# so point it AT latent_dir rather than output_dir.)
from underfit.backends import mlx_engine
latent_dir.mkdir(parents=True, exist_ok=True)
- encode_args = mlx_engine.build_encode_cmd(input_path, latent_dir, model)
- if exclude_set:
- print(f"[encode] mlx encoder does not yet honor the exclude list "
- f"({len(exclude_set)} file(s)) — encoding the full folder.")
+ encode_args = mlx_engine.build_encode_cmd(
+ input_path, latent_dir, model, exclude_file=exclude_file)
encode_env = os.environ.copy()
encode_env.update({"PYTHONUNBUFFERED": "1"})
else:
diff --git a/underfit/backends/mlx_engine.py b/underfit/backends/mlx_engine.py
index 16416f1..7810b99 100644
--- a/underfit/backends/mlx_engine.py
+++ b/underfit/backends/mlx_engine.py
@@ -435,11 +435,12 @@ def build_gradio_cmd(model, lora_ckpt_paths, share=False, port=None):
return cmd
-def build_encode_cmd(input_dir, output_dir, base_model):
+def build_encode_cmd(input_dir, output_dir, base_model, exclude_file=None):
"""Map the dashboard's encode request onto pre_encode_mlx.py (the torch-free
MLX SAME encoder). base_model → codec: small (sm-music/sm-sfx) → same-s,
medium → same-l. The encoder weights (SAME npz) are auto-downloaded / local
- in the MLX checkout — no torch base checkpoint needed."""
+ in the MLX checkout — no torch base checkpoint needed. exclude_file is a text
+ file of relpaths to skip (same format the torch pre_encode.py uses)."""
mlx_root, mlx_python = resolve_mlx_paths()
script = os.path.join(mlx_root, "scripts", "pre_encode_mlx.py")
if not os.path.isfile(script):
@@ -448,8 +449,11 @@ def build_encode_cmd(input_dir, output_dir, base_model):
f"UNDERFIT_MLX_ROOT at the stable-audio-3/optimized/mlx checkout."
)
codec = "same-l" if map_model_name(base_model) == "medium" else "same-s"
- return [mlx_python, script, "--audio-dir", str(input_dir),
- "--output-dir", str(output_dir), "--codec", codec]
+ cmd = [mlx_python, script, "--audio-dir", str(input_dir),
+ "--output-dir", str(output_dir), "--codec", codec]
+ if exclude_file:
+ cmd.extend(["--exclude-file", str(exclude_file)])
+ return cmd
# --------------------------------------------------------------------------- #
From c1669ee8897439fc36a7f2a12eb34e3d3484cb2a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 17:17:40 -0400
Subject: [PATCH 10/19] mlx-engine: pass ARC demos, forward grad-clip, inherit
trainer stdout
run_mlx_training now inherits the trainer's stdout/stderr instead of
capturing and re-emitting a synthetic step line. The trainer emits the
same tqdm output as underfit's torch loop, so writing it straight to
the run log preserves tqdm's carriage-return progress updates (the
dashboard collapses on them) and the grad_norm / lora_magnitude postfix
(the metric charts parse it). Drops the translation loop and the
now-unused _STEP_LINE_RE.
_build_demo_entries passes ARC entries through (arc=true, cfg=1, steps=8
defaults) instead of skipping them - the MLX trainer renders them now.
build_trainer_cmd forwards --gradient-clip-val (the dashboard passes
1.0) so the MLX trainer clips identically.
---
underfit/backends/mlx_engine.py | 81 +++++++++++++++------------------
1 file changed, 37 insertions(+), 44 deletions(-)
diff --git a/underfit/backends/mlx_engine.py b/underfit/backends/mlx_engine.py
index 7810b99..30560fd 100644
--- a/underfit/backends/mlx_engine.py
+++ b/underfit/backends/mlx_engine.py
@@ -35,12 +35,6 @@
_SUPPORTED_DITS = ("sm-music", "sm-sfx", "medium")
-# Underfit's tqdm-postfix step line the MLX trainer emits, per its locked
-# contract: "step {N} train/loss {L:.6f} train/lr {lr:.3e} epoch {E} (…)"
-_STEP_LINE_RE = re.compile(
- r"^step (\d+)\s+train/loss (\S+)\s+train/lr (\S+)\s+epoch (\d+)"
-)
-
_STEP_TOKEN_RE = re.compile(r"step=(\d+)")
_EPOCH_TOKEN_RE = re.compile(r"epoch=(\d+)")
@@ -285,19 +279,29 @@ def _resolve_offsets(training, resume_path):
def _build_demo_entries(demo):
"""Map underfit's `training.demo` block onto the MLX trainer's --demo-config,
- which is a flat LIST of entries {prompt, cfg, seed, steps, lora_strength?,
- lora_interval_max?}. ARC entries are skipped: the MLX trainer finetunes the
- BASE model, and ARC demos need a weight-swap/second-model it doesn't do."""
+ which is a flat LIST of entries {prompt, cfg, seed, steps, arc?,
+ lora_strength?, lora_interval_max?}.
+
+ Both RF/base and ARC entries are passed through. ARC entries carry
+ ``arc: true`` and default to cfg=1 / steps=8 (the distilled rf_denoiser
+ convention, matching underfit's demo_step._run_one_arc_entry); the MLX
+ trainer renders them by loading the shipped ARC weights with the trained
+ LoRA merged in and sampling with the pingpong integrator. RF entries default
+ to the run's demo_cfg_scales[0] / demo_steps."""
cfg_scales = demo.get("demo_cfg_scales") or [7]
default_cfg = cfg_scales[0] if cfg_scales else 7
default_steps = demo.get("demo_steps", 50)
entries = []
for e in demo.get("demo_cond") or []:
if e.get("arc"):
- continue
- entry = {"prompt": e.get("prompt", ""),
- "cfg": e.get("cfg", default_cfg),
- "steps": e.get("steps", default_steps)}
+ entry = {"prompt": e.get("prompt", ""),
+ "cfg": e.get("cfg", 1), # ARC: distilled, cfg=1
+ "steps": e.get("steps", 8), # ARC: 8-step pingpong
+ "arc": True}
+ else:
+ entry = {"prompt": e.get("prompt", ""),
+ "cfg": e.get("cfg", default_cfg),
+ "steps": e.get("steps", default_steps)}
for k in ("seed", "lora_strength", "lora_interval_max"):
if e.get(k) is not None:
entry[k] = e[k]
@@ -376,6 +380,13 @@ def build_trainer_cmd(args, base_weights=None):
_add(cmd, "--eps", opt_cfg.get("eps"))
_add(cmd, "--weight-decay", opt_cfg.get("weight_decay"))
+ # gradient clipping: the dashboard passes --gradient-clip-val (1.0 by
+ # default, server.py restart_cmd) to lora_train.py; forward it so the MLX
+ # trainer clips identically and the grad_norm metric is post-clip like torch.
+ gcv = getattr(args, "gradient_clip_val", None)
+ if gcv:
+ _add(cmd, "--gradient-clip-val", gcv)
+
# LR scheduler: only InverseLR is supported (the SA3 templates' scheduler)
sched = _scheduler_config(training)
if isinstance(sched, dict) and sched.get("type") == "InverseLR":
@@ -460,13 +471,18 @@ def build_encode_cmd(input_dir, output_dir, base_model, exclude_file=None):
# Runners
# --------------------------------------------------------------------------- #
def run_mlx_training(args):
- """Launch the MLX trainer, stream its stdout, and return its exit code.
+ """Launch the MLX trainer, inheriting its stdout/stderr, and return its exit
+ code.
cwd is the current process cwd (the dashboard already cd's into the run's
demo dir before invoking lora_train.py), so loss_by_timestep.bin and demo
- files land where the dashboard expects. Every trainer step line is echoed
- verbatim and ALSO re-emitted in the dashboard's tqdm-progress format so the
- step/loss/lr parsers pick it up.
+ files land where the dashboard expects. The trainer emits the SAME tqdm
+ output as underfit's torch loop (desc "Step N, Epoch E"; postfix
+ train/loss, train/lr, train/grad_norm, train/lora_magnitude), so we let it
+ write straight through to lora_train.py's stdout (which the dashboard has
+ redirected to the run log). No line interception/translation — that would
+ break tqdm's \\r in-place updates (the dashboard collapses on them) and drop
+ the grad_norm / lora_magnitude postfix the charts parse.
"""
model_config = _load_json(args.model_config)
dit = _dit_model_from_config(model_config)
@@ -490,33 +506,10 @@ def run_mlx_training(args):
env = dict(os.environ)
env.setdefault("PYTHONUNBUFFERED", "1")
- proc = subprocess.Popen(
- cmd,
- cwd=os.getcwd(),
- stdout=subprocess.PIPE,
- stderr=subprocess.STDOUT,
- text=True,
- bufsize=1,
- env=env,
- )
- try:
- for line in proc.stdout:
- line = line.rstrip("\n")
- print(line, flush=True)
- m = _STEP_LINE_RE.match(line)
- if m:
- n, loss, lr, ep = m.group(1), m.group(2), m.group(3), m.group(4)
- # Format satisfies dashboard/server.py:
- # _parse_latest_step (Epoch N: … a/b + "Step N," prefix)
- # _HISTORY_RE / _LR_RE (train/loss= and train/lr=)
- print(
- f"Step {n}, Epoch {ep}: 100%|##########| 1/1 "
- f"train/loss={loss} train/lr={lr}",
- flush=True,
- )
- finally:
- if proc.stdout:
- proc.stdout.close()
+ # Inherit stdout/stderr: the trainer's tqdm writes straight to the run log
+ # (via lora_train.py's redirected fd), preserving \r so the dashboard
+ # collapses the progress bar and parses every step's postfix metrics.
+ proc = subprocess.Popen(cmd, cwd=os.getcwd(), env=env)
code = proc.wait()
print(f"[mlx-engine] MLX trainer exited with code {code}", flush=True)
return code
From b7f6dfdf4fae21f12cdbadfba20285fe8ccf7a0b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 17:33:59 -0400
Subject: [PATCH 11/19] dashboard: fix Apple/MLX gaps in gradio launch, device
panel, resume, model gate
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
(1) The /api/gradio gpu validation now only applies on CUDA. Apple/CPU have one
implicit accelerator (index 0) and nvidia-smi reports 0 GPUs, so the old
gpu >= _get_gpu_count() check rejected every Apple launch with "gpu must be
0--1". The MLX/MPS launch ignores the index anyway.
(2) The device-detail panel had no data on Apple: _query_gpu_mem returned {}
(no nvidia-smi) so the VRAM history never recorded ("Collecting VRAM data..."),
and _get_gpu_processes used nvidia-smi + /proc (empty on macOS → "No
processes"). Now _query_gpu_mem reports unified memory as GPU 0 (vm_stat), and
_apple_gpu_processes lists the managed training / gradio / encoding runs with
their whole process-tree RSS (ps/pgrep, no /proc).
(3) The "Base model not downloaded" popup (_warnUnregisteredModel) fired on
Apple for both MLX and MPS even though the MLX npz is present. Guarded with
_platform==='apple', matching the existing encode and validateNewFt guards.
(4) Resume now skips the torch-backend gate for engine=mlx, mirroring the
fresh-launch path — killing and reviving an MLX run no longer errors "requires
the sa3 backend".
---
dashboard/index.html | 5 ++
dashboard/server.py | 111 +++++++++++++++++++++++++++++++++++++++----
2 files changed, 108 insertions(+), 8 deletions(-)
diff --git a/dashboard/index.html b/dashboard/index.html
index 7d36cea..218b8eb 100644
--- a/dashboard/index.html
+++ b/dashboard/index.html
@@ -5567,6 +5567,11 @@ Download Audio Selection
// disabled-Next/disabled-Encode logic still runs alongside this.
window._warnUnregisteredModel = function(modelKey) {
if (!modelKey) return;
+ // On Apple the MLX engine (New Finetune) and the MLX SAME encoder (Encode
+ // Dataset) bring their own base npz — auto-downloaded / local in the MLX
+ // checkout — so the torch 'registered' flag doesn't gate usage here. Matches
+ // the validateNewFt (engine!=mlx) and encode (_platform!=apple) guards.
+ if (_platform === 'apple') return;
const meta = _remoteModels[modelKey] || {};
if (meta.registered !== false) return;
const label = _escAttr(meta.label || modelKey);
diff --git a/dashboard/server.py b/dashboard/server.py
index d9219f7..916e6e9 100644
--- a/dashboard/server.py
+++ b/dashboard/server.py
@@ -2737,7 +2737,11 @@ def _get_gpu_count(force_refresh=False):
def _query_gpu_mem():
- """Return dict of gpu_index -> used_mb from nvidia-smi."""
+ """Return dict of gpu_index -> used_mb. CUDA: nvidia-smi per GPU. Apple: the
+ single unified-memory accelerator (index 0) via vm_stat, so the VRAM history
+ sampler and per-run VRAM chart populate on Apple Silicon too."""
+ if _detect_platform() == "apple":
+ return {0: _apple_gpu_mem_used_mb()}
try:
out = subprocess.check_output(
["nvidia-smi", "--query-gpu=index,memory.used",
@@ -3574,7 +3578,11 @@ def do_POST(self):
self._json_response({"error": "missing checkpoint_path or gpu"}, status=400)
return
gpu = int(gpu)
- if gpu < 0 or gpu >= _get_gpu_count():
+ # On CUDA, validate against the nvidia-smi GPU count. On Apple/CPU
+ # there's a single implicit accelerator (index 0) and nvidia-smi
+ # reports 0 GPUs, so skip the check — the MLX/MPS launch ignores the
+ # index anyway (CUDA_VISIBLE_DEVICES is only set on CUDA).
+ if _detect_platform() == "cuda" and (gpu < 0 or gpu >= _get_gpu_count()):
self._json_response({"error": f"gpu must be 0-{_get_gpu_count()-1}"}, status=400)
return
instance_id, err = gradio_manager.launch(
@@ -4478,11 +4486,14 @@ def _handle_resume(self, run_id, body):
gpu = body.get("gpu", run.get("gpu"))
gpu_env = _cuda_env_prefix(gpu)
# Refuse the resume up front if the run's declared backends aren't
- # installed (see _missing_backend_error_for_model docstring).
- backend_err = _missing_backend_error_for_model(run.get("base_model"))
- if backend_err:
- self._json_response({"error": backend_err}, status=400)
- return
+ # installed (see _missing_backend_error_for_model docstring). engine=mlx
+ # trains in a separate MLX venv/checkout, so the in-venv torch backend
+ # need not be installed — skip the gate (mirrors the fresh-launch path).
+ if run.get("engine", "torch") != "mlx":
+ backend_err = _missing_backend_error_for_model(run.get("base_model"))
+ if backend_err:
+ self._json_response({"error": backend_err}, status=400)
+ return
backend_env = _backend_env_for_model(run.get("base_model"))
demo_dir = run.get("demo_source_dir", str(RUNS_DIR))
os.makedirs(demo_dir, exist_ok=True)
@@ -5240,8 +5251,92 @@ def _get_gpu_info(self):
resp["cuda_visible_devices"] = visible_gpus
return resp
+ def _descendant_pids(self, pid):
+ """Set of all descendant PIDs of pid (recursive pgrep -P)."""
+ found = set()
+ frontier = [int(pid)]
+ while frontier:
+ p = frontier.pop()
+ try:
+ kids = subprocess.check_output(
+ ["pgrep", "-P", str(p)], stderr=subprocess.DEVNULL
+ ).decode().split()
+ except Exception:
+ kids = []
+ for k in kids:
+ ki = int(k)
+ if ki not in found:
+ found.add(ki)
+ frontier.append(ki)
+ return found
+
+ def _apple_gpu_processes(self):
+ """Apple unified-memory accelerator processes: the dashboard-managed
+ training / gradio / encoding runs that are alive, each reporting the RSS
+ of its whole process tree as used_mb (Apple has no per-process VRAM —
+ RSS on unified memory is the closest analog). Uses ps/pgrep, not
+ nvidia-smi or /proc, so it works on macOS."""
+ def _alive(pid):
+ try:
+ os.kill(int(pid), 0)
+ return True
+ except Exception:
+ return False
+
+ def _tree_rss_mb(pid):
+ try:
+ pids = self._descendant_pids(pid) | {int(pid)}
+ total_kb = 0
+ for p in pids:
+ out = subprocess.check_output(
+ ["ps", "-o", "rss=", "-p", str(p)],
+ stderr=subprocess.DEVNULL,
+ ).decode().strip()
+ if out:
+ total_kb += int(out.split()[0])
+ return total_kb // 1024
+ except Exception:
+ return 0
+
+ procs = []
+ for r in registry.list_runs():
+ if r.get("status") in ("completed", "killed"):
+ continue
+ pid = r.get("pid")
+ if pid and _alive(pid):
+ procs.append({
+ "pid": pid, "used_mb": _tree_rss_mb(pid), "type": "training",
+ "name": r.get("display_name", r["id"]), "run_id": r["id"],
+ })
+ for inst in gradio_manager.list_instances():
+ if inst["status"] not in ("starting", "ready"):
+ continue
+ pid = inst.get("pid")
+ if pid and _alive(pid):
+ procs.append({
+ "pid": pid, "used_mb": _tree_rss_mb(pid), "type": "gradio",
+ "name": inst.get("title") or inst.get("checkpoint_name", "?"),
+ "run_id": inst.get("run_id"),
+ "checkpoint_path": inst.get("checkpoint_path"),
+ "instance_id": inst["id"],
+ })
+ for ds in datasets_registry.list_datasets():
+ if ds["status"] != "encoding":
+ continue
+ pid = ds.get("encoding_pid")
+ if pid and _alive(pid):
+ procs.append({
+ "pid": pid, "used_mb": _tree_rss_mb(pid), "type": "encoding",
+ "name": ds["name"], "dataset_id": ds["id"],
+ })
+ procs.sort(key=lambda p: p["used_mb"], reverse=True)
+ return procs
+
def _get_gpu_processes(self, gpu_idx):
- """Return classified processes running on a specific GPU."""
+ """Return classified processes running on a specific GPU (CUDA), or the
+ dashboard-managed accelerator processes on Apple unified memory."""
+ if _detect_platform() == "apple":
+ return self._apple_gpu_processes()
# Query nvidia-smi for compute processes
try:
out = subprocess.check_output(
From f8e606fb7e096b0807202c7b5e12d4bd5748a4ef Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 18:01:45 -0400
Subject: [PATCH 12/19] dashboard/installer: on-demand + install-time MLX model
download
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
When a base model's MLX weight pack (base + ARC + SAME codec + encoder + shared
t5gemma) isn't present, offer to download it instead of a dead-end "run
installer again" message.
- mlx_engine: shared pack helpers — mlx_pack_rel_paths / mlx_model_available /
mlx_missing_pack (files + approx GB) / build_mlx_download_cmd (runs the MLX
checkout's weights.ensure_local, the download source of truth, emitting
DL/OK/ALL_DONE lines for progress). mlx_root_or_none resolves without raising.
- server: /api/models annotates each model with mlx_available + mlx_download_gb
on Apple; POST /api/models/download starts a tracked background download and
GET /api/models/download?model= reports {state, done/total files, current
file, remaining GB}.
- index.html: New Finetune shows a "Download base + ARC (~X GB)" button when the
selected model's MLX pack is missing (engine=mlx on Apple). It downloads with
a live progress line, refreshes availability, and clears the gate — Next stays
disabled until the pack lands.
- setup.py: on Apple, the model phase pre-downloads the MLX packs into the
existing MLX checkout (reusing the same mlx_engine download) and skips the
torch packs an MLX-only Mac won't use. Degrades gracefully if the checkout /
venv isn't set up (weights then download lazily on first use).
---
dashboard/index.html | 83 +++++++++++++++++++++-
dashboard/server.py | 121 +++++++++++++++++++++++++++++++-
underfit/backends/mlx_engine.py | 82 ++++++++++++++++++++++
underfit/cli/setup.py | 61 ++++++++++++++++
4 files changed, 345 insertions(+), 2 deletions(-)
diff --git a/dashboard/index.html b/dashboard/index.html
index 218b8eb..68b02ac 100644
--- a/dashboard/index.html
+++ b/dashboard/index.html
@@ -1280,6 +1280,8 @@ New Finetune
+
+
Engine
@@ -5581,6 +5583,77 @@
Download Audio Selection
_showSeedDialog('Base model not downloaded', bodyHtml, [{label: 'OK', primary: true}]);
};
+// ── On-demand MLX model-pack download (item 5) ──────────────────────
+// When a base model's MLX weights (base + ARC + codec + encoder + t5) aren't
+// present, offer to download them into the MLX checkout instead of a dead-end
+// "run installer" message. Renders into `containerId`; `onDone()` re-runs once
+// the pack lands so the gate clears. Apple + engine=mlx only.
+const _mlxDlState = {}; // modelKey -> 'idle'|'downloading'|'done'|'error'
+const _mlxDlCb = {}; // modelKey -> onDone callback
+
+window._renderMlxDownload = function(containerId, modelKey, meta, show, onDone) {
+ const el = document.getElementById(containerId);
+ if (!el) return;
+ if (!show) { el.style.display = 'none'; el.innerHTML = ''; return; }
+ el.style.display = 'block';
+ _mlxDlCb[modelKey] = onDone;
+ if (_mlxDlState[modelKey] === 'downloading') return; // poller owns the panel
+ const label = _escAttr(meta.label || modelKey);
+ const gb = meta.mlx_download_gb ? ` (~${meta.mlx_download_gb} GB)` : '';
+ el.innerHTML =
+ `
`
+ + `⚠ ${label} not downloaded for MLX${gb} `
+ + `Download base + ARC `
+ + `
`;
+};
+
+window._startModelDownload = async function(modelKey, containerId) {
+ _mlxDlState[modelKey] = 'downloading';
+ const el = document.getElementById(containerId);
+ if (el) el.innerHTML = `
Starting download…
`;
+ try {
+ const res = await fetch('/api/models/download', {
+ method: 'POST', headers: {'Content-Type': 'application/json'},
+ body: JSON.stringify({model: modelKey}),
+ });
+ const data = await res.json();
+ if (!res.ok) throw new Error(data.error || 'failed to start');
+ } catch (e) {
+ _mlxDlState[modelKey] = 'error';
+ if (el) el.innerHTML = `
Download failed: ${_escAttr(e.message)}
`;
+ return;
+ }
+ _pollModelDownload(modelKey, containerId);
+};
+
+window._pollModelDownload = async function(modelKey, containerId) {
+ const el = document.getElementById(containerId);
+ let data;
+ try {
+ const res = await fetch('/api/models/download?model=' + encodeURIComponent(modelKey));
+ data = await res.json();
+ } catch (e) { setTimeout(() => _pollModelDownload(modelKey, containerId), 2500); return; }
+ if (data.state === 'downloading') {
+ const cur = data.current_file ? ` — ${_escAttr(data.current_file)}` : '';
+ const rem = data.remaining_gb ? `, ${data.remaining_gb} GB left` : '';
+ if (el) el.innerHTML = `
`
+ + `Downloading ${data.done_files || 0}/${data.total_files || '?'} files${rem}${cur}…
`;
+ setTimeout(() => _pollModelDownload(modelKey, containerId), 2500);
+ return;
+ }
+ if (data.state === 'done' || data.available) {
+ _mlxDlState[modelKey] = 'done';
+ if (el) { el.style.display = 'none'; el.innerHTML = ''; }
+ await _loadRemoteModels(); // refresh mlx_available
+ const cb = _mlxDlCb[modelKey];
+ if (typeof cb === 'function') cb(); // re-validate → gate clears
+ return;
+ }
+ _mlxDlState[modelKey] = 'error';
+ if (el) el.innerHTML = `
Download failed — see server log
`;
+};
+
// Open the file picker for the seed-LoRA upload. Always clear the input's
// value first — otherwise re-picking the same file (after Cancel) doesn't
// re-fire the change event and the picker appears to do nothing.
@@ -6210,10 +6283,17 @@
Download Audio Selection
// its own base npz (auto-downloaded / local in the MLX checkout), so the torch
// download isn't required for an mlx run.
const _ftEngine = (document.getElementById('new-ft-engine') || {}).value || 'torch';
+ // MLX engine: gate on the MLX pack (base+ARC+codec+encoder+t5), not the torch
+ // 'registered' flag. If missing, the download panel offers an on-demand fetch;
+ // Next stays disabled until it lands (item 5).
+ const _mlxUnavail = _platform === 'apple' && _ftEngine === 'mlx'
+ && baseModelMeta.mlx_available === false;
+ _renderMlxDownload('new-ft-download', baseModelKey, baseModelMeta, _mlxUnavail, validateNewFt);
let error = '';
if (_newFtNoDatasets) error = 'Create a dataset first';
else if (baseModelKey && baseModelMeta.registered === false && _ftEngine !== 'mlx')
error = `Base model "${baseModelMeta.label || baseModelKey}" not downloaded — please run installer again`;
+ else if (_mlxUnavail) error = '_mlx_download';
else if (!name) error = '_name_empty';
else if (runsData.some(r => r.display_name === name)) error = `Run "${name}" already exists`;
else if (isNaN(steps) || steps < 100) error = 'Steps must be >= 100';
@@ -6230,13 +6310,14 @@
Download Audio Selection
// Show name errors next to the label; other errors in the estimate area
const nameLabelEl = document.getElementById('new-ft-name').parentElement.querySelector('.modal-label');
const isNameError = error === '_name_empty' || error.startsWith('Run "');
+ const isMlxDl = error === '_mlx_download'; // shown via the download panel
if (nameLabelEl) {
if (error === '_name_empty') nameLabelEl.innerHTML = 'Run name
* ';
else if (error.startsWith('Run "')) nameLabelEl.innerHTML = `Run name
${error} `;
else nameLabelEl.innerHTML = 'Run name';
}
btn.disabled = !!error;
- if (error && !isNameError) { el.style.color = 'var(--red)'; el.textContent = error; }
+ if (error && !isNameError && !isMlxDl) { el.style.color = 'var(--red)'; el.textContent = error; }
else { el.style.color = 'var(--yellow)'; updateCkptEstimate(); }
updateLayerPreview();
};
diff --git a/dashboard/server.py b/dashboard/server.py
index 916e6e9..2f7f467 100644
--- a/dashboard/server.py
+++ b/dashboard/server.py
@@ -2564,6 +2564,14 @@ def _query_gpu_compute_caps() -> dict:
_gradio_vram = {"load_mb": 10000, "peak_mb": 12000, "n_samples": 0}
_gradio_vram_baselines = {} # instance_id -> {"gpu": int, "before_mb": int, "measured": bool}
+# MLX model-pack downloads (Apple on-demand, item 5). dit -> {proc, log, started}.
+# Each is a subprocess running mlx_engine.build_mlx_download_cmd (weights.
+# ensure_local for the base+arc+codec+t5 pack) with stdout tee'd to a log the
+# status endpoint tails for progress.
+_mlx_downloads = {}
+_mlx_downloads_lock = threading.Lock()
+_MLX_DOWNLOAD_DIR = STATE_FILES_DIR / "mlx_downloads"
+
def _load_gradio_vram_estimate():
global _gradio_vram
@@ -3268,7 +3276,32 @@ def do_GET(self):
self.send_response(404)
self.end_headers()
elif path == "/api/models":
- self._json_response({"models": MODELS_UI_PAYLOAD})
+ payload = MODELS_UI_PAYLOAD
+ # On Apple, annotate each model with whether its MLX weight pack
+ # (base + arc + codec + encoder + t5) is present, so the UI can
+ # offer an on-demand download instead of the torch 'registered' flag.
+ if _detect_platform() == "apple":
+ try:
+ from underfit.backends import mlx_engine
+ payload = {}
+ for k, v in MODELS_UI_PAYLOAD.items():
+ vv = dict(v)
+ try:
+ vv["mlx_available"] = mlx_engine.mlx_model_available(k)
+ _miss, _gb = mlx_engine.mlx_missing_pack(
+ mlx_engine.map_model_name(k))
+ vv["mlx_download_gb"] = _gb
+ except Exception:
+ vv["mlx_available"] = None
+ payload[k] = vv
+ except Exception:
+ payload = MODELS_UI_PAYLOAD
+ self._json_response({"models": payload})
+ elif path == "/api/models/download":
+ # Status of an on-demand MLX pack download. ?model=
+ qs = parse_qs(parsed.query)
+ model = (qs.get("model") or [None])[0]
+ self._json_response(self._mlx_download_status(model))
elif path == "/api/runs":
self._json_response(self._get_runs())
elif path == "/api/status":
@@ -3598,6 +3631,10 @@ def do_POST(self):
self._json_response({"error": err}, status=409)
else:
self._json_response({"id": instance_id}, status=201)
+ elif parsed.path == "/api/models/download":
+ length = int(self.headers.get("Content-Length", 0))
+ body = json.loads(self.rfile.read(length)) if length else {}
+ self._start_mlx_download(body.get("model"))
elif parsed.path == "/api/runs/new":
length = int(self.headers.get("Content-Length", 0))
body = json.loads(self.rfile.read(length)) if length else {}
@@ -5251,6 +5288,88 @@ def _get_gpu_info(self):
resp["cuda_visible_devices"] = visible_gpus
return resp
+ def _start_mlx_download(self, model):
+ """Kick off (or report) an on-demand MLX pack download for `model`.
+ Apple-only; the pack (base + arc + codec + encoder + t5) is fetched into
+ the existing MLX checkout via the MLX venv's weights.ensure_local."""
+ if _detect_platform() != "apple":
+ self._json_response(
+ {"error": "MLX model download is Apple-only."}, status=400)
+ return
+ try:
+ from underfit.backends import mlx_engine
+ dit = mlx_engine.map_model_name(model)
+ except Exception as e:
+ self._json_response({"error": str(e)}, status=400)
+ return
+ if mlx_engine.mlx_model_available(model):
+ self._json_response({"ok": True, "state": "available"})
+ return
+ with _mlx_downloads_lock:
+ cur = _mlx_downloads.get(dit)
+ if cur and cur["proc"].poll() is None:
+ self._json_response({"ok": True, "state": "downloading", "dit": dit})
+ return
+ try:
+ cmd = mlx_engine.build_mlx_download_cmd(dit)
+ mlx_root, _ = mlx_engine.resolve_mlx_paths()
+ except FileNotFoundError as e:
+ self._json_response({"error": str(e)}, status=400)
+ return
+ _MLX_DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)
+ log_path = _MLX_DOWNLOAD_DIR / f"{dit}.log"
+ try:
+ logf = open(log_path, "w")
+ except OSError as e:
+ self._json_response({"error": str(e)}, status=500)
+ return
+ env = dict(os.environ)
+ env.setdefault("PYTHONUNBUFFERED", "1")
+ proc = subprocess.Popen(
+ cmd, cwd=mlx_root, stdout=logf, stderr=subprocess.STDOUT, env=env)
+ _mlx_downloads[dit] = {
+ "proc": proc, "log": str(log_path), "started": time.time()}
+ self._json_response(
+ {"ok": True, "state": "downloading", "dit": dit}, status=201)
+
+ def _mlx_download_status(self, model):
+ """Progress of an on-demand MLX pack download for `model`."""
+ try:
+ from underfit.backends import mlx_engine
+ dit = mlx_engine.map_model_name(model)
+ except Exception as e:
+ return {"state": "error", "error": str(e), "available": False}
+ available = mlx_engine.mlx_model_available(model)
+ total = len(mlx_engine.mlx_pack_rel_paths(dit))
+ with _mlx_downloads_lock:
+ cur = _mlx_downloads.get(dit)
+ state = "available" if available else "idle"
+ done_files, current = 0, None
+ if cur:
+ rc = cur["proc"].poll()
+ try:
+ log = Path(cur["log"]).read_text()
+ except Exception:
+ log = ""
+ oks = set(re.findall(r"OK (\S+)", log))
+ dls = re.findall(r"DL (\S+)", log)
+ done_files = len(oks)
+ pending = [d for d in dls if d not in oks]
+ current = os.path.basename(pending[-1]) if pending else None
+ if rc is None:
+ state = "downloading"
+ elif "ALL_DONE" in log or available:
+ state = "done"
+ else:
+ state = "error"
+ return {
+ "state": state, "available": available,
+ "done_files": done_files, "total_files": total,
+ "current_file": current,
+ "remaining_gb": (mlx_engine.mlx_missing_pack(dit)[1]
+ if not available else 0),
+ }
+
def _descendant_pids(self, pid):
"""Set of all descendant PIDs of pid (recursive pgrep -P)."""
found = set()
diff --git a/underfit/backends/mlx_engine.py b/underfit/backends/mlx_engine.py
index 30560fd..76aa47b 100644
--- a/underfit/backends/mlx_engine.py
+++ b/underfit/backends/mlx_engine.py
@@ -94,6 +94,88 @@ def map_model_name(base_model):
return name
+# --------------------------------------------------------------------------- #
+# MLX model packs (weights) — availability + download
+# --------------------------------------------------------------------------- #
+# A model's MLX "pack": base npz (LoRA training) + arc npz (inference / ARC
+# demos) + SAME decoder + SAME encoder (a2a / dataset encode) + the shared
+# t5gemma conditioner. Mirrors scripts/weights.py (DIT_BUNDLES + SHARED +
+# TRAINING_BASE) in the MLX checkout — that file stays the download source of
+# truth (ensure_local); this list drives availability + size for the
+# dashboard/installer WITHOUT importing the MLX venv. Sizes (GB) approximate.
+_MLX_FILE_GB = {
+ "dit_medium-base_f16.npz": 2.70, "dit_medium_f16.npz": 2.70,
+ "same_l_decoder_f32.npz": 1.58, "same_l_encoder_f32.npz": 1.58,
+ "dit_sm-music-base_f16.npz": 0.85, "dit_sm-music_f16.npz": 0.85,
+ "dit_sm-sfx-base_f16.npz": 0.85, "dit_sm-sfx_f16.npz": 0.85,
+ "same_s_decoder_f32.npz": 0.20, "same_s_encoder_f32.npz": 0.20,
+ "t5gemma_f16.npz": 0.52,
+}
+
+
+def mlx_pack_rel_paths(dit):
+ """The models/mlx/*.npz a --dit model needs (base + arc + codec + t5)."""
+ codec = "same_l" if dit == "medium" else "same_s"
+ return [
+ f"models/mlx/dit_{dit}-base_f16.npz", # base — LoRA training
+ f"models/mlx/dit_{dit}_f16.npz", # arc — inference / ARC demos
+ f"models/mlx/{codec}_decoder_f32.npz", # decode
+ f"models/mlx/{codec}_encoder_f32.npz", # encode (a2a / dataset)
+ "models/mlx/t5gemma_f16.npz", # shared conditioner
+ ]
+
+
+def mlx_root_or_none():
+ """MLX checkout root if it resolves on disk, else None (no raise) — for
+ availability checks before the checkout/venv are known to exist."""
+ root = os.environ.get("UNDERFIT_MLX_ROOT") or os.path.join(
+ _SA3_LOCAL, "optimized", "mlx")
+ root = os.path.abspath(root)
+ return root if os.path.isdir(root) else None
+
+
+def mlx_missing_pack(dit):
+ """(missing_rel_paths, approx_gb) for a --dit model's MLX pack. Everything
+ counts as missing if the checkout itself isn't present yet."""
+ root = mlx_root_or_none()
+ paths = mlx_pack_rel_paths(dit)
+ missing = [p for p in paths
+ if not (root and os.path.isfile(os.path.join(root, p)))]
+ gb = round(sum(_MLX_FILE_GB.get(os.path.basename(p), 0.0) for p in missing), 1)
+ return missing, gb
+
+
+def mlx_model_available(base_model):
+ """True if the model's MLX pack is fully present on disk (base + arc + codec
+ + encoder + t5). base_model may be 'sa3-medium' or 'medium'."""
+ try:
+ dit = map_model_name(base_model)
+ except ValueError:
+ return False
+ missing, _ = mlx_missing_pack(dit)
+ return not missing
+
+
+def build_mlx_download_cmd(dit):
+ """MLX-venv-python command that downloads a model's whole pack via the MLX
+ checkout's weights.ensure_local (the download source of truth). Emits
+ 'DL ' / 'OK ' per file and 'ALL_DONE' at the end so callers can
+ show progress. Run with cwd=."""
+ _, mlx_python = resolve_mlx_paths()
+ paths = mlx_pack_rel_paths(dit)
+ prog = (
+ "import sys; sys.path.insert(0, 'scripts')\n"
+ "from weights import ensure_local\n"
+ f"paths = {paths!r}\n"
+ "for p in paths:\n"
+ " print('DL ' + p, flush=True)\n"
+ " ensure_local(p, verbose=False)\n"
+ " print('OK ' + p, flush=True)\n"
+ "print('ALL_DONE', flush=True)\n"
+ )
+ return [mlx_python, "-c", prog]
+
+
def resolve_base_weights(dit_model):
"""Path to the MLX BASE-checkpoint DiT weights npz for a --dit value.
diff --git a/underfit/cli/setup.py b/underfit/cli/setup.py
index fc9b3db..d8a2668 100644
--- a/underfit/cli/setup.py
+++ b/underfit/cli/setup.py
@@ -826,6 +826,59 @@ def run_model_phase(args, backend: Backend) -> int:
# ── ENTRY POINT ──────────────────────────────────────────────────────────────
+def maybe_download_mlx_packs(args) -> int:
+ """Apple-only: pre-download the MLX weight packs (base + ARC + SAME codec +
+ encoder + shared t5gemma) into the existing MLX checkout so the first MLX
+ finetune / dataset-encode doesn't stall. Reuses underfit.backends.mlx_engine
+ — the same download the dashboard's on-demand button uses (item 5). Degrades
+ gracefully: if the MLX checkout / venv isn't set up, the weights just
+ download lazily on first use instead."""
+ import subprocess as _sp
+ try:
+ from underfit.backends import mlx_engine
+ except Exception as e:
+ print(f"\n(could not import the MLX engine — {e}. MLX weights will "
+ f"download on first use.)")
+ return 0
+ root = mlx_engine.mlx_root_or_none()
+ if not root:
+ print("\n▸ Apple Silicon detected, but the MLX checkout wasn't found.\n"
+ " Set UNDERFIT_MLX_ROOT (or clone stable-audio-3 as a sibling and\n"
+ " create its .venv). MLX weights will download on first use.")
+ return 0
+ all_keys = [p.key for p in SA3_PACKS]
+ if args.models:
+ wanted = {s.strip() for s in args.models.split(",") if s.strip()}
+ keys = [k for k in all_keys if k in wanted]
+ else:
+ keys = all_keys
+ print(f"\n▸ MLX weight packs → {os.path.join(root, 'models', 'mlx')}")
+ rc = 0
+ for key in keys:
+ try:
+ dit = mlx_engine.map_model_name(key)
+ except Exception:
+ continue
+ if mlx_engine.mlx_model_available(dit):
+ print(f" ✓ {dit}: already present")
+ continue
+ _missing, gb = mlx_engine.mlx_missing_pack(dit)
+ print(f" ↓ {dit}: downloading ~{gb} GB …", flush=True)
+ try:
+ cmd = mlx_engine.build_mlx_download_cmd(dit)
+ except FileNotFoundError as e:
+ print(f" skipped — {e}")
+ rc = 1
+ continue
+ try:
+ _sp.run(cmd, cwd=root, check=True)
+ print(f" ✓ {dit}: done")
+ except Exception as e:
+ print(f" download failed ({e}) — will retry on first use")
+ rc = 1
+ return rc
+
+
def main(argv: list[str] | None = None) -> int:
p = argparse.ArgumentParser(
prog="underfit-setup",
@@ -873,6 +926,14 @@ def main(argv: list[str] | None = None) -> int:
print("\n(skipping model phase as requested)")
return 0
+ import platform as _pf
+ if _pf.system() == "Darwin" and _pf.machine() == "arm64":
+ # Apple Silicon: the MLX engine (Metal GPU) is the primary path, so
+ # download the MLX weight packs. The torch packs (run_model_phase) are
+ # only needed for the torch/MPS engine — skip them here to avoid pulling
+ # ~10 GB of torch weights an MLX-only Mac won't use.
+ return maybe_download_mlx_packs(args)
+
return run_model_phase(args, backend)
From 08c0003576ff1a9ac96dff1d27c1687ac889b06c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 18:26:25 -0400
Subject: [PATCH 13/19] mlx-engine: pass demo `duration` through to the MLX
trainer
_build_demo_entries dropped the per-demo `duration` field, so full-length demos
collapsed to the crop length in the MLX trainer. Forward it (like seed /
lora_strength / lora_interval_max) so each demo renders at its configured length.
---
underfit/backends/mlx_engine.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/underfit/backends/mlx_engine.py b/underfit/backends/mlx_engine.py
index 76aa47b..ecdfdb8 100644
--- a/underfit/backends/mlx_engine.py
+++ b/underfit/backends/mlx_engine.py
@@ -384,7 +384,7 @@ def _build_demo_entries(demo):
entry = {"prompt": e.get("prompt", ""),
"cfg": e.get("cfg", default_cfg),
"steps": e.get("steps", default_steps)}
- for k in ("seed", "lora_strength", "lora_interval_max"):
+ for k in ("seed", "lora_strength", "lora_interval_max", "duration"):
if e.get(k) is not None:
entry[k] = e[k]
entries.append(entry)
From ecc15e1e4ca5e592f0c8f73da8575c907a9a0bca Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 18:33:39 -0400
Subject: [PATCH 14/19] dashboard: don't offer MPS without a torch backend; fix
greyed-silent gate
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
On Apple the engine dropdown always offered "MPS" (torch) even with no torch
backend installed, so starting an MPS run greyed out Next with no warning — and
only a name change surfaced "not downloaded — run installer again".
- server: /api/gpu `engines` now reflects installed backends — mlx on Apple
(its checkout brings its own stack) plus torch only where stable_audio_3 /
stable_audio_tools is importable. An MLX-only Mac reports ["mlx"], so the UI
omits the unusable MPS option entirely rather than showing a greyed dead-end.
- index.html: _applyEngineOptions hides/disables the torch (MPS) option when
torch isn't in _engines (it previously only handled mlx).
- Fixed the greyed-silent gate itself: validateNewFt now flags a blocking error
and updateCkptEstimate bails instead of overwriting it. The onchanges call
both, so the size estimate was clobbering the "not downloaded" message, which
then only reappeared on a name change. Any validation error shows immediately.
---
dashboard/index.html | 13 +++++++++++++
dashboard/server.py | 16 +++++++++++++++-
2 files changed, 28 insertions(+), 1 deletion(-)
diff --git a/dashboard/index.html b/dashboard/index.html
index 68b02ac..38bdd1e 100644
--- a/dashboard/index.html
+++ b/dashboard/index.html
@@ -3806,6 +3806,7 @@ Download Audio Selection
const sel = document.getElementById('new-ft-engine');
if (!sel) return;
const hasMlx = _engines.includes('mlx');
+ const hasTorch = _engines.includes('torch');
const isApple = _platform === 'apple';
Array.from(sel.options).forEach(o => {
if (o.value === 'mlx') {
@@ -3813,6 +3814,10 @@ Download Audio Selection
o.disabled = !hasMlx;
o.textContent = isApple ? 'MLX' : 'mlx';
} else if (o.value === 'torch') {
+ // Hide MPS/torch when no torch backend is installed — an unusable engine
+ // option is worse than its absence (server omits it from _engines).
+ o.hidden = !hasTorch;
+ o.disabled = !hasTorch;
o.textContent = isApple ? 'MPS' : 'torch';
}
});
@@ -5833,6 +5838,10 @@ Download Audio Selection
};
function updateCkptEstimate() {
+ // Don't overwrite a validation error currently shown in this element — the
+ // onchanges call validateNewFt() then updateCkptEstimate(), so the error
+ // (set by validateNewFt) must win over the size estimate.
+ if (window._newFtHasBlockingErr) return;
const el = document.getElementById('new-ft-ckpt-estimate');
const model = document.getElementById('new-ft-base-model').value;
const type = document.getElementById('new-ft-lora-type').value;
@@ -6317,6 +6326,10 @@ Download Audio Selection
else nameLabelEl.innerHTML = 'Run name';
}
btn.disabled = !!error;
+ // Flag a blocking error shown in `el` so a following updateCkptEstimate()
+ // (onchanges call both) doesn't clobber the message — previously the error
+ // only surfaced on a name change, leaving Next greyed with no warning.
+ window._newFtHasBlockingErr = !!(error && !isNameError && !isMlxDl);
if (error && !isNameError && !isMlxDl) { el.style.color = 'var(--red)'; el.textContent = error; }
else { el.style.color = 'var(--yellow)'; updateCkptEstimate(); }
updateLayerPreview();
diff --git a/dashboard/server.py b/dashboard/server.py
index 2f7f467..c5c9417 100644
--- a/dashboard/server.py
+++ b/dashboard/server.py
@@ -2726,6 +2726,20 @@ def _cuda_env_prefix(gpu):
return f"CUDA_VISIBLE_DEVICES={gpu} "
+def _available_engines(plat):
+ """Training engines the dashboard can actually launch. 'mlx' on Apple (the
+ MLX checkout brings its own stack); 'torch' only where a torch backend
+ (stable_audio_3 / stable_audio_tools) is importable — never list an engine
+ whose backend isn't installed, since a greyed, unusable option (e.g. MPS on
+ an MLX-only Mac) is worse than its absence."""
+ import importlib.util
+ torch_ok = any(importlib.util.find_spec(m) is not None
+ for m in ("stable_audio_3", "stable_audio_tools"))
+ if plat == "apple":
+ return ["mlx"] + (["torch"] if torch_ok else [])
+ return ["torch"]
+
+
def _get_gpu_count(force_refresh=False):
"""Return number of CUDA GPUs from nvidia-smi (cached after the first
*successful* call). Returns 0 if nvidia-smi is missing or fails — the UI
@@ -5283,7 +5297,7 @@ def _get_gpu_info(self):
resp = {"gpus": gpus, "gradio_estimate": _gradio_vram, "arc_info": arc_info,
"platform": plat,
- "engines": ["torch", "mlx"] if plat == "apple" else ["torch"]}
+ "engines": _available_engines(plat)}
if visible_gpus is not None:
resp["cuda_visible_devices"] = visible_gpus
return resp
From fba6445587221708496cf681596ce7666fe8ee96 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 18:42:56 -0400
Subject: [PATCH 15/19] dashboard: pass base_model to the MLX gradio launch
(--pretrained-name)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Launching an MLX gradio failed with "could not determine the --dit model":
resolve_dit_model reads 'base_model' from the run's _model.json (which doesn't
carry it) and the launch never passed --pretrained-name. The launch already
knows base_model from the run record — forward it so the MLX gradio picks the
--dit value (the torch path ignores --pretrained-name).
---
dashboard/server.py | 4 ++++
1 file changed, 4 insertions(+)
diff --git a/dashboard/server.py b/dashboard/server.py
index c5c9417..2f0ce00 100644
--- a/dashboard/server.py
+++ b/dashboard/server.py
@@ -1535,6 +1535,10 @@ def launch(self, checkpoint_path, gpu, run_id=None, checkpoint_name=None, title=
f"--ckpt-path {_bash_quote(ckpt_path_model)} "
f"{lora_args} "
f"--engine {engine} "
+ # base_model name for engine=mlx: the run's _model.json has no
+ # top-level 'base_model', so resolve_dit_model needs it here to
+ # pick the --dit value (the torch path ignores --pretrained-name).
+ f"--pretrained-name {shlex.quote(base_model)} "
f"--model-half "
f"--title {shlex.quote(title)}"
f"{default_prompt_arg}"
From 054928026f9c17900f06f5aa96e589f15d348048 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 19:01:26 -0400
Subject: [PATCH 16/19] dashboard: default MLX demos to ARC-only (skip the slow
RF demos)
RF/Base demos are 50 steps (~8 min each at full length on MLX); ARC is 8 steps.
For the MLX engine the suggested default now drops the RF demos and keeps only
the ARC ones (all presets), falling back to RF only when the model has no ARC
variant. _defaultNewFtDemos / _defaultDemoDurationLatents / the
unconditional-prompt indices (_demoEmptyIdx) are engine-aware, and newFtGoStep4
rebuilds the default with the current engine on entry.
---
dashboard/index.html | 67 +++++++++++++++++++++++++++-----------------
1 file changed, 42 insertions(+), 25 deletions(-)
diff --git a/dashboard/index.html b/dashboard/index.html
index 38bdd1e..a9ba442 100644
--- a/dashboard/index.html
+++ b/dashboard/index.html
@@ -4758,23 +4758,34 @@ Download Audio Selection
let _newFtPreset = 8;
let _newFtPresetUserSet = false; // true once the user clicks a Preset button
+// True when the New Finetune engine is MLX (Apple). MLX RF/Base demos take ~50
+// forwards — ~8 min each at full length — so with MLX we suggest ARC-only demos
+// (8 steps). Falls back to RF only when the model has no ARC variant.
+function _engineIsMlx() {
+ return ((document.getElementById('new-ft-engine') || {}).value || 'torch') === 'mlx';
+}
+function _modelHasArc() {
+ const mi = _arcInfo[(document.getElementById('new-ft-base-model') || {}).value] || {};
+ return mi.arc_type === 'full_model' || mi.arc_type === 'lora';
+}
+
function _defaultNewFtDemos(preset) {
const model = document.getElementById('new-ft-base-model').value || 'sa3';
const mi = _arcInfo[model] || {};
const hasArc = mi.arc_type === 'full_model' || mi.arc_type === 'lora';
const sampler = mi.diffusion_objective === 'v' ? 'V' : 'Base';
const baseSteps = mi.diffusion_objective === 'v' ? 100 : 50;
+ // MLX: drop the RF/Base demos (too slow) and suggest ARC-only.
+ const arcOnly = _engineIsMlx() && hasArc;
if (preset === undefined || preset === null) preset = _newFtPreset;
if (preset === 2) {
- const demos = [{prompt: '', sampler, cfg: 7, steps: baseSteps, seed: _randomSeed()}];
+ const demos = arcOnly ? [] : [{prompt: '', sampler, cfg: 7, steps: baseSteps, seed: _randomSeed()}];
if (hasArc) demos.push({prompt: '', sampler: 'ARC', cfg: 1, steps: 8, seed: _randomSeed()});
return demos;
}
if (preset === 4) {
- // Order matters: indices used by _populateDemoPrompts to clear prompt for "no prompt" slots.
- // idx 0: RF + prompt idx 1: RF + no prompt
- // idx 2: ARC + prompt idx 3: ARC + no prompt
- const demos = [
+ // Order matters: indices used by _demoEmptyIdx to clear prompt for "no prompt" slots.
+ const demos = arcOnly ? [] : [
{prompt: '', sampler, cfg: 7, steps: baseSteps, seed: _randomSeed()},
{prompt: '', sampler, cfg: 7, steps: baseSteps, seed: _randomSeed()},
];
@@ -4787,7 +4798,7 @@ Download Audio Selection
return demos;
}
// preset === 8 (default suite)
- const demos = [
+ const demos = arcOnly ? [] : [
{prompt: '', sampler: sampler, cfg: 7, steps: baseSteps, seed: _randomSeed()},
{prompt: '', sampler: sampler, cfg: 7, steps: baseSteps, seed: _randomSeed()},
{prompt: '', sampler: sampler, cfg: 1, steps: baseSteps, seed: _randomSeed()},
@@ -4805,19 +4816,25 @@ Download Audio Selection
}
// "Unconditional" indices per preset — emptied by _populateDemoPrompts after the
-// dataset-driven prompt fill. Two: none. Four: 2nd of each sampler. Eight: 4th
-// RF + 3rd ARC (legacy positions).
-const _DEMO_PRESET_EMPTY_IDX = {2: [], 4: [1, 3], 8: [3, 6]};
-
-// Default per-demo duration depending on preset:
-// Eight — demos 0 (first RF) and 4 (first ARC) default to the model's native
-// length so the user always has a long-form reference; the rest use
-// the training crop length.
-// Two/Four — every demo defaults to the training crop length. These presets
-// exist for quick previews on slow GPUs; full-length defeats the point.
+// dataset-driven prompt fill. RF+ARC layout: Two none, Four 2nd of each sampler,
+// Eight 4th RF + 3rd ARC. ARC-only (MLX) drops the RF slots, so the
+// unconditional slot shifts down accordingly.
+function _demoEmptyIdx(preset) {
+ const p = preset || _newFtPreset;
+ const arcOnly = _engineIsMlx() && _modelHasArc();
+ const layout = arcOnly ? {2: [], 4: [1], 8: [2]} : {2: [], 4: [1, 3], 8: [3, 6]};
+ return layout[p] || [];
+}
+
+// Default per-demo duration: the first demo of each sampler gets the model's
+// native (full) length as a long-form reference (preset 8 only); the rest use
+// the training crop length. RF+ARC → idx 0 (RF) + 4 (ARC); ARC-only (MLX) → the
+// first ARC is idx 0.
function _defaultDemoDurationLatents(idx, cropLen, nativeLen) {
- if (_newFtPreset === 8 && (idx === 0 || idx === 4)) return nativeLen;
- return cropLen;
+ if (_newFtPreset !== 8) return cropLen;
+ const arcOnly = _engineIsMlx() && _modelHasArc();
+ const nativeIdx = arcOnly ? [0] : [0, 4];
+ return nativeIdx.includes(idx) ? nativeLen : cropLen;
}
let _newFtDemos = _defaultNewFtDemos(8);
@@ -5084,13 +5101,13 @@ Download Audio Selection
};
window.newFtGoStep4 = function() {
- // First entry (not cloned, user hasn't manually picked): auto-pick by GPU.
+ // First entry (not cloned, user hasn't manually picked): auto-pick by GPU and
+ // rebuild the default demos with the CURRENT engine (MLX → ARC-only). Always
+ // rebuild here (not only when the preset changes) so the engine-aware default
+ // applies even when the preset is unchanged; prompts are repopulated below.
if (!_newFtDemosCloned && !_newFtPresetUserSet) {
- const auto = _pickPresetForGpu();
- if (auto !== _newFtPreset) {
- _newFtPreset = auto;
- _newFtDemos = _defaultNewFtDemos(auto);
- }
+ _newFtPreset = _pickPresetForGpu();
+ _newFtDemos = _defaultNewFtDemos(_newFtPreset);
}
_updatePresetUI();
// Populate demo prompts from the prompt config heuristic (skip if cloned)
@@ -5182,7 +5199,7 @@ Download Audio Selection
}
// Clear prompts for the preset's "unconditional" slots — Two has none,
// Four clears the second of each sampler, Eight clears idx 3 + 6.
- const emptyIdx = _DEMO_PRESET_EMPTY_IDX[preset || _newFtPreset] || [];
+ const emptyIdx = _demoEmptyIdx(preset);
for (const i of emptyIdx) {
if (_newFtDemos.length > i) _newFtDemos[i].prompt = '';
}
From d8a6a913d7175ecf727262e46ee29674e634c440 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 19:10:51 -0400
Subject: [PATCH 17/19] dashboard: SA3-medium "Use SAME-S for demos (faster)"
toggle
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
New Finetune → Demos shows a "Use SAME-S for demos (faster)" checkbox for
SA3-medium only — default ON for Apple/MLX, off elsewhere. When on, demos decode
with the small SAME-S autoencoder instead of the model's SAME-L (~1.6× faster,
near-identical: the two share the latent space, corr 0.99). Wired
payload.demo_decoder → training.demo.demo_decoder → the MLX trainer's
--demo-decoder (mlx_engine); the torch loop keeps its own decoder.
---
dashboard/index.html | 33 +++++++++++++++++++++++++++++++++
dashboard/server.py | 5 +++++
underfit/backends/mlx_engine.py | 2 ++
3 files changed, 40 insertions(+)
diff --git a/dashboard/index.html b/dashboard/index.html
index a9ba442..8d698ce 100644
--- a/dashboard/index.html
+++ b/dashboard/index.html
@@ -1481,6 +1481,16 @@ New Finetune Demos
steps
+
+
+
+
+
+
@@ -4769,6 +4779,24 @@
Download Audio Selection
return mi.arc_type === 'full_model' || mi.arc_type === 'lora';
}
+// SA3-medium only: the "Use SAME-S for demos (faster)" checkbox. Default on for
+// Apple/MLX (SAME-L decode is a big share of demo time there), off elsewhere;
+// _sameSUserSet preserves an explicit toggle for the rest of the modal session.
+let _sameSUserSet = false;
+function _updateSameSDemoOption() {
+ const row = document.getElementById('new-ft-same-s-row');
+ const cb = document.getElementById('new-ft-same-s-demos');
+ if (!row || !cb) return;
+ const isMedium = (document.getElementById('new-ft-base-model') || {}).value === 'sa3-medium';
+ row.style.display = isMedium ? 'flex' : 'none';
+ if (isMedium && !_sameSUserSet) cb.checked = (_platform === 'apple');
+}
+function _useSameSForDemos() {
+ const cb = document.getElementById('new-ft-same-s-demos');
+ return (document.getElementById('new-ft-base-model') || {}).value === 'sa3-medium'
+ && cb && cb.checked;
+}
+
function _defaultNewFtDemos(preset) {
const model = document.getElementById('new-ft-base-model').value || 'sa3';
const mi = _arcInfo[model] || {};
@@ -5117,6 +5145,7 @@ Download Audio Selection
document.getElementById('new-ft-step4').style.display = '';
document.getElementById('new-ft-demo-every-display').textContent =
document.getElementById('new-ft-demo-every').value || '1000';
+ _updateSameSDemoOption(); // SA3-medium: show + default the SAME-S demo toggle
_renderDemoEntries();
};
@@ -6206,6 +6235,7 @@ Download Audio Selection
// Reset demos state — auto-pick preset will run when user reaches Step 4
_newFtPreset = 8;
_newFtPresetUserSet = false;
+ _sameSUserSet = false; // re-default the SAME-S demo toggle per modal session
_newFtDemos = _defaultNewFtDemos(8);
_newFtDemosCloned = false;
// Reset Step 2 state
@@ -6420,6 +6450,9 @@ Download Audio Selection
}
return entry;
});
+ // SA3-medium: decode demos with SAME-S instead of the model's SAME-L when the
+ // "faster demos" toggle is on (default on for Apple/MLX).
+ if (_useSameSForDemos()) payload.demo_decoder = 'same-s';
// Build ground_truth from source files used to generate each demo prompt
payload.ground_truth = _newFtDemos.map(d => {
const f = d.sourceFile;
diff --git a/dashboard/server.py b/dashboard/server.py
index 2f0ce00..75b52f8 100644
--- a/dashboard/server.py
+++ b/dashboard/server.py
@@ -4023,6 +4023,11 @@ def _handle_new_finetune(self, body):
cfg["training"].setdefault("demo", {})["demo_every"] = demo_every
cfg["training"]["demo"]["demo_mode"] = "lora_dashboard"
cfg["training"]["demo"]["latent_crop_length"] = body.get("latent_crop_length", mi["latent_crop_length"])
+ # SA3-medium (MLX): decode demos with SAME-S instead of SAME-L when
+ # the "faster demos" toggle is on — the MLX trainer honors this via
+ # --demo-decoder (mlx_engine); the torch loop keeps its own decoder.
+ if body.get("demo_decoder"):
+ cfg["training"]["demo"]["demo_decoder"] = body["demo_decoder"]
# Apply custom demo_cond from frontend if provided
if custom_demo_cond:
# Compute seconds_total from the actual latent crop length
diff --git a/underfit/backends/mlx_engine.py b/underfit/backends/mlx_engine.py
index ecdfdb8..c05abb9 100644
--- a/underfit/backends/mlx_engine.py
+++ b/underfit/backends/mlx_engine.py
@@ -499,6 +499,8 @@ def build_trainer_cmd(args, base_weights=None):
if entries:
_add(cmd, "--demo-config", _write_demo_config(entries))
_add(cmd, "--demo-every", demo_every)
+ # SA3-medium "faster demos": decode with SAME-S instead of SAME-L.
+ _add(cmd, "--demo-decoder", demo.get("demo_decoder"))
return cmd
From 899dfc863a9692849b2f420d98c454485d3ad32c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 19:44:12 -0400
Subject: [PATCH 18/19] dashboard: relabel demo presets One/Two/Four on MLX
(ARC-only counts)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
MLX demos are ARC-only, so a preset yields half the demos of the RF+ARC layout
(2→1, 4→2, 8→4). Relabel the preset buttons One/Two/Four on MLX to match the
actual ARC-demo count; non-MLX keeps Two/Four/Eight (half RF + half ARC). The
internal preset values (2/4/8) are unchanged — labels only.
---
dashboard/index.html | 7 +++++++
1 file changed, 7 insertions(+)
diff --git a/dashboard/index.html b/dashboard/index.html
index 8d698ce..12e08a8 100644
--- a/dashboard/index.html
+++ b/dashboard/index.html
@@ -5105,9 +5105,16 @@ Download Audio Selection
}
function _updatePresetUI() {
+ // MLX (ARC-only) demos drop the RF half, so a preset yields half the demos —
+ // relabel the buttons One/Two/Four to match the actual count. Non-MLX keeps
+ // Two/Four/Eight (half RF + half ARC).
+ const arcOnly = _engineIsMlx() && _modelHasArc();
+ const labels = arcOnly ? {2: 'One', 4: 'Two', 8: 'Four'}
+ : {2: 'Two', 4: 'Four', 8: 'Eight'};
for (const k of [2, 4, 8]) {
const btn = document.getElementById('new-ft-preset-' + k);
if (!btn) continue;
+ btn.textContent = labels[k];
if (_newFtPreset === k) {
btn.classList.remove('modal-cancel');
btn.classList.add('modal-launch');
From de117eab01ef944c02245ecabb5ac34c72ee2689 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E2=80=9CCortexelus=E2=80=9D?= <“emperorcj@gmail.com”>
Date: Wed, 15 Jul 2026 19:47:32 -0400
Subject: [PATCH 19/19] dashboard: hide dotfile scratch from the checkpoint
lists
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
_find_checkpoints included any *.safetensors, so the MLX ARC-demo scratch LoRA
(.arc_demo_tmp_*.safetensors) showed up in the Checkpoints box and the
Revive/launch checkpoint pickers. Skip dotfiles in the scan — hides that
scratch (and any hidden temp) from every checkpoint list.
---
dashboard/server.py | 2 ++
1 file changed, 2 insertions(+)
diff --git a/dashboard/server.py b/dashboard/server.py
index 75b52f8..425b4bc 100644
--- a/dashboard/server.py
+++ b/dashboard/server.py
@@ -6840,6 +6840,8 @@ def _find_checkpoints(ckpts_dir, run_id, dataset_history=None):
if not os.path.isdir(ckpt_subdir):
continue
for f in os.scandir(ckpt_subdir):
+ if f.name.startswith("."):
+ continue # hidden scratch (e.g. .arc_demo_tmp — MLX ARC-demo temp LoRA)
if f.name.endswith(".safetensors"):
all_ckpts[f.path] = (Path(f.path), f.stat())
elif f.name.endswith(".ckpt"):