From 289c7dd7bb20d95d96bc8dd237cf9d3f12293beb Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Wed, 25 Mar 2026 15:12:56 +0000
Subject: [PATCH] Add --local and --package flags to install.ps1

Windows install.ps1 had no way to install from a local repo checkout,
unlike install.sh which supports ./install.sh --local. This adds:

- --local: install from the local repo via editable install (-e . --no-deps)
  after installing deps from PyPI, mirroring install.sh behavior
- --package: install a different package name for testing

The --local flag:
1. Validates pyproject.toml exists at the script's directory
2. Installs torch + unsloth deps normally
3. Overlays the local checkout with uv pip install -e <repo> --no-deps
4. Passes STUDIO_LOCAL_INSTALL and STUDIO_LOCAL_REPO to setup.ps1
---
 install.ps1 | 54 ++++++++++++++++++++++++++++++++++++++++++++++++++---
 1 file changed, 51 insertions(+), 3 deletions(-)

diff --git a/install.ps1 b/install.ps1
index ed656850d4..9405072bd7 100644
--- a/install.ps1
+++ b/install.ps1
@@ -1,10 +1,37 @@
 # Unsloth Studio Installer for Windows PowerShell
 # Usage:  irm https://raw.githubusercontent.com/unslothai/unsloth/main/install.ps1 | iex
-# Local:  Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass; .\install.ps1
+# Local:  Set-ExecutionPolicy -Scope Process -ExecutionPolicy Bypass; .\install.ps1 --local
+# Test:   .\install.ps1 --package roland-sloth
 
 function Install-UnslothStudio {
     $ErrorActionPreference = "Stop"
 
+    # ── Parse flags ──
+    $StudioLocalInstall = $false
+    $PackageName = "unsloth"
+    $RepoRoot = ""
+    $argList = $args
+    for ($i = 0; $i -lt $argList.Count; $i++) {
+        switch ($argList[$i]) {
+            "--local"   { $StudioLocalInstall = $true }
+            "--package" {
+                $i++
+                if ($i -ge $argList.Count) {
+                    Write-Host "[ERROR] --package requires an argument." -ForegroundColor Red
+                    return
+                }
+                $PackageName = $argList[$i]
+            }
+        }
+    }
+    if ($StudioLocalInstall) {
+        $RepoRoot = (Resolve-Path (Split-Path -Parent $PSCommandPath)).Path
+        if (-not (Test-Path (Join-Path $RepoRoot "pyproject.toml"))) {
+            Write-Host "[ERROR] --local must be run from the unsloth repo root (pyproject.toml not found at $RepoRoot)" -ForegroundColor Red
+            return
+        }
+    }
+
     $PythonVersion = "3.13"
     $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio"
     $VenvDir = Join-Path $StudioHome "unsloth_studio"
@@ -581,6 +608,10 @@ shell.Run cmd, 0, False
         # in the new venv location, while preserving existing torch/CUDA
         Write-Host "==> Upgrading unsloth in migrated environment..."
         uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.3.11" unsloth-zoo
+        if ($StudioLocalInstall) {
+            Write-Host "==> Overlaying local repo (editable)..."
+            uv pip install --python $VenvPython -e $RepoRoot --no-deps
+        }
     } elseif ($TorchIndexUrl) {
         Write-Host "==> Installing PyTorch ($TorchIndexUrl)..."
         uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl
@@ -590,11 +621,23 @@ shell.Run cmd, 0, False
         }
 
         Write-Host "==> Installing unsloth (this may take a few minutes)..."
-        uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.11"
+        if ($StudioLocalInstall) {
+            uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.3.11" unsloth-zoo
+            Write-Host "==> Overlaying local repo (editable)..."
+            uv pip install --python $VenvPython -e $RepoRoot --no-deps
+        } else {
+            uv pip install --python $VenvPython --upgrade-package unsloth "$PackageName"
+        }
     } else {
         # Fallback: GPU detection failed to produce a URL -- let uv resolve torch
         Write-Host "==> Installing unsloth (this may take a few minutes)..."
-        uv pip install --python $VenvPython "unsloth>=2026.3.11" --torch-backend=auto
+        if ($StudioLocalInstall) {
+            uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.3.11" --torch-backend=auto
+            Write-Host "==> Overlaying local repo (editable)..."
+            uv pip install --python $VenvPython -e $RepoRoot --no-deps
+        } else {
+            uv pip install --python $VenvPython "$PackageName" --torch-backend=auto
+        }
     }
     if ($LASTEXITCODE -ne 0) {
         Write-Host "[ERROR] Failed to install unsloth (exit code $LASTEXITCODE)" -ForegroundColor Red
@@ -615,6 +658,11 @@ shell.Run cmd, 0, False
     }
     # Tell setup.ps1 to skip base package installation (install.ps1 already did it)
     $env:SKIP_STUDIO_BASE = "1"
+    $env:STUDIO_PACKAGE_NAME = $PackageName
+    if ($StudioLocalInstall) {
+        $env:STUDIO_LOCAL_INSTALL = "1"
+        $env:STUDIO_LOCAL_REPO = $RepoRoot
+    }
     & $UnslothExe studio setup
     if ($LASTEXITCODE -ne 0) {
         Write-Host "[ERROR] unsloth studio setup failed (exit code $LASTEXITCODE)" -ForegroundColor Red