From 5262d93b5875caf15e86e399a80f6a1118281b6b Mon Sep 17 00:00:00 2001
From: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Date: Sat, 2 May 2026 08:51:56 +0400
Subject: [PATCH] studio: add --local to setup.sh + overlay unsloth-zoo from
 git main (#5252)

* studio: add --local to setup.sh + overlay unsloth-zoo from git main

setup.sh now accepts --local, which exports STUDIO_LOCAL_INSTALL=1 and
STUDIO_LOCAL_REPO=$REPO_ROOT. install_python_stack.py overlays unsloth-zoo
from git main on top of the editable unsloth checkout in both local_repo
branches (no-torch and with-torch).

The Colab notebook now invokes ./studio/setup.sh --local so the cloned
repo is used in editable mode and unsloth-zoo tracks main, matching the
behavior of install.sh --local on a VM. install.sh --local is unchanged:
it still sets SKIP_STUDIO_BASE=1, which short-circuits the local_repo
branches in install_python_stack.py, so the overlay is not run twice.

* studio: make --local overlays visible + guard empty arg parsing

- setup.sh: gate the --local flag loop on $# > 0 (defensive against any
  shell that surfaces unset $@ under set -u) and emit a substep when local
  mode is detected so the user can confirm the flag was parsed.
- install_python_stack.py: emit explicit _step lines before each overlay
  pip_install in both local_repo branches so overlays appear in the static
  log instead of being overwritten by the in-place progress bar.
---
 studio/Unsloth_Studio_Colab.ipynb |  2 +-
 studio/install_python_stack.py    | 20 ++++++++++++++++++++
 studio/setup.sh                   | 18 ++++++++++++++++++
 3 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/studio/Unsloth_Studio_Colab.ipynb b/studio/Unsloth_Studio_Colab.ipynb
index 46e2067ba7..c3aec04820 100644
--- a/studio/Unsloth_Studio_Colab.ipynb
+++ b/studio/Unsloth_Studio_Colab.ipynb
@@ -64,7 +64,7 @@
     "id": "27e68f91"
    },
    "outputs": [],
-   "source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh"
+   "source": "!git clone --depth 1 --branch main https://github.com/unslothai/unsloth.git\n%cd /content/unsloth\n!chmod +x studio/setup.sh && ./studio/setup.sh --local"
   },
   {
    "cell_type": "markdown",
diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py
index 8ac59d389d..25d950fd3e 100644
--- a/studio/install_python_stack.py
+++ b/studio/install_python_stack.py
@@ -972,6 +972,7 @@ def install_python_stack() -> int:
             req = REQ_ROOT / "no-torch-runtime.txt",
         )
         if local_repo:
+            _step(_LABEL, f"overlaying local repo (editable): {local_repo}")
             pip_install(
                 "Overlaying local repo (editable)",
                 "--no-cache-dir",
@@ -980,6 +981,15 @@ def install_python_stack() -> int:
                 local_repo,
                 constrain = False,
             )
+            _step(_LABEL, "overlaying unsloth-zoo from git main")
+            pip_install(
+                "Overlaying unsloth-zoo from git main",
+                "--no-cache-dir",
+                "--no-deps",
+                "--force-reinstall",
+                "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo",
+                constrain = False,
+            )
     elif local_repo:
         # Local dev install: update deps from base.txt, then overlay the
         # local checkout as an editable install (--no-deps so torch is
@@ -994,6 +1004,7 @@ def install_python_stack() -> int:
             "unsloth-zoo",
             req = REQ_ROOT / "base.txt",
         )
+        _step(_LABEL, f"overlaying local repo (editable): {local_repo}")
         pip_install(
             "Overlaying local repo (editable)",
             "--no-cache-dir",
@@ -1002,6 +1013,15 @@ def install_python_stack() -> int:
             local_repo,
             constrain = False,
         )
+        _step(_LABEL, "overlaying unsloth-zoo from git main")
+        pip_install(
+            "Overlaying unsloth-zoo from git main",
+            "--no-cache-dir",
+            "--no-deps",
+            "--force-reinstall",
+            "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo",
+            constrain = False,
+        )
     elif package_name != "unsloth":
         # Custom package name (e.g. roland-sloth for testing) — install directly
         _progress("base packages")
diff --git a/studio/setup.sh b/studio/setup.sh
index 142d253554..9575f143d5 100755
--- a/studio/setup.sh
+++ b/studio/setup.sh
@@ -8,6 +8,21 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
 REPO_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
 RULE=$(printf '\342\224\200%.0s' {1..52})
 
+# ── Parse flags ──
+# --local: install from the local repo checkout (overlays unsloth as editable
+# and unsloth-zoo from git main). Mirrors install.sh --local for the Colab
+# path that runs setup.sh directly without going through install.sh.
+if [ "$#" -gt 0 ]; then
+    for _arg in "$@"; do
+        case "$_arg" in
+            --local)
+                export STUDIO_LOCAL_INSTALL=1
+                export STUDIO_LOCAL_REPO="$REPO_ROOT"
+                ;;
+        esac
+    done
+fi
+
 # ── Maintainer-editable defaults ──────────────────────────────────────────
 # Change these in the GitHub-hosted script so all users get updated defaults.
 # User environment variables always override these baked-in values.
@@ -170,6 +185,9 @@ _LLAMA_ONLY="${UNSLOTH_STUDIO_LLAMA_ONLY:-0}"
 if [ "$_LLAMA_ONLY" = "1" ]; then
     substep "llama.cpp only mode"
 fi
+if [ "${STUDIO_LOCAL_INSTALL:-0}" = "1" ]; then
+    substep "local mode: overlaying $REPO_ROOT (editable) + unsloth-zoo from git main"
+fi
 # ── Clean up stale caches ──
 rm -rf "$REPO_ROOT/unsloth_compiled_cache"
 rm -rf "$SCRIPT_DIR/backend/unsloth_compiled_cache"