MLX engine for training + gradio (Apple Silicon)

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
 

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,13 +132,27 @@ 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
 
 
 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()

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)