Merge branch 'image-generation' into video-diffusion-improvements

.github/scripts/agent-guides-drive.sh

@@ -527,6 +527,154 @@ case "$MODE" in
     echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)"
     ;;
 
+  # ── resume: does a launched agent's session survive exit and resume? ────
+  # Unlike the other modes, this drives the real LAUNCH path (`unsloth start
+  # <agent> ...`, the interactive default), not the --no-launch recipe. That
+  # path relocates each agent's home to a throwaway temp dir wiped on exit, so
+  # a session cannot be resumed -- unless --persist routes it to the stable
+  # Unsloth agents dir instead. We run one headless turn per pass and check
+  # whether the turn left a session in a persistent store (deterministic, no
+  # reliance on the model recalling anything), for a baseline pass and a
+  # --persist pass, and assert the expected split for this agent.
+  resume)
+    CODEWORD="PLATYPUS7"
+    T1="Remember this codeword for later: ${CODEWORD}. Reply with just the word OK."
+    T2="What codeword did I ask you to remember? Reply with just that word."
+    WORK="$WORKDIR_BASE/${AGENT}-resume"
+
+    # STABLE_HOME: the stable dir that --no-launch (and --persist) relocate to.
+    # Read it from a --no-launch probe (which also writes the agent's config
+    # there). codex/pi relocate their whole home/HOME here; opencode/claude keep
+    # their session data in a fixed user dir, so STABLE_HOME stays empty for them.
+    parse_connect
+    case "$AGENT" in
+      codex)    STABLE_HOME="$(raw_env CODEX_HOME)" ;;
+      pi)       STABLE_HOME="$(raw_env HOME)" ;;
+      *)        STABLE_HOME="" ;;
+    esac
+
+    # The persistent stores a session would land in if it were NOT wiped. We
+    # count files here before/after each turn; a positive delta means the
+    # session persisted (is resumable), zero means it went to a wiped temp dir.
+    resume_tracked_dirs() {
+      case "$AGENT" in
+        codex)    printf '%s\n' "$HOME/.codex" ;;
+        opencode) printf '%s\n' "$HOME/.local/share/opencode" "$HOME/.config/opencode" ;;
+        claude)   printf '%s\n' "$HOME/.claude" ;;
+        pi)       printf '%s\n' "$HOME/.pi" ;;
+        *)        : ;;
+      esac
+      [ -n "$STABLE_HOME" ] && printf '%s\n' "$STABLE_HOME"
+    }
+    count_session_files() {
+      local total=0 d n
+      while IFS= read -r d; do
+        [ -n "$d" ] && [ -d "$d" ] || continue
+        n="$(find "$d" -type f 2>/dev/null | wc -l)"; total=$((total + n))
+      done < <(resume_tracked_dirs)
+      echo "$total"
+    }
+
+    # The headless first-turn subcommand per agent (mirrors file-edit's map),
+    # forwarded verbatim through the launch path as passthrough args.
+    set_t1_cmd() {
+      case "$AGENT" in
+        claude)   T1_CMD=("${CLAUDE_CONNECT_FLAGS[@]}" -p "$T1") ;;
+        codex)    T1_CMD=(exec "$T1") ;;
+        opencode) T1_CMD=(run "$T1") ;;
+        pi)       T1_CMD=(-p "$T1") ;;
+        *)        guide_fail "resume mode does not cover agent '$AGENT'" ;;
+      esac
+    }
+
+    # Run one headless turn through the launch path. $1=outfile, $2="" or
+    # "--persist", rest = the agent subcommand. --yolo auto-approves so no tool
+    # prompt can hang; --api-key attaches to the already-served CI model.
+    launch_turn() {
+      local out="$1" rflag="$2"; shift 2
+      local flag=(); [ -n "$rflag" ] && flag=("$rflag")
+      run_timed "$out" unsloth start "$AGENT" "${flag[@]}" --yolo \
+        --api-key "$UNSLOTH_API_KEY" "$@"
+      local rc=$?
+      redact "$out"
+      return "$rc"
+    }
+
+    # One pass: fresh work dir, one planting turn, set RESULT to PERSISTED/WIPED
+    # from the session-store delta. Runs in the main shell (not a command
+    # substitution) so a hang's guide_fail actually fails the job and the
+    # progress lines reach the CI log. $1 = "" (baseline) or "--persist".
+    RESULT=""
+    run_pass() {
+      local rflag="$1" label="baseline"
+      [ -n "$rflag" ] && label="resume"
+      rm -rf "$WORK"; mkdir -p "$WORK"
+      set_t1_cmd
+      local out="$LOGS_DIR/${AGENT}-resume-${label}.txt"
+      local before after rc
+      before="$(count_session_files)"
+      pushd "$WORK" >/dev/null || guide_fail "could not enter work dir $WORK"
+      launch_turn "$out" "$rflag" "${T1_CMD[@]}"; rc=$?
+      popd >/dev/null || true
+      after="$(count_session_files)"
+      echo "[$AGENT] ${label}: session files ${before} -> ${after} (rc=${rc})"
+      # The turn must succeed for the delta to mean anything: an agent that writes a
+      # session file then errors would otherwise be misread as PERSISTED. Mirror the
+      # file-edit mode and fail the pass on a non-zero launch (the flagship codex recall
+      # below stays WARN-only, driven by its own launch_turn calls).
+      [ "$rc" -eq 0 ] || { echo "[$AGENT] ${label} transcript (tail):"; tail -30 "$out" 2>/dev/null || true; \
+        guide_fail "resume ${label} turn for ${AGENT} exited non-zero (rc=${rc})"; }
+      if [ "$after" -gt "$before" ]; then RESULT="PERSISTED"; else RESULT="WIPED"; fi
+    }
+
+    run_pass ""; BASELINE="$RESULT"
+    # Only the temp-dir agents (codex/pi) need the --persist pass to prove the fix.
+    # opencode/claude persist either way, so the baseline already proves it and a
+    # second full CPU turn only risks a timeout; skip it for them.
+    case "$AGENT" in
+      codex|pi) run_pass "--persist"; RESUME="$RESULT" ;;
+      *)        RESUME="n/a (persists either way)" ;;
+    esac
+
+    # Expected: codex/pi relocate their whole home to the temp dir, so a plain
+    # launch is WIPED and only --persist PERSISTS. opencode/claude keep their
+    # session data in a fixed user dir, so the baseline already PERSISTS.
+    case "$AGENT" in
+      codex|pi)        EXPECT_BASELINE="WIPED" ;;
+      opencode|claude) EXPECT_BASELINE="PERSISTED" ;;
+    esac
+
+    echo "──────────────────────────────────────────────"
+    echo "[$AGENT] RESUME EXPERIMENT"
+    echo "  baseline (unsloth start ${AGENT}):                 ${BASELINE}  (expected ${EXPECT_BASELINE})"
+    echo "  with --persist (unsloth start ${AGENT} --persist): ${RESUME}"
+    echo "──────────────────────────────────────────────"
+
+    [ "$BASELINE" = "$EXPECT_BASELINE" ] \
+      || guide_fail "baseline resume behavior for ${AGENT} was ${BASELINE}, expected ${EXPECT_BASELINE}"
+    case "$AGENT" in
+      codex|pi)
+        [ "$RESUME" = "PERSISTED" ] \
+          || guide_fail "--persist did not persist ${AGENT}'s session (got ${RESUME}); the session dir is still not stable" ;;
+    esac
+
+    # Flagship behavioral proof (codex only, WARN-only): after a --persist plant,
+    # resume the session and check the model actually recalls the codeword. A
+    # miss is not a failure (the CI model is small); the mechanism gate above is
+    # the real assertion.
+    if [ "$AGENT" = "codex" ]; then
+      rm -rf "$WORK"; mkdir -p "$WORK"
+      ( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-plant.txt" "--persist" exec "$T1" ) || true
+      ( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-recall.txt" "--persist" exec resume --last "$T2" ) || true
+      if grep -q "$CODEWORD" "$LOGS_DIR/codex-resume-recall.txt" 2>/dev/null; then
+        echo "[codex] behavioral recall HIT: resumed session remembered ${CODEWORD}"
+      else
+        echo "::warning::[codex] behavioral recall MISS (small CI model); mechanism gate still passed"
+      fi
+    fi
+    echo "[$AGENT] resume OK"
+    ;;
+
   *)
     echo "agent-guides-drive.sh: unknown mode '$MODE'" >&2
     exit 2

.github/workflows/local-agent-guides-ci.yml

@@ -471,6 +471,176 @@ jobs:
             redacted-configs/
           retention-days: 7
 
+  # ═════════════════════════════════════════════════════════════════════
+  # Job: resume
+  #   Does a conversation started with `unsloth start <agent>` survive exit
+  #   and resume? This drives the REAL launch path (not the --no-launch
+  #   recipe the other jobs use). A plain launch relocates the agent home to
+  #   a temp dir wiped on exit, so codex/pi cannot resume; --persist routes the
+  #   session to the stable Unsloth agents dir so it persists. opencode/claude
+  #   keep their session data in a fixed user dir, so they persist either way.
+  #   Dispatch-only: it is an end-to-end experiment, not a PR gate.
+  # ═════════════════════════════════════════════════════════════════════
+  resume:
+    name: resume (${{ matrix.agent }})
+    if: github.event_name == 'workflow_dispatch'
+    runs-on: ubuntu-latest
+    timeout-minutes: 60
+    strategy:
+      fail-fast: false
+      matrix:
+        # codex/pi relocate their whole home (resume broken without --persist);
+        # opencode/claude keep session data in a fixed dir (resume already works).
+        # One agent from each class proves the split end to end; openclaw/hermes
+        # share codex's relocation mechanism and are covered by the unit tests.
+        agent: [codex, opencode, claude, pi]
+    env:
+      GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF
+      GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf
+      STUDIO_PORT: '18904'
+    steps:
+      - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd  # v6.0.2
+        with:
+          persist-credentials: false
+
+      - name: Linux deps for llama.cpp prebuilt
+        run: |
+          sudo apt-get update
+          sudo apt-get install -y --no-install-recommends \
+            libcurl4-openssl-dev libssl-dev jq
+
+      - uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e  # v6.4.0
+        with:
+          node-version: '22'
+
+      - uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405  # v6.2.0
+        with:
+          python-version: '3.12'
+          cache: 'pip'
+
+      - name: Restore GGUF model file
+        id: cache-gguf
+        uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae  # v5.0.5
+        continue-on-error: true
+        with:
+          path: gguf-cache
+          key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
+
+      - name: Download GGUF if cache miss
+        id: download-gguf
+        if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success'
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: |
+          python -m pip install --upgrade huggingface_hub
+          mkdir -p gguf-cache
+          bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache
+
+      - name: Save GGUF model file
+        if: always() && steps.download-gguf.outcome == 'success'
+        uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae  # v5.0.5
+        with:
+          path: gguf-cache
+          key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
+
+      - name: Install Studio (--local, --no-torch)
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
+        run: |
+          mkdir -p logs
+          set -o pipefail
+          bash install.sh --local --no-torch 2>&1 | tee logs/install.log
+
+      - name: Serve unsloth run --disable-tools (gemma-4-E4B)
+        run: |
+          unsloth studio reset-password
+          bash .github/scripts/serve-unsloth-run.sh \
+            --gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \
+            --port "$STUDIO_PORT" --log-dir logs \
+            --extra "--seed $UNSLOTH_SEED --temp 0" \
+            --health-timeout 900
+
+      - name: Preflight the agent's API dialect (class-a isolation)
+        env:
+          AGENT: ${{ matrix.agent }}
+        run: |
+          set -uo pipefail
+          B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY"
+          preflight_fail() {
+            echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect). Endpoint contract lives in studio/backend/routes/**.";
+            exit 1
+          }
+          code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \
+            -H "Authorization: Bearer $K") || true
+          [ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code"
+          case "$AGENT" in
+            claude)
+              code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \
+                -H "Authorization: Bearer $K" -H 'content-type: application/json' \
+                --max-time 120 \
+                -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true
+              [ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code"
+              ;;
+            codex)
+              code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \
+                -H "Authorization: Bearer $K" -H 'content-type: application/json' \
+                --max-time 120 \
+                -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true
+              [ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code"
+              ;;
+            *)
+              code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \
+                -H "Authorization: Bearer $K" -H 'content-type: application/json' \
+                --max-time 120 \
+                -d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true
+              [ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code"
+              ;;
+          esac
+          echo "preflight OK for $AGENT"
+
+      - name: Install agent CLI (class-b isolation)
+        env:
+          AGENT: ${{ matrix.agent }}
+        run: bash .github/scripts/agent-guides-install.sh "$AGENT"
+
+      - name: Resume experiment (launch path)
+        env:
+          AGENT: ${{ matrix.agent }}
+        run: bash .github/scripts/agent-guides-drive.sh resume "$AGENT"
+
+      - name: Collect server logs (debug)
+        if: always()
+        run: |
+          mkdir -p logs/studio-logs
+          cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true
+          if [ -n "${UNSLOTH_API_KEY:-}" ]; then
+            grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do
+              sed -i "s#${UNSLOTH_API_KEY}#<REDACTED>#g" "$f" 2>/dev/null || true
+            done
+          fi
+
+      - name: Stop Studio
+        if: always()
+        run: |
+          if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then
+            kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true
+          fi
+          sleep 2
+          ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true
+
+      - name: Upload logs
+        if: always()
+        continue-on-error: true
+        uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a  # v7.0.1
+        with:
+          name: resume-${{ matrix.agent }}-log
+          path: |
+            logs/
+            agent-workdir/
+            redacted-configs/
+          retention-days: 7
+
   # ═════════════════════════════════════════════════════════════════════
   # Job 3: prompt-cache
   #   (a) curl 2-turn /v1/chat/completions: assert turn-2 cached_tokens > 0

.github/workflows/security-audit.yml

@@ -2,8 +2,8 @@
 # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
 
 # Multi-language supply-chain audit. Triggers:
-#   - PRs touching any dependency manifest (Python / npm / Cargo) or
-#     this workflow file,
+#   - PRs touching any dependency manifest (Python / npm / Cargo), a
+#     scanner or its allowlist baseline, or this workflow file,
 #   - push to main / pip,
 #   - nightly @ 04:13 UTC so newly-published advisories surface even
 #     when no PR opens,
@@ -57,7 +57,9 @@ on:
       - 'studio/src-tauri/Cargo.lock'
       - 'pyproject.toml'
       - 'scripts/scan_packages.py'
+      - 'scripts/scan_packages_baseline.json'
       - 'scripts/scan_npm_packages.py'
+      - 'scripts/scan_npm_packages_baseline.json'
       - '.github/workflows/security-audit.yml'
   push:
     branches: [main, pip]

README.md

@@ -212,7 +212,7 @@ By default `unsloth studio` binds to `127.0.0.1` (this machine only). To reach i
 ```bash
 unsloth studio --secure -p 8888
 ```
-- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. Only use this on a trusted network.
+- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. This also starts a public Cloudflare quick tunnel by default, which publishes an internet-reachable `https://*.trycloudflare.com` URL even behind a firewall. Both the raw port and the tunnel expose Studio beyond this machine, so only use this on a network you trust; pass `--no-cloudflare` to drop the public link while keeping the network bind.
 ```bash
 unsloth studio -H 0.0.0.0 -p 8888
 ```

install.ps1

@@ -469,6 +469,17 @@ function Install-UnslothStudio {
         param(
             [Parameter(Mandatory = $true)][ScriptBlock]$Command
         )
+        # Installer-pinned index installs (torch) must beat an inherited uv mirror
+        # (#6898): when the command pins an index, clear every uv index env var so
+        # it wins, then restore in finally. Other installs keep the user's mirror.
+        $savedUvIndex = $null
+        if ($Command.ToString() -match '--default-index') {
+            $savedUvIndex = @{}
+            foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
+                $savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n)
+                Remove-Item "Env:$n" -ErrorAction SilentlyContinue
+            }
+        }
         $prevEap = $ErrorActionPreference
         $ErrorActionPreference = "Continue"
         try {
@@ -488,6 +499,7 @@ function Install-UnslothStudio {
             return [int]$LASTEXITCODE
         } finally {
             $ErrorActionPreference = $prevEap
+            if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } }
         }
     }
 
@@ -2200,7 +2212,7 @@ exit 0
             # ABI-incompatible torchvision/torchaudio on AMD's per-arch index.
             $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
             $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
-            $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
+            $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
             if ($torchInstallExit -ne 0) {
                 # Transient AMD-index failure: fall back to a CPU base so the install
                 # still completes; Studio setup retries ROCm afterwards.
@@ -2209,7 +2221,7 @@ exit 0
                 # torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU
                 # torch>= range, so without it uv would keep the ROCm build and only swap
                 # the companions -- a mismatched venv the flavor-repair block won't fix.
-                $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
+                $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
                 if ($torchInstallExit -ne 0) {
                     Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red
                     return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@@ -2223,7 +2235,7 @@ exit 0
         } else {
             Write-TauriLog "STEP" "Installing PyTorch"
             substep "installing PyTorch ($TorchIndexUrl)..."
-            $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
+            $torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
             if ($torchInstallExit -ne 0) {
                 Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
                 return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@@ -2306,7 +2318,7 @@ exit 0
     # keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on
     # "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is
     # expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx*
-    # is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install
+    # is a PEP 503 index uv resolves via --default-index, same URL the fresh ROCm install
     # above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops.
     if (-not $SkipTorch) {
         $expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl
@@ -2322,7 +2334,7 @@ exit 0
                     $visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
                     $audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
                     substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
-                    $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
+                    $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
                     if ($torchFixExit -ne 0) {
                         Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red
                         return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit)
@@ -2331,7 +2343,7 @@ exit 0
                 } elseif ($expectedTorchTag -ne 'rocm') {
                     # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet.
                     substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow"
-                    $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
+                    $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
                     if ($torchFixExit -ne 0) {
                         Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red
                         return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit)

install.sh

@@ -159,6 +159,12 @@ run_maybe_quiet() {
 run_install_cmd() {
     _label="$1"
     shift
+    # Installer-pinned index installs (torch) must beat an inherited uv mirror
+    # (#6898): when we pass --default-index, neutralize every uv index env var so
+    # the pinned index wins. Other installs keep the user's mirror.
+    case " $* " in
+        *" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;;
+    esac
     if _is_verbose; then
         "$@" && return 0
         _rc=$?
@@ -2190,9 +2196,9 @@ _expected_torch_flavor_tag() {
     esac
 }
 
-# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX /
+# Whether index ($1) supports a plain --default-index reinstall. pytorch.org cuXXX /
 # rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv
-# resolves (torch + every transitive dep) via --index-url -- the same URLs the
+# resolves (torch + every transitive dep) via --default-index -- the same URLs the
 # fresh-install paths above already use -- so a stale wheel is auto-repairable.
 # Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall.
 _torch_index_repairable() {
@@ -2744,7 +2750,7 @@ if [ "$_MIGRATED" = true ]; then
                     substep "repairing ROCm torch (overwritten by dependency resolution)..."
                     run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
                         "$TORCH_CONSTRAINT" torchvision torchaudio \
-                        --index-url "$TORCH_INDEX_URL" \
+                        --default-index "$TORCH_INDEX_URL" \
                         --force-reinstall
                 fi
                 ;;
@@ -2870,7 +2876,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
                     substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
                     run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
                         "$TORCH_CONSTRAINT" torchvision torchaudio \
-                        --index-url "$TORCH_INDEX_URL"
+                        --default-index "$TORCH_INDEX_URL"
                 else
                     substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..."
                     # Pass explicit wheel URLs so the matched trio is
@@ -2893,18 +2899,18 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
                 substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
                 run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
                     "$TORCH_CONSTRAINT" torchvision torchaudio \
-                    --index-url "$TORCH_INDEX_URL"
+                    --default-index "$TORCH_INDEX_URL"
             fi
         else
             substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN"
             run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
                 "$TORCH_CONSTRAINT" torchvision torchaudio \
-                --index-url "$TORCH_INDEX_URL"
+                --default-index "$TORCH_INDEX_URL"
         fi
     else
         substep "installing PyTorch ($TORCH_INDEX_URL)..."
         run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
-            --index-url "$TORCH_INDEX_URL"
+            --default-index "$TORCH_INDEX_URL"
     fi
     # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths).
     # Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm
@@ -2964,7 +2970,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
                     substep "repairing ROCm torch (overwritten by dependency resolution)..."
                     run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
                         "$TORCH_CONSTRAINT" torchvision torchaudio \
-                        --index-url "$TORCH_INDEX_URL" \
+                        --default-index "$TORCH_INDEX_URL" \
                         --force-reinstall
                 fi
                 ;;
@@ -2999,14 +3005,14 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
         _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
         _installed_torch_tag=""
         [ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver")
-        # Repair when flavor is wrong AND the index is plain --index-url reinstallable
+        # Repair when flavor is wrong AND the index is plain --default-index reinstallable
         # (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only.
         if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \
            && [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then
             substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..."
             run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \
                 "$TORCH_CONSTRAINT" torchvision torchaudio \
-                --index-url "$TORCH_INDEX_URL" \
+                --default-index "$TORCH_INDEX_URL" \
                 --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio
             _installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
             _installed_torch_tag=""
@@ -3017,7 +3023,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
             substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN"
             substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN"
             substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN"
-            substep "[WARN]   uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
+            substep "[WARN]   uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
         fi
     fi
 fi

scripts/image_speedmem_bench.py

@@ -0,0 +1,429 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Speed + accuracy lever benchmark for the IMAGE diffusion backend (per-lever LPIPS).
+
+Drives the SAME production lever functions the image loader calls -- ``apply_step_cache``,
+``apply_attention_backend``, ``apply_speed_optims``, ``quantize_text_encoders``, the
+compile-safe eager patches -- with the loader's own default arguments and order, so each
+measured configuration reflects a real load. For each config it loads the pipeline fresh
+(quant/compile mutate irreversibly), warms up (to pay the one-time compile), renders a
+fixed prompt set at a fixed seed, and reports total latency, median per-step ms, peak
+resident GB, and mean LPIPS(AlexNet) vs the bit-exact reference config (speed off,
+native attention, uncached, dense) rendered at the same seed/settings.
+
+Lever isolation knobs (for before/after measurement of shipped fixes):
+  --no-epc         force torch._inductor.config.emulate_precision_casts back off after
+                   the speed layer enables it (the pre-fix compile numerics).
+  --unarm-cache    restore the cache hooks' eager inner forwards after the speed layer
+                   arms them (the pre-fix cache x compile composition).
+
+Example:
+    CUDA_VISIBLE_DEVICES=3 python scripts/image_speedmem_bench.py --family flux.1-dev \\
+        --config compile --out outputs/image_speedmem
+"""
+
+from __future__ import annotations
+
+import argparse
+import gc
+import json
+import os
+import sys
+import time
+import types
+from pathlib import Path
+from typing import Any, Optional
+
+os.environ.setdefault("BITSANDBYTES_NOWELCOME", "1")
+
+_REPO_ROOT = Path(__file__).resolve().parent.parent
+_BACKEND_ROOT = _REPO_ROOT / "studio" / "backend"
+for _p in (str(_BACKEND_ROOT), str(_REPO_ROOT / "scripts")):
+    if _p not in sys.path:
+        sys.path.insert(0, _p)
+
+# Fixed prompt set (the diffusion_quality.py defaults + one photographic subject) so the
+# LPIPS mean is not hostage to a single composition.
+PROMPTS = [
+    "A cozy reading nook by a rain-streaked window, warm lamplight, a cat asleep on a stack of books",
+    "A lone lighthouse on a rocky cliff at sunset, dramatic clouds, crashing waves, highly detailed",
+    "A bustling night market street in the rain, neon signs reflected in puddles, cinematic",
+    "A photograph of an astronaut riding a horse on the surface of the moon, detailed, 8k",
+]
+
+# Production defaults per family (diffusion_families.default_generation_params).
+_FAMILIES: dict[str, dict[str, Any]] = {
+    "qwen-image": {"repo": "Qwen/Qwen-Image", "family": "qwen-image"},
+    "flux.1-dev": {"repo": "black-forest-labs/FLUX.1-dev", "family": "flux.1"},
+    "flux.2-klein-4b": {"repo": "black-forest-labs/FLUX.2-klein-4B", "family": "flux.2-klein"},
+    "sdxl": {"repo": "stabilityai/stable-diffusion-xl-base-1.0", "family": "sdxl"},
+}
+
+#                       te        speed       attn       cache
+_CONFIGS: dict[str, dict[str, Any]] = {
+    # bit-exact reference: everything off / native / dense.
+    "reference": dict(te = "none", speed = "off", attn = "native", cache = "off"),
+    # the non-compile floor: eager patches + attention auto-upgrade, no compile.
+    "eager": dict(te = "none", speed = "eager", attn = "auto", cache = "off"),
+    # the default dense tier (regional compile), uncached.
+    "compile": dict(te = "none", speed = "default", attn = "auto", cache = "off"),
+    # max tier (max-autotune regional compile + TF32 + fused QKV), uncached.
+    "speedmax": dict(te = "none", speed = "max", attn = "auto", cache = "off"),
+    # the default tier + FBCache (the auto path for 20+ step schedules).
+    "fbcache": dict(te = "none", speed = "default", attn = "auto", cache = "fbcache"),
+    # FBCache without compile (isolates the cache's own drift from the compile floor).
+    "fbcache_eager": dict(te = "none", speed = "eager", attn = "auto", cache = "fbcache"),
+    # TE quant isolation on the bit-exact stack: the conditioning perturbation ALONE.
+    "te_fp8dyn": dict(te = "fp8_dynamic", speed = "off", attn = "native", cache = "off"),
+    "te_fp8": dict(te = "fp8", speed = "off", attn = "native", cache = "off"),
+}
+
+
+def _sync() -> None:
+    import torch
+    if torch.cuda.is_available():
+        torch.cuda.synchronize()
+
+
+def _reset_peak() -> None:
+    import torch
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats()
+
+
+def _alloc_gb() -> float:
+    import torch
+    return torch.cuda.memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0
+
+
+def _peak_gb() -> float:
+    import torch
+    return torch.cuda.max_memory_allocated() / 1e9 if torch.cuda.is_available() else 0.0
+
+
+def _empty() -> None:
+    import torch
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+
+
+_LP: dict = {}
+
+
+def _lpips_alex(ref_arr, arr) -> Optional[float]:
+    """LPIPS(AlexNet) between two HxWx3 uint8 images (net on CPU). None if lpips missing."""
+    try:
+        import lpips
+        import torch
+
+        fn = _LP.get("fn")
+        if fn is None:
+            fn = lpips.LPIPS(net = "alex", verbose = False).eval()
+            _LP["fn"] = fn
+
+        def _t(a):
+            import torch as _torch
+            return _torch.from_numpy(a).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0
+
+        with torch.no_grad():
+            return float(fn(_t(ref_arr), _t(arr)).item())
+    except Exception:
+        return None
+
+
+def _import_diffusers():
+    import torch  # noqa: F401
+    import torchao  # noqa: F401
+    import diffusers.utils.import_utils as iu
+
+    iu._bitsandbytes_available = False
+    import diffusers
+
+    return diffusers
+
+
+def _target():
+    """Stand-in for DiffusionDeviceTarget: what the real lever functions read."""
+    import torch
+    return types.SimpleNamespace(
+        device = "cuda",
+        dtype = torch.bfloat16,
+        supports_default_torch_compile = True,
+    )
+
+
+def _find_family(name: str):
+    from core.inference.diffusion_families import _FAMILIES as ALL
+    for fam in ALL:
+        if fam.name == name:
+            return fam
+    raise SystemExit(f"unknown family '{name}'")
+
+
+def _apply_levers(
+    pipe,
+    cfg: dict,
+    *,
+    fam_obj,
+    no_epc: bool = False,
+    unarm_cache: bool = False,
+    logger = None,
+) -> dict:
+    """Apply the configured levers with the loader's own argument values, in the loader's
+    order (diffusion.py): TE quant -> attention -> step cache -> eager patches -> speed."""
+    from core.inference.diffusion_precision import quantize_text_encoders
+    from core.inference.diffusion_attention import (
+        apply_attention_backend,
+        select_attention_backend,
+    )
+    from core.inference.diffusion_cache import apply_step_cache, _restore_hooked_block_inners
+    from core.inference.diffusion_eager_patches import (
+        install_compile_safe_patches,
+        uninstall_patches,
+    )
+    from core.inference.diffusion_arch_patches import (
+        install_arch_patches,
+        uninstall_arch_patches,
+    )
+    from core.inference.diffusion_speed import apply_speed_optims
+
+    tgt = _target()
+    engaged: dict[str, Any] = {"te": None, "attn": None, "cache": None, "speed_optims": {}}
+
+    if cfg["te"] != "none":
+        engaged["te"] = quantize_text_encoders(
+            pipe, tgt, mode = cfg["te"], family = fam_obj.name, logger = logger
+        )
+
+    speed_mode = cfg["speed"]
+    engaged["attn"] = apply_attention_backend(
+        pipe,
+        select_attention_backend(
+            tgt, None if cfg["attn"] == "auto" else cfg["attn"], speed_active = speed_mode != "off"
+        ),
+        logger = logger,
+    )
+
+    if cfg["cache"] != "off":
+        engaged["cache"] = apply_step_cache(
+            pipe, mode = cfg["cache"], quant_active = False, logger = logger
+        )
+
+    if speed_mode != "off":
+        install_compile_safe_patches()
+        install_arch_patches()
+    else:
+        uninstall_patches()
+        uninstall_arch_patches()
+
+    engaged["speed_optims"] = apply_speed_optims(
+        pipe,
+        tgt,
+        is_gguf = False,
+        family = fam_obj,
+        speed_mode = speed_mode,
+        cache_active = engaged["cache"] is not None,
+        offload_active = False,
+    )
+
+    if no_epc:
+        import torch
+        cfg_ind = getattr(getattr(torch, "_inductor", None), "config", None)
+        if cfg_ind is not None and hasattr(cfg_ind, "emulate_precision_casts"):
+            cfg_ind.emulate_precision_casts = False
+            engaged["epc_forced_off"] = True
+    if unarm_cache:
+        transformer = getattr(pipe, "transformer", None)
+        if transformer is not None:
+            _restore_hooked_block_inners(transformer)
+            engaged["cache_unarmed"] = True
+    return engaged
+
+
+def _generate(
+    pipe,
+    fam_obj,
+    *,
+    steps: int,
+    guidance: float,
+    size: int,
+    seed: int,
+    limit: Optional[int] = None,
+) -> tuple:
+    """Render every prompt at a fixed per-prompt seed; returns (arrays, total_s, step_ms)."""
+    import numpy as np
+    import torch
+
+    call_params = {}
+    try:
+        import inspect
+        call_params = inspect.signature(pipe.__call__).parameters
+    except (TypeError, ValueError):
+        pass
+
+    step_times: list[float] = []
+    last: dict[str, float] = {}
+
+    def _cb(p, i, t, kw):
+        now = time.perf_counter()
+        if "t" in last:
+            step_times.append(now - last["t"])
+        last["t"] = now
+        return kw
+
+    arrs = []
+    total = 0.0
+    for idx, prompt in enumerate(PROMPTS[: limit or len(PROMPTS)]):
+        kwargs: dict[str, Any] = {
+            "prompt": prompt,
+            "num_inference_steps": steps,
+            "width": size,
+            "height": size,
+            "generator": torch.Generator("cuda").manual_seed(seed + idx),
+        }
+        if fam_obj.cfg_kwarg in call_params:
+            kwargs[fam_obj.cfg_kwarg] = guidance
+        if "callback_on_step_end" in call_params:
+            kwargs["callback_on_step_end"] = _cb
+        last.clear()
+        _sync()
+        t0 = time.perf_counter()
+        with torch.inference_mode():
+            image = pipe(**kwargs).images[0]
+        _sync()
+        total += time.perf_counter() - t0
+        arrs.append(np.array(image.convert("RGB")))
+    med_step = sorted(step_times)[len(step_times) // 2] * 1000.0 if step_times else None
+    return arrs, total, med_step, step_times
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description = __doc__.splitlines()[0])
+    ap.add_argument("--family", required = True, choices = sorted(_FAMILIES))
+    ap.add_argument("--config", required = True, choices = sorted(_CONFIGS))
+    ap.add_argument("--steps", type = int, default = None, help = "override the family default")
+    ap.add_argument("--size", type = int, default = 1024)
+    ap.add_argument("--seed", type = int, default = 42)
+    ap.add_argument("--out", default = "outputs/image_speedmem")
+    ap.add_argument("--no-epc", action = "store_true")
+    ap.add_argument("--unarm-cache", action = "store_true")
+    ap.add_argument("--tag", default = None, help = "output row name (default: config name)")
+    args = ap.parse_args()
+
+    import logging
+
+    logging.basicConfig(level = logging.INFO, format = "%(levelname)s %(name)s: %(message)s")
+    logger = logging.getLogger("image_speedmem")
+
+    fam_spec = _FAMILIES[args.family]
+    cfg = _CONFIGS[args.config]
+    tag = args.tag or args.config
+
+    import numpy as np
+    import torch
+
+    diffusers = _import_diffusers()
+    from core.inference.diffusion_families import default_generation_params
+
+    fam_obj = _find_family(fam_spec["family"])
+    steps, guidance = default_generation_params(fam_spec["repo"])
+    if args.steps is not None:
+        steps = args.steps
+
+    out_dir = Path(args.out) / args.family
+    out_dir.mkdir(parents = True, exist_ok = True)
+    ref_npz = out_dir / f"ref_seed{args.seed}_st{steps}_{args.size}.npz"
+
+    logger.info(
+        "family=%s config=%s steps=%d guidance=%s size=%d seed=%d",
+        args.family,
+        args.config,
+        steps,
+        guidance,
+        args.size,
+        args.seed,
+    )
+
+    _reset_peak()
+    t0 = time.perf_counter()
+    pipe = diffusers.DiffusionPipeline.from_pretrained(fam_spec["repo"], torch_dtype = torch.bfloat16)
+    load_s = time.perf_counter() - t0
+
+    engaged = _apply_levers(
+        pipe,
+        cfg,
+        fam_obj = fam_obj,
+        no_epc = args.no_epc,
+        unarm_cache = args.unarm_cache,
+        logger = logger,
+    )
+    pipe.to("cuda")
+    weights_gb = _alloc_gb()
+
+    # Warmup: pays the one-time compile (and the cuDNN autotune) outside the timed runs.
+    wt0 = time.perf_counter()
+    _generate(
+        pipe,
+        fam_obj,
+        steps = steps,
+        guidance = guidance,
+        size = args.size,
+        seed = args.seed + 1000,
+        limit = 1,
+    )
+    warmup_s = time.perf_counter() - wt0
+
+    _reset_peak()
+    arrs, total_s, med_step_ms, step_times = _generate(
+        pipe, fam_obj, steps = steps, guidance = guidance, size = args.size, seed = args.seed
+    )
+    gen_peak = _peak_gb()
+
+    # Persist / score against the reference.
+    lpips_vals: list[float] = []
+    if args.config == "reference" and not (args.no_epc or args.unarm_cache):
+        np.savez_compressed(ref_npz, *arrs)
+    if ref_npz.exists():
+        ref = np.load(ref_npz)
+        refs = [ref[k] for k in ref.files]
+        for r, a in zip(refs, arrs):
+            v = _lpips_alex(r, a)
+            if v is not None:
+                lpips_vals.append(v)
+
+    from PIL import Image
+
+    for i, a in enumerate(arrs):
+        Image.fromarray(a).save(out_dir / f"{tag}_p{i}.png")
+
+    row = {
+        "family": args.family,
+        "config": args.config,
+        "tag": tag,
+        "steps": steps,
+        "guidance": guidance,
+        "size": args.size,
+        "seed": args.seed,
+        "engaged": {k: v for k, v in engaged.items()},
+        "load_s": round(load_s, 2),
+        "warmup_s": round(warmup_s, 2),
+        "total_gen_s": round(total_s, 2),
+        "per_image_s": round(total_s / len(PROMPTS), 3),
+        "median_step_ms": round(med_step_ms, 1) if med_step_ms else None,
+        "step_times_s": [round(t, 4) for t in step_times],
+        "weights_gb": round(weights_gb, 2),
+        "gen_peak_gb": round(gen_peak, 2),
+        "lpips_vs_ref_mean": round(sum(lpips_vals) / len(lpips_vals), 4) if lpips_vals else None,
+        "lpips_vs_ref_per_prompt": [round(v, 4) for v in lpips_vals] or None,
+    }
+    (out_dir / f"{tag}.json").write_text(json.dumps(row, indent = 2, default = str))
+    print(json.dumps(row, indent = 2, default = str))
+
+    del pipe
+    _empty()
+
+
+if __name__ == "__main__":
+    main()

studio/backend/core/inference/_vulkan_probe.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Standalone free-VRAM probe for the bundled ggml Vulkan backend.
+
+Run in a short-lived subprocess (``python _vulkan_probe.py <bindir>``) so the
+Vulkan instance never lives in the long-running backend process. Loads the
+bundled ggml Vulkan backend from ``<bindir>`` and prints one
+``<idx>\\t<free_bytes>\\t<is_igpu>\\t<total_bytes>`` line per device to stdout.
+Indices are ggml's own Vulkan device ordinals, which need not match nvidia-smi
+order. ``is_igpu`` (from ggml's device type) is ``1`` for an integrated GPU
+sharing system RAM. ``total_bytes`` is the device-local heap; the reader uses
+it to reserve absolute headroom on a discrete card (parity with the CUDA/ROCm
+fit) and ignores it for an iGPU, whose "VRAM" is shared system RAM.
+
+Uses only the standard library so it stays runnable as a bare script.
+"""
+
+import ctypes
+import os
+import sys
+
+# ggml_backend_dev_type enum (ggml-backend.h): CPU=0, GPU=1, IGPU=2, ...
+_GGML_BACKEND_DEVICE_TYPE_IGPU = 2
+
+
+def _igpu_flags(base, lib, count: int) -> list[bool]:
+    """Per-device integrated-GPU flags via ggml's backend registry.
+
+    The Vulkan reg enumerates devices in the same order as
+    ``ggml_backend_vk_get_device_memory`` (each context uses ``ctx->device =
+    i``), so reg index == device ordinal. Returns all-False on any failure so
+    the reader never over-caps a discrete card.
+    """
+    flags = [False] * count
+    try:
+        lib.ggml_backend_vk_reg.restype = ctypes.c_void_p
+        lib.ggml_backend_vk_reg.argtypes = []
+        base.ggml_backend_reg_dev_count.restype = ctypes.c_size_t
+        base.ggml_backend_reg_dev_count.argtypes = [ctypes.c_void_p]
+        base.ggml_backend_reg_dev_get.restype = ctypes.c_void_p
+        base.ggml_backend_reg_dev_get.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
+        base.ggml_backend_dev_type.restype = ctypes.c_int
+        base.ggml_backend_dev_type.argtypes = [ctypes.c_void_p]
+
+        reg = lib.ggml_backend_vk_reg()
+        if not reg:
+            return flags
+        dev_count = base.ggml_backend_reg_dev_count(reg)
+        for i in range(min(count, dev_count)):
+            dev = base.ggml_backend_reg_dev_get(reg, i)
+            if dev:
+                flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU
+    except Exception:
+        # Best-effort: any failure degrades to "discrete" so the memory
+        # readings still get through instead of crashing the probe.
+        pass
+    return flags
+
+
+def main() -> int:
+    if len(sys.argv) < 2:
+        return 0
+    bindir = sys.argv[1]
+
+    # Hold add_dll_directory's handle for the rest of main() (the documented
+    # idiom) so bindir stays on the search path while the sibling ggml DLLs
+    # resolve below.
+    _dll_dir = None
+    if sys.platform == "win32":
+        base_name, vk_name = "ggml-base.dll", "ggml-vulkan.dll"
+        try:
+            _dll_dir = os.add_dll_directory(bindir)
+        except Exception:
+            pass
+    else:
+        base_name, vk_name = "libggml-base.so", "libggml-vulkan.so"
+
+    # RTLD_GLOBAL exposes ggml-base's symbols to ggml-vulkan on POSIX. getattr
+    # falls back to 0 where the flag doesn't exist (Windows CDLL ignores mode).
+    _rtld_global = getattr(ctypes, "RTLD_GLOBAL", 0)
+    try:
+        base = ctypes.CDLL(os.path.join(bindir, base_name), mode = _rtld_global)
+        lib = ctypes.CDLL(os.path.join(bindir, vk_name), mode = _rtld_global)
+    except OSError as e:
+        print(f"ggml-vulkan load failed: {e}", file = sys.stderr)
+        return 1
+
+    lib.ggml_backend_vk_get_device_count.restype = ctypes.c_int
+    lib.ggml_backend_vk_get_device_count.argtypes = []
+    lib.ggml_backend_vk_get_device_memory.restype = None
+    lib.ggml_backend_vk_get_device_memory.argtypes = [
+        ctypes.c_int,
+        ctypes.POINTER(ctypes.c_size_t),
+        ctypes.POINTER(ctypes.c_size_t),
+    ]
+
+    count = lib.ggml_backend_vk_get_device_count()
+    igpu = _igpu_flags(base, lib, count)
+    rows = []
+    for i in range(count):
+        free, total = ctypes.c_size_t(0), ctypes.c_size_t(0)
+        lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total))
+        rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value))
+    sys.stdout.write("\n".join(rows))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

studio/backend/core/inference/diffusion_cache.py

@@ -471,17 +471,23 @@ def _compile_hooked_block_inners(transformer: Any, logger: Any = None) -> int:
     """Restore the regional compile on cache-hooked blocks' COMPUTED steps.
 
     ``enable_cache`` replaces each block's ``forward`` with the hook's ``new_forward``
-    (stashing the pre-hook bound method in ``fn_ref.original_forward``), and every cache
-    ``new_forward`` is ``@torch.compiler.disable``d because its skip decision is
-    data-dependent Python. The disable is recursive, so the compute branch's call into
-    ``original_forward`` runs EAGER and the block's regional compile artifact
-    (``_compiled_call_impl``) is never reached: measured 1.69 vs 1.09 s/step on
-    HunyuanVideo-1.5-720p, i.e. the cache forfeited the whole compile win on every
-    non-skipped step. An explicitly ``torch.compile``d callable re-enables dynamo for
-    its own extent even inside a disabled frame, so re-pointing
+    (stashing the pre-hook bound method in ``fn_ref.original_forward``), whose skip
+    decision is data-dependent Python: MagCache ``@torch.compiler.disable``s the whole
+    ``new_forward`` (recursive -- the compute branch runs EAGER), and even FBCache's
+    traceable ``new_forward`` graph-breaks around its disabled threshold decision,
+    which on some archs (measured: Qwen-Image) drops the compute branch's call into
+    ``original_forward`` out of the compiled region -- the block's regional compile
+    artifact (``_compiled_call_impl``) is never reached and the cache forfeits the
+    compile win on every non-skipped step. An explicitly ``torch.compile``d callable
+    re-enables
+    dynamo for its own extent even inside a disabled frame, so re-pointing
     ``fn_ref.original_forward`` at a compiled wrapper of the same bound method restores
-    compiled compute steps while the skip decision stays eager exactly as designed
-    (measured: identical skip counts, balanced MagCache 39.4 -> 26.9 s at 50 steps).
+    compiled compute steps while the skip decision stays eager exactly as designed.
+    Measured (B200, scripts/image_speedmem_bench.py): Qwen-Image FBCache computed steps
+    91.8 -> 71.2 ms (= the uncached compiled rate), 1.21x end to end; FLUX.1-dev is
+    neutral (its FBCache ``new_forward`` happens to trace, so computed steps were
+    already compiled -- same-process armed vs unarmed latents bit-identical); on the
+    video DiT balanced MagCache went 39.4 -> 26.9 s at 50 steps.
 
     Only blocks the speed layer actually compiled are armed (``_compiled_call_impl``
     guard -- eager tiers stay untouched), and only when ``original_forward`` is a plain

studio/backend/core/inference/diffusion_precision.py

@@ -421,6 +421,24 @@ def _cast_int8_selective(encoder: Any, target: Any, skip_first: int, skip_last:
     quantize_(encoder, _make_quant_config(TQ_INT8), filter_fn = filter_fn)
 
 
+def _weight_has_zero_output_row(module: Any) -> bool:
+    """True when a Linear's weight contains an all-zero OUTPUT row. torchao's per-row
+    fp8 scheme derives a per-output-channel scale from that row's amax, so a dead row
+    yields scale 0 -> 0/0 = NaN through the whole forward. Real checkpoints ship such
+    rows: SDXL's text_encoder_2 (OpenCLIP ViT-bigG) has one in
+    ``text_model.encoder.layers.2.self_attn.out_proj`` -- measured on B200: every
+    fp8_dynamic SDXL render came out black (NaN embeddings) until this Linear is left
+    dense. Cheap (one amax per Linear, once per load); False on any error so the
+    caster's own failure handling stays in charge."""
+    try:
+        weight = getattr(module, "weight", None)
+        if weight is None or weight.ndim != 2:
+            return False
+        return bool((weight.abs().amax(dim = -1) == 0).any().item())
+    except Exception:  # noqa: BLE001 -- unreadable weight: let quantize_ decide
+        return False
+
+
 def _cast_fp8_dynamic(encoder: Any, target: Any) -> None:
     # torchao dynamic fp8 COMPUTE, per-row (per-token activation + per-output-channel weight ->
     # torch._scaled_mm on the fp8 tensor cores). Unlike the layerwise `fp8` backend this keeps the
@@ -436,9 +454,15 @@ def _cast_fp8_dynamic(encoder: Any, target: Any) -> None:
 
     # require_bf16: scaled_mm asserts a bf16 weight, so skip any stray non-bf16 Linear the encoder
     # keeps (belt-and-suspenders over the named T5 wo exclusion) rather than aborting the pass.
-    filter_fn = make_filter_fn(
+    base = make_filter_fn(
         DEFAULT_MIN_LINEAR_FEATURES, _te_exclude_tokens(encoder), require_bf16 = True
     )
+
+    # A Linear with an all-zero output row NaNs under per-row scaling (scale 0 -> 0/0);
+    # keep exactly those Linears dense so one dead row cannot black out every render.
+    def filter_fn(module: Any, fqn: str = "") -> bool:
+        return base(module, fqn) and not _weight_has_zero_output_row(module)
+
     quantize_(encoder, _make_quant_config(TQ_FP8), filter_fn = filter_fn)
 
 

studio/backend/core/inference/diffusion_speed.py

@@ -360,11 +360,14 @@ def _compile_repeated_blocks(
                     setattr(dynamo_cfg, _limit_attr, max(getattr(dynamo_cfg, _limit_attr) or 0, 64))
         # Match eager's intermediate rounding inside inductor's fused pointwise kernels:
         # by default they keep chains in fp32 where eager materialises bf16 between ops,
-        # a per-forward rounding delta (max abs ~0.09 on the HunyuanVideo-1.5 DiT) that a
-        # multi-step denoise amplifies chaotically -- measured full-clip LPIPS vs the
-        # bit-exact reference drops 0.221 -> 0.052 at ZERO speed cost (1.093 vs 1.089
-        # s/step on a B200, 720p/33f). Process-global, so snapshot_backend_flags carries
-        # it and unload restores the prior value.
+        # a per-forward rounding delta that a multi-step denoise amplifies chaotically.
+        # Measured (B200, scripts/image_speedmem_bench.py, pairwise LPIPS of the
+        # compiled tier vs the same-stack eager tier): Qwen-Image 0.019 -> 0.006 at
+        # identical speed, FLUX.1-dev 0.046 -> 0.029 at +2% step time, FLUX.2-klein
+        # 0.018 -> 0.017 at identical speed; on the video DiT (HunyuanVideo-1.5-720p)
+        # full-clip LPIPS vs bit-exact drops 0.221 -> 0.052 at zero cost. Process-
+        # global, so snapshot_backend_flags carries it and unload restores the prior
+        # value.
         inductor_cfg = _inductor_config()
         if inductor_cfg is not None and hasattr(inductor_cfg, "emulate_precision_casts"):
             inductor_cfg.emulate_precision_casts = True

studio/backend/core/inference/llama_cpp.py

@@ -100,6 +100,10 @@ class LlamaServerNotFoundError(RuntimeError):
     Subclasses RuntimeError so existing handlers still catch it."""
 
 
+class _LlamaStreamCancelled(Exception):
+    """Internal signal for an expected client/request cancellation."""
+
+
 # Shared so the from_identifier preflight and the load-time raise stay in sync.
 LLAMA_SERVER_NOT_FOUND_DETAIL = (
     "This is a GGUF model, but the llama.cpp runtime (llama-server) is not "
@@ -1436,6 +1440,50 @@ def _backfill_usage_from_timings(usage, timings):
     return out
 
 
+def _vulkan_lib_filename() -> str:
+    return "ggml-vulkan.dll" if sys.platform == "win32" else "libggml-vulkan.so"
+
+
+# Host RAM to leave free on an integrated GPU, matching llama.cpp's own --fit
+# margin (default 1024 MiB per device). ggml reports an iGPU's "VRAM" as shared
+# system RAM, so hold back the same margin rather than inventing a larger one.
+_IGPU_HOST_RESERVE_MIB = 1024
+
+
+def _apply_igpu_host_reserve_mib(free_mib: int, is_igpu: bool) -> int:
+    """Reserve host headroom on an integrated (shared-memory) Vulkan GPU.
+
+    An iGPU's reported free "VRAM" is really free system RAM, so sizing
+    context/offload against all of it would push the host into swap or the OOM
+    killer. Leave the same margin llama.cpp's --fit uses. ``is_igpu`` comes from
+    ggml's device type, so a discrete card is never touched; only ever reduces.
+    """
+    if not is_igpu:
+        return free_mib
+    return max(0, free_mib - _IGPU_HOST_RESERVE_MIB)
+
+
+def _llama_lib_dir(binary: str) -> Path:
+    # The installer exposes llama-server as a top-level entrypoint into build/bin/,
+    # where the ggml backend libs live, so callers looking for sibling libs (Vulkan
+    # detection, LD_LIBRARY_PATH, probe bindir) need the real dir. It is normally a
+    # symlink (resolve() reaches build/bin), but create_exec_entrypoint falls back to
+    # a shell wrapper (exec "$(dirname "$0")/build/bin/llama-server" "$@") when it
+    # cannot symlink, and resolve() stops at the wrapper file. Follow the wrapper's
+    # exec target too, so a wrapper-based install still finds build/bin.
+    resolved = Path(binary).resolve()
+    try:
+        with open(resolved, "rb") as _f:
+            _head = _f.read(256)
+        if _head.startswith(b"#!"):
+            _m = re.search(r'exec "\$\(dirname "\$0"\)/([^"]+)"', _head.decode("utf-8", "ignore"))
+            if _m:
+                return (resolved.parent / _m.group(1)).resolve().parent
+    except OSError:
+        pass
+    return resolved.parent
+
+
 def _is_external_link(path: Path) -> bool:
     """True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink
     or a Windows directory junction / reparse point. Such a link resolves into
@@ -2278,6 +2326,30 @@ class LlamaCppBackend:
 
         return total
 
+    @staticmethod
+    def _is_vulkan_backend(binary: Optional[str] = None) -> bool:
+        """True if the installed llama.cpp build is Vulkan-only.
+
+        The official prebuilts are single-backend, so the Vulkan ggml lib next
+        to llama-server identifies a Vulkan build. Keeps the free-memory probe
+        and GPU pin in ggml's Vulkan device-index space. For a custom
+        multi-backend build with a CUDA or HIP ggml lib alongside Vulkan, defer
+        to that backend (torch-usable, better-understood probe/pin).
+        """
+        binary = binary or LlamaCppBackend._find_llama_server_binary()
+        if not binary:
+            return False
+        lib_dir = _llama_lib_dir(binary)
+        if not (lib_dir / _vulkan_lib_filename()).is_file():
+            return False
+        for _backend in ("cuda", "hip"):
+            sibling = (
+                f"ggml-{_backend}.dll" if sys.platform == "win32" else f"libggml-{_backend}.so"
+            )
+            if (lib_dir / sibling).is_file():
+                return False
+        return True
+
     @staticmethod
     def _resolve_visible_physical_ids() -> Optional[list[int]]:
         """Physical GPU ids behind the active visibility mask (HIP/ROCR/CUDA on
@@ -2440,11 +2512,42 @@ class LlamaCppBackend:
         return True
 
     @staticmethod
-    def _get_gpu_free_memory() -> list[tuple[int, int]]:
+    def _visible_devices_mask(env_name: str) -> Optional[set[int]]:
+        """Physical indices a ``*_VISIBLE_DEVICES`` mask permits, or None if unset.
+
+        ``if x.strip()`` filters trailing-comma masks ("0,1,"); an empty mask
+        ("") yields an empty set (all devices hidden), distinct from an unset
+        var (None, no mask). Used by the nvidia-smi probe.
+        """
+        raw = os.environ.get(env_name)
+        if raw is None:
+            return None
+        try:
+            return set(int(x.strip()) for x in raw.split(",") if x.strip())
+        except ValueError:
+            return None
+
+    @staticmethod
+    def _vulkan_pin_args(gpu_indices: Optional[Iterable[int]]) -> list[str]:
+        """``--device Vulkan<i>,...`` to pin a Vulkan launch to selected GPUs.
+
+        The indices are ggml's compact Vulkan ordinals (as _get_gpu_free_memory
+        reports and the registry names ``Vulkan<i>``). Pin by that name, NOT via
+        GGML_VK_VISIBLE_DEVICES: ggml parses that env var in the raw
+        vkEnumeratePhysicalDevices space (before dropping CPU/llvmpipe devices
+        and deduplicating ICDs), so a compact ordinal there could select a
+        different physical device or the CPU rasterizer.
+        """
+        if not gpu_indices:
+            return []
+        return ["--device", ",".join(f"Vulkan{i}" for i in gpu_indices)]
+
+    @staticmethod
+    def _get_gpu_free_memory(binary: Optional[str] = None) -> list[tuple[int, int]]:
         """Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by
         index; empty if no supported GPU is reachable. Thin wrapper over
         ``_get_gpu_memory`` for callers that only need free VRAM."""
-        return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()]
+        return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory(binary)]
 
     @staticmethod
     def _apple_metal_memory_budget_bytes() -> int:
@@ -2475,7 +2578,7 @@ class LlamaCppBackend:
         return int(rec_bytes * _APPLE_UNIFIED_MEMORY_FRACTION)
 
     @staticmethod
-    def _get_gpu_memory() -> list[tuple[int, int, int]]:
+    def _get_gpu_memory(binary: Optional[str] = None) -> list[tuple[int, int, int]]:
         """Query free AND total memory per GPU.
 
         Order:
@@ -2487,9 +2590,18 @@ class LlamaCppBackend:
              probe returned [] on AMD) and NVIDIA hosts missing
              ``nvidia-smi`` from PATH.
 
+        On a Vulkan build the ggml Vulkan probe is authoritative, so the indices
+        are ggml's compact Vulkan ordinals (the space the pin selects via
+        ``--device Vulkan<i>``). It reports ``total`` for discrete cards and 0
+        for an iGPU (shared RAM) so the fit falls back to free*frac there.
+        Otherwise nvidia-smi / torch cover NVIDIA + AMD ROCm.
+
         Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no
-        supported GPU is reachable. ``total`` lets the fit reserve absolute headroom.
+        supported GPU is reachable.
         """
+        binary = binary or LlamaCppBackend._find_llama_server_binary()
+        if LlamaCppBackend._is_vulkan_backend(binary):
+            return LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
         # ── NVIDIA via nvidia-smi ────────────────────────────────────
         try:
             result = subprocess.run(
@@ -2505,16 +2617,7 @@ class LlamaCppBackend:
                 **_windows_hidden_subprocess_kwargs(),
             )
             if result.returncode == 0:
-                allowed: Optional[set[int]] = None
-                cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
-                if cvd is not None:
-                    try:
-                        # `if x.strip()` filters trailing-comma masks ("0,1,").
-                        # Empty mask (CVD="") yields an empty set -> all GPUs
-                        # filtered out, per codebase convention.
-                        allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip())
-                    except ValueError:
-                        pass
+                allowed = LlamaCppBackend._visible_devices_mask("CUDA_VISIBLE_DEVICES")
                 gpus: list[tuple[int, int, int]] = []
                 for line in result.stdout.strip().splitlines():
                     parts = [p.strip() for p in line.split(",")]
@@ -2579,6 +2682,91 @@ class LlamaCppBackend:
             logger.debug(f"torch GPU probe failed: {e}")
             return []
 
+    @staticmethod
+    def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]:
+        """Query free (and total) VRAM per device via the bundled ggml Vulkan backend.
+
+        Loads ``libggml-vulkan`` in a short-lived subprocess (no Vulkan instance
+        in this process) and returns (device_index, free_mib, total_mib) sorted
+        by index. The index is ggml's compact Vulkan ordinal -- the one the
+        registry names ``Vulkan<index>`` and load_model pins with ``--device``,
+        NOT the raw ``GGML_VK_VISIBLE_DEVICES`` space. A user-set
+        ``GGML_VK_VISIBLE_DEVICES`` is honored by ggml (passed through), so the
+        list already reflects it. iGPUs leave a host-RAM margin (see
+        ``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass
+        their real total through. [] when no Vulkan build or device is reachable.
+        """
+        binary = binary or LlamaCppBackend._find_llama_server_binary()
+        if not binary:
+            return []
+        binary_dir = _llama_lib_dir(binary)
+        if not (binary_dir / _vulkan_lib_filename()).is_file():
+            return []
+
+        env = child_env_without_native_path_secret()
+        # Pass any inherited GGML_VK_VISIBLE_DEVICES through to ggml unchanged so
+        # the probe enumerates the same device list the launch will, named
+        # Vulkan0..N in the compact order reported here and pinned by that name
+        # via --device -- probe, mask, and pin stay in one index space. Do NOT
+        # filter the mask in Python: ggml parses the env var in raw
+        # vkEnumeratePhysicalDevices space while this probe reports the compact
+        # post-filter ordinal, so a Python filter would compare mismatched spaces.
+        if sys.platform != "win32":
+            # Let the loader resolve sibling ggml libs next to the binary.
+            existing_ld = env.get("LD_LIBRARY_PATH", "")
+            env["LD_LIBRARY_PATH"] = (
+                f"{binary_dir}:{existing_ld}" if existing_ld else str(binary_dir)
+            )
+        probe_script = Path(__file__).with_name("_vulkan_probe.py")
+        try:
+            result = subprocess.run(
+                [sys.executable, str(probe_script), str(binary_dir)],
+                capture_output = True,
+                text = True,
+                timeout = 15,
+                env = env,
+                **_windows_hidden_subprocess_kwargs(),
+            )
+            if result.returncode != 0:
+                logger.debug(
+                    f"vulkan GPU probe exited {result.returncode}: {result.stderr.strip()}"
+                )
+                return []
+        except Exception as e:
+            logger.debug(f"vulkan GPU probe failed: {e}")
+            return []
+
+        gpus: list[tuple[int, int, int]] = []
+        for line in result.stdout.strip().splitlines():
+            parts = line.split("\t")
+            if len(parts) != 4:
+                continue
+            try:
+                idx = int(parts[0])
+                free_mib = int(parts[1]) // (1024 * 1024)
+                is_igpu = parts[2] == "1"
+                # iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the
+                # fit stays on free*frac (the host reserve below is its
+                # headroom); a discrete card passes its real total through.
+                total_mib = 0 if is_igpu else int(parts[3]) // (1024 * 1024)
+            except ValueError:
+                continue
+            capped = _apply_igpu_host_reserve_mib(free_mib, is_igpu)
+            if capped < free_mib:
+                logger.info(
+                    f"Vulkan device VK{idx} is an integrated GPU sharing system "
+                    f"RAM; reserving {free_mib - capped}MiB host headroom "
+                    f"({free_mib}->{capped}MiB usable)"
+                )
+            gpus.append((idx, capped, total_mib))
+        gpus.sort(key = lambda g: g[0])
+        if gpus:
+            logger.info(
+                "Vulkan GPU memory detected: "
+                + ", ".join(f"VK{idx}={free}MiB" for idx, free, _total in gpus)
+            )
+        return gpus
+
     @staticmethod
     def _available_system_memory_mib() -> Optional[int]:
         """Available system RAM in MiB (psutil, then /proc/meminfo), or None if
@@ -2807,7 +2995,8 @@ class LlamaCppBackend:
     def _llama_server_env_for_binary(binary: str) -> dict[str, str]:
         """Build a subprocess env that lets llama-server resolve native libs."""
         env = child_env_without_native_path_secret()
-        binary_dir = str(Path(binary).parent)
+        # _llama_lib_dir resolves the llama-server symlink to the real build/bin.
+        binary_dir = str(_llama_lib_dir(binary))
 
         if sys.platform == "win32":
             # Ordering: see _build_windows_path_dirs. #5106.
@@ -4488,6 +4677,29 @@ class LlamaCppBackend:
             cancel_event = cancel_event,
         )
 
+    def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]:
+        """A drafter already in this repo's local HF cache, reused offline when a
+        fresh copy can't be fetched. Prefers a repo-root ``mtp-*.gguf`` across all
+        cached snapshots; else an existing ``MTP/`` copy (any precision -- the
+        target verifies every drafted token). None if none is cached."""
+        try:
+            from utils.models.model_config import _iter_hf_cache_snapshots
+
+            roots: list[Path] = []
+            subdirs: list[Path] = []
+            for snap in _iter_hf_cache_snapshots(hf_repo):  # newest first
+                for f in sorted(_gguf_snapshot_files(snap)):
+                    if _is_companion_gguf_path(f) and "mmproj" not in f.lower():
+                        (roots if "/" not in f else subdirs).append(snap / f)
+            # Keep snapshot order (newest first), root before any MTP/ copy, so a
+            # newer main GGUF pairs with the newest cached drafter, not a stale one.
+            for cand in roots + subdirs:
+                if cand.is_file():
+                    return str(cand)
+        except Exception as e:
+            logger.debug("Cached MTP drafter lookup failed for %s: %s", hf_repo, e)
+        return None
+
     def _download_mtp(
         self,
         *,
@@ -4504,11 +4716,25 @@ class LlamaCppBackend:
         are intentionally skipped. Returns the local path, or None.
         """
 
+        # Offline, reuse any drafter already on disk (a fresh copy can't be
+        # fetched). Online, _download_companion_gguf/hf_hub_download reuse the
+        # current cached file and refetch a changed one, so skip the probe here
+        # rather than pair new weights with a stale draft.
+        if _hf_env_offline():
+            cached = self._cached_repo_mtp_drafter(hf_repo)
+            if cached:
+                logger.info(f"Reusing cached MTP drafter (offline): {cached}")
+                return cached
+
         def _pick_mtp(candidates: list[str]) -> Optional[str]:
+            # Root-level only: MTP/ subdir copies now share the mtp- prefix but
+            # are explicit-selection, not auto-fetch (they'd sort ahead of root).
             mtp_files = sorted(
                 f
                 for f in candidates
-                if f.lower().endswith(".gguf") and Path(f).name.lower().startswith("mtp-")
+                if f.lower().endswith(".gguf")
+                and "/" not in f
+                and Path(f).name.lower().startswith("mtp-")
             )
             return mtp_files[0] if mtp_files else None
 
@@ -5210,6 +5436,7 @@ class LlamaCppBackend:
             # Resolve llama-server now but defer a not-found error: a block-diffusion
             # GGUF uses the diffusion runner, and its arch is only known after the header.
             binary = self._find_llama_server_binary()
+            is_vulkan_backend = self._is_vulkan_backend(binary)
 
             # ── Phase 2: download (NO lock held, so cancel can proceed) ──
             # mtp_draft_path arrives set for local Gemma loads (detected
@@ -5449,7 +5676,8 @@ class LlamaCppBackend:
                     model_size = gguf_size + mmproj_size
                     # 2-tuple gpus for existing logic + a total map for the absolute
                     # per-GPU headroom (correct when the GPU is already partly used).
-                    _gpu_mem = self._get_gpu_memory()
+                    # Pass binary so a Vulkan build probes ggml's Vulkan ordinals.
+                    _gpu_mem = self._get_gpu_memory(binary)
                     gpus = [(idx, free) for idx, free, _t in _gpu_mem]
                     total_by_idx = {idx: total for idx, _f, total in _gpu_mem}
 
@@ -6222,7 +6450,12 @@ class LlamaCppBackend:
                 # cap, not the ROCm-reported VRAM, is the real ceiling); refuse an
                 # oversize load the OS would otherwise kill mid-flight. Base model
                 # only: an optional MTP drafter is dropped by the MTP-drop fallback.
-                if model_size is not None and self._amd_apu_wants_unified_memory(gpu_indices):
+                # CUDA/ROCm ids only; a Vulkan build's gpu_indices are ggml ordinals.
+                if (
+                    model_size is not None
+                    and not is_vulkan_backend
+                    and self._amd_apu_wants_unified_memory(gpu_indices)
+                ):
                     _ram_msg = self._apu_ram_shortfall_message(
                         model_size, self._available_system_memory_mib()
                     )
@@ -6485,6 +6718,12 @@ class LlamaCppBackend:
                             ", ".join(unsupported_cache_flags),
                         )
 
+                # Vulkan pins via --device (a cmd arg, unlike the env-based
+                # CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's
+                # last-wins parsing lets a user --device override Studio's pick.
+                if is_vulkan_backend and gpu_indices is not None:
+                    cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices)
+
                 # User pass-through args go last so llama.cpp's last-wins parsing
                 # lets the user override Studio's auto-set flags. Already
                 # validated by the route via validate_extra_args().
@@ -6536,23 +6775,25 @@ class LlamaCppBackend:
                         env.setdefault("OMP_NUM_THREADS", "2")
 
                 # AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use
-                # shared system RAM. setdefault so a user value wins.
-                if self._amd_apu_wants_unified_memory(gpu_indices):
+                # shared system RAM. setdefault so a user value wins. Not on Vulkan
+                # (nor DC below): gpu_indices are ggml ordinals, not CUDA/ROCm ids.
+                if not is_vulkan_backend and self._amd_apu_wants_unified_memory(gpu_indices):
                     env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1")
                     logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1")
 
                 # DC NVIDIA GPUs: FP32 accum (+ P2P / launch queues for multi-GPU).
                 # See _apply_datacenter_env; opt out with UNSLOTH_DISABLE_DC_TUNING=1.
-                if self._apply_datacenter_env(env, gpu_indices):
+                if not is_vulkan_backend and self._apply_datacenter_env(env, gpu_indices):
                     multi_gpu = self._effective_gpu_count(gpu_indices) > 1
                     logger.info(
                         f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})"
                     )
 
                 # Pin to selected GPU(s). On ROCm, narrowing only
-                # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full
-                # set, so set HIP_VISIBLE_DEVICES too.
-                if gpu_indices is not None:
+                # CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so
+                # set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device
+                # (above), not here.
+                if gpu_indices is not None and not is_vulkan_backend:
                     pinned = ",".join(str(i) for i in gpu_indices)
                     env["CUDA_VISIBLE_DEVICES"] = pinned
                     try:
@@ -8379,7 +8620,7 @@ class LlamaCppBackend:
     ):
         """Open one streaming POST and let cancel interrupt prefill or reads."""
         if cancel_event is not None and cancel_event.is_set():
-            raise GeneratorExit
+            raise _LlamaStreamCancelled
 
         _cancel_closed = threading.Event()
         _response_ref: list = [None]
@@ -8424,13 +8665,13 @@ class LlamaCppBackend:
             ) as response:
                 _response_ref[0] = response
                 if cancel_event is not None and cancel_event.is_set():
-                    raise GeneratorExit
+                    raise _LlamaStreamCancelled
                 yield response
                 return
         except (httpx.RequestError, RuntimeError):
             # Response was closed by the cancel watcher
             if cancel_event is not None and cancel_event.is_set():
-                raise GeneratorExit
+                raise _LlamaStreamCancelled
             raise
         finally:
             _cancel_closed.set()
@@ -8633,6 +8874,8 @@ class LlamaCppBackend:
                         "finish_reason": _metadata_finish_reason,
                     }
 
+        except _LlamaStreamCancelled:
+            return
         except httpx.ConnectError as e:
             # Server already down. If this was an MTP+tensor crash, recover by
             # reloading without MTP (scheduled in the background) and fail this
@@ -9757,6 +10000,8 @@ class LlamaCppBackend:
                         break
                 continue
 
+            except _LlamaStreamCancelled:
+                return
             except httpx.ConnectError:
                 # Mark unresolved provisional cards as failed before raising.
                 for _pid, _pname in provisional_started_tool_calls.items():
@@ -9939,6 +10184,8 @@ class LlamaCppBackend:
                 if _meta is not None:
                     yield _meta
 
+        except _LlamaStreamCancelled:
+            return
         except httpx.ConnectError:
             raise RuntimeError("Lost connection to llama-server")
         except Exception as e:

studio/backend/core/inference/llama_keepwarm.py

@@ -61,12 +61,15 @@ _INFERENCE_SUFFIXES = (
     "/responses",
     "/generate/stream",  # Studio's own streaming route on the same llama-server
     "/audio/generate",  # direct GGUF TTS; can outlive the idle TTL
-    # Image/video generation holds a multi-GB diffusion/video pipeline for the whole request.
-    # Tracking them here lets other_inference_request_count() see an in-flight generation, so an
+    # Image generation holds a multi-GB diffusion pipeline for the whole request.
+    # Tracking it here lets other_inference_request_count() see an in-flight generation, so an
     # API-key training start is refused (409) before its unload cancels the generation. endswith
     # so the GET *-progress and */cancel variants are not matched.
     "/images/generate",  # /api/inference/images/generate
     "/images/generations",  # /v1/images/generations (+ /api/inference/images/generations)
+    # Video generation runs as a background job (the POST returns at once), so this entry only
+    # covers the brief accept request; the training-start guards additionally probe the video
+    # backend's generate-progress for an in-flight background clip.
     "/video/generate",  # /api/inference/video/generate
 )
 

studio/backend/core/inference/video.py

@@ -96,6 +96,7 @@ from .diffusion_precision import TE_QUANT_AUTO, normalize_te_quant, quantize_tex
 from .diffusion_vae_quant import VAE_QUANT_AUTO, normalize_vae_quant, quantize_vae
 from .video_families import (
     VIDEO_CANCELLED_MSG,
+    VIDEO_GENERATION_BUSY_MSG,
     VIDEO_NOT_LOADED_MSG,
     VideoFamily,
     default_video_generation_params,
@@ -409,6 +410,10 @@ class VideoBackend:
         self._active_generate_cancel: Optional[threading.Event] = None
         # Generation progress, written by the step callback / phase transitions.
         self._gen: dict[str, Any] = {"active": False}
+        # True from begin_generate() until its worker records a terminal state, so
+        # a second begin_generate() is refused while the first still runs (or is
+        # about to run: generate() only sets _gen after taking its locks).
+        self._generate_job_active = False
 
     # ── validation ───────────────────────────────────────────────────────────
 
@@ -1713,6 +1718,162 @@ class VideoBackend:
                 except Exception:  # noqa: BLE001 -- reset is best-effort, never fail a generation
                     pass
 
+    def begin_generate(
+        self,
+        *,
+        prompt: str,
+        negative_prompt: Optional[str] = None,
+        width: Optional[int] = None,
+        height: Optional[int] = None,
+        num_frames: Optional[int] = None,
+        fps: Optional[int] = None,
+        steps: Optional[int] = None,
+        guidance: Optional[float] = None,
+        guidance_2: Optional[float] = None,
+        seed: Optional[int] = None,
+    ) -> None:
+        """Validate cheaply, then run generate + gallery persist on a daemon thread.
+
+        Returns at once, mirroring begin_load: a clip takes minutes to denoise, and
+        a proxy in front of Studio (secure mode's Cloudflare tunnel) caps the origin
+        response window near 100 seconds, so the HTTP call must not span the
+        generation. The terminal outcome (phase "completed" with the saved gallery
+        record, or "failed" with a client-safe error) is reported by
+        generate_progress(); cancel_generate() keeps working against the job.
+        Raises RuntimeError with VIDEO_NOT_LOADED_MSG / VIDEO_GENERATION_BUSY_MSG
+        sentinels the route maps to 409.
+        """
+        cancel = threading.Event()
+        with self._lock:
+            if self._state is None:
+                raise RuntimeError(VIDEO_NOT_LOADED_MSG)
+            if self._generate_job_active:
+                raise RuntimeError(VIDEO_GENERATION_BUSY_MSG)
+            self._generate_job_active = True
+            # Register the cancel event BEFORE the worker starts so a cancel (or an
+            # unload) that lands in the spawn window still stops the run instead of
+            # returning "nothing to cancel".
+            self._active_generate_cancel = cancel
+            self._gen = {
+                "active": True,
+                "phase": "queued",
+                "step": 0,
+                "total": 0,
+                "eta_seconds": None,
+            }
+        threading.Thread(
+            target = self._run_generate,
+            kwargs = dict(
+                prompt = prompt,
+                negative_prompt = negative_prompt,
+                width = width,
+                height = height,
+                num_frames = num_frames,
+                fps = fps,
+                steps = steps,
+                guidance = guidance,
+                guidance_2 = guidance_2,
+                seed = seed,
+                cancel_event = cancel,
+            ),
+            daemon = True,
+        ).start()
+
+    def _run_generate(self, *, cancel_event: threading.Event, **gen_kwargs: Any) -> None:
+        """begin_generate's worker: generate, persist to the gallery, record the
+        terminal state where generate_progress() reports it. The error mapping is
+        the exact one the route applied when the call was synchronous: ValueError
+        text is client input feedback, sentinel RuntimeErrors pass through, and any
+        other failure is logged server-side and reported as a generic message so
+        internals (CUDA state, paths) never reach the client."""
+        from . import video_gallery
+
+        try:
+            result = self.generate(cancel_event = cancel_event, **gen_kwargs)
+        except ValueError as exc:
+            self._finish_generate_job(cancel_event = cancel_event, error = str(exc))
+            return
+        except RuntimeError as exc:
+            msg = str(exc)
+            if msg not in (VIDEO_NOT_LOADED_MSG, VIDEO_CANCELLED_MSG):
+                logger.error("video.generate_failed: %s", exc, exc_info = True)
+                msg = "Video generation failed."
+            self._finish_generate_job(cancel_event = cancel_event, error = msg)
+            return
+        except Exception as exc:  # noqa: BLE001 -- worker thread: never propagate
+            logger.error("video.generate_failed: %s", exc, exc_info = True)
+            self._finish_generate_job(cancel_event = cancel_event, error = "Video generation failed.")
+            return
+
+        # Persist the clip with its full recipe as the JSON sidecar the gallery reads back.
+        created_at = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
+        try:
+            record = video_gallery.save(
+                result["mp4_bytes"],
+                {
+                    "prompt": gen_kwargs["prompt"],
+                    "negative_prompt": gen_kwargs.get("negative_prompt"),
+                    "width": result["width"],
+                    "height": result["height"],
+                    "num_frames": result["num_frames"],
+                    "fps": result["fps"],
+                    "duration_s": result["duration_s"],
+                    "steps": result["steps"],
+                    "guidance": result["guidance"],
+                    "guidance_2": gen_kwargs.get("guidance_2"),
+                    "seed": result["seed"],
+                    "has_audio": result["has_audio"],
+                    "model": result["repo_id"],
+                    "created_at": created_at,
+                },
+            )
+        except Exception as exc:  # noqa: BLE001 -- disk failure must reach the poller
+            logger.error("video.persist_failed: %s", exc)
+            self._finish_generate_job(
+                cancel_event = cancel_event, error = "Failed to save the generated video."
+            )
+            return
+        self._finish_generate_job(cancel_event = cancel_event, video = record, total = result["steps"])
+
+    def _finish_generate_job(
+        self,
+        *,
+        cancel_event: Optional[threading.Event] = None,
+        video: Optional[dict] = None,
+        error: Optional[str] = None,
+        total: int = 0,
+    ) -> None:
+        """Record a job's terminal state as one atomic swap. The terminal dict
+        replaces the live-progress one so a poll can never mix fields from both,
+        and the busy flag drops in the same critical section so the earliest
+        moment a new begin_generate() can start is after the outcome is visible."""
+        with self._lock:
+            self._generate_job_active = False
+            if cancel_event is not None and self._active_generate_cancel is cancel_event:
+                # generate() clears its own registration; this covers a job whose
+                # worker failed before (or without) reaching generate()'s finally.
+                # Identity-guarded so a direct generate() that registered its own
+                # event in the meantime keeps its cancel handle.
+                self._active_generate_cancel = None
+            if error is not None:
+                self._gen = {
+                    "active": False,
+                    "phase": "failed",
+                    "error": error,
+                    "step": 0,
+                    "total": 0,
+                    "eta_seconds": None,
+                }
+            else:
+                self._gen = {
+                    "active": False,
+                    "phase": "completed",
+                    "video": video,
+                    "step": total,
+                    "total": total,
+                    "eta_seconds": None,
+                }
+
     def generate(
         self,
         *,
@@ -1726,9 +1887,13 @@ class VideoBackend:
         guidance: Optional[float] = None,
         guidance_2: Optional[float] = None,
         seed: Optional[int] = None,
+        cancel_event: Optional[threading.Event] = None,
     ) -> dict[str, Any]:
         import torch
-        cancel = threading.Event()
+
+        # begin_generate passes the event it already registered (so a cancel in the
+        # spawn window is honoured); a direct call makes its own.
+        cancel = cancel_event if cancel_event is not None else threading.Event()
         with self._generate_lock:
             with self._lock:
                 state = self._state
@@ -2013,7 +2178,14 @@ class VideoBackend:
                 pass
 
     def generate_progress(self) -> dict[str, Any]:
-        gen = dict(self._gen)
+        with self._lock:
+            gen = dict(self._gen)
+            # generate() swaps in a bare {"active": False} on its own exit paths
+            # before the job worker records the terminal dict; report the job as
+            # still active across that gap so a poller only sees active drop
+            # together with a terminal phase ("completed" / "failed").
+            if self._generate_job_active:
+                gen["active"] = True
         gen.setdefault("active", False)
         return gen
 

studio/backend/core/inference/video_families.py

@@ -25,6 +25,7 @@ from typing import Optional
 # these EXACTLY to return 409 (client-recoverable) instead of a sanitized 500.
 VIDEO_NOT_LOADED_MSG = "No video model is loaded."
 VIDEO_CANCELLED_MSG = "Video generation was cancelled."
+VIDEO_GENERATION_BUSY_MSG = "A video generation is already in progress."
 
 
 @dataclass(frozen = True)

studio/backend/core/training/training.py

@@ -739,6 +739,10 @@ class TrainingBackend:
         # True while a pump thread should be running; cleared on intended exits.
         # Left True after an abnormal death so _ensure_pump_alive spots a crash.
         self._pump_running: bool = False
+        # True from the start_training() guard passing until its spawn attempt
+        # finishes; blocks a second concurrent start (routes call start_training
+        # from a worker thread, so overlapping requests are possible).
+        self._start_in_progress: bool = False
         self._lock = threading.Lock()
 
         # Progress state (updated by pump thread from subprocess events)
@@ -800,11 +804,32 @@ class TrainingBackend:
         still letting auto-selection place training against the freed memory.
         Hook failures never block the start.
         """
+        # Compare-and-set start guard: the route runs this whole method on a worker
+        # thread (asyncio.to_thread), so two overlapping /train/start requests can
+        # reach it concurrently. Without the flag both would pass the alive-check
+        # below (the proc is only assigned at the end) and double-spawn. Mirrors the
+        # diffusion training service's reserve().
         with self._lock:
+            if self._start_in_progress:
+                logger.warning("Training start already in progress")
+                return False
             if self._proc is not None and self._proc.is_alive():
                 logger.warning("Training subprocess already running")
                 return False
+            self._start_in_progress = True
+        try:
+            return self._start_training_impl(job_id, before_spawn = before_spawn, **kwargs)
+        finally:
+            with self._lock:
+                self._start_in_progress = False
 
+    def _start_training_impl(
+        self,
+        job_id: str,
+        *,
+        before_spawn = None,
+        **kwargs,
+    ) -> bool:
         # Join prior pump thread — refuse to start if it won't die
         if self._pump_thread is not None and self._pump_thread.is_alive():
             self._pump_thread.join(timeout = 5.0)

studio/backend/hub/utils/gguf_plan.py

@@ -102,16 +102,17 @@ def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]:
 def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]:
     """The separate MTP drafter to fetch with every variant: the repo-root
     ``mtp-*.gguf`` copy unsloth ships for llama.cpp ``-hf`` auto-discovery
-    (Gemma 4). Same pick as the loader's drafter resolution (``mtp-`` basename
-    prefix, first in sort order) so download and load resolve the same file;
-    the higher-precision ``MTP/`` subdir copies are for explicit selection and
-    are not auto-fetched. None for repos with the head baked into the main
-    GGUF (Qwen)."""
+    (Gemma 4). Same pick as the loader's drafter resolution (root-level
+    ``mtp-`` prefix, first in sort order) so download and load resolve the same
+    file; the higher-precision ``MTP/`` subdir copies are for explicit
+    selection and are not auto-fetched. None for repos with the head baked into
+    the main GGUF (Qwen)."""
+    # Root-level only: the MTP/ subdir copies now share the mtp- prefix too.
     candidates = sorted(
         (
             s
             for s in siblings
-            if (name := _gguf_rfilename(s)) and name.lower().rsplit("/", 1)[-1].startswith("mtp-")
+            if (name := _gguf_rfilename(s)) and "/" not in name and name.lower().startswith("mtp-")
         ),
         key = lambda s: getattr(s, "rfilename"),
     )

studio/backend/main.py

@@ -767,9 +767,14 @@ _BODY_UPLOAD_PASSTHROUGH_EXACT_PATHS = (_DIFFUSION_DATASET_UPLOAD_PATH,)
 def _get_upload_passthrough_request_max_bytes(path: str) -> int:
     if path.startswith(_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX):
         return upload_request_limit_bytes(UNSTRUCTURED_RECIPE_UPLOAD_MAX_BYTES)
+    # The trailing-slash variant (/api/train/diffusion/dataset/) reaches this middleware
+    # BEFORE the router's redirect_slashes 307, so it must resolve to the same upload cap
+    # as the canonical path or a large upload 413s on the default /api/train body cap.
+    # Stripping slashes cannot promote a JSON sub-route: those all keep extra path
+    # components after normalization and still miss the exact match.
     if (
         path.startswith(_DATASET_UPLOAD_PASSTHROUGH_PREFIX)
-        or path == _DIFFUSION_DATASET_UPLOAD_PATH
+        or path.rstrip("/") == _DIFFUSION_DATASET_UPLOAD_PATH
     ):
         return upload_request_limit_bytes()
     return default_request_body_limit_bytes()
@@ -831,7 +836,10 @@ class MaxBodyMiddleware:
         self.upload_passthrough_exact_paths = upload_passthrough_exact_paths
 
     def _is_upload_passthrough(self, path: str) -> bool:
-        return path in self.upload_passthrough_exact_paths or any(
+        # Exact paths also match their trailing-slash variant: the middleware runs
+        # before the router's redirect_slashes 307, and a JSON sub-route can never
+        # normalize down to the exact path (it keeps extra components).
+        return path.rstrip("/") in self.upload_passthrough_exact_paths or any(
             path.startswith(p) for p in self.upload_passthrough_prefixes
         )
 

studio/backend/models/inference.py

@@ -1533,12 +1533,41 @@ class AnthropicToolResultBlock(BaseModel):
     tool_use_id: str
     content: Union[str, list] = ""
 
+    @field_validator("content", mode = "before")
+    @classmethod
+    def _coerce_null_content(cls, v):
+        # Some clients send null content for an empty tool result; the str|list
+        # union would 400 on it, so treat null as "".
+        return "" if v is None else v
+
+
+# Block types the converter translates explicitly. Anything else (thinking /
+# redacted_thinking, a provider block a resumed session replays, or a future type)
+# is accepted as an unknown block and dropped by the converter, rather than 400-ing
+# the whole request on strict validation.
+_KNOWN_ANTHROPIC_BLOCK_TYPES = frozenset({"text", "image", "tool_use", "tool_result"})
+
+
+class AnthropicUnknownBlock(BaseModel):
+    type: str
+    model_config = {"extra": "allow"}
+
+    @field_validator("type")
+    @classmethod
+    def _only_unknown_types(cls, v):
+        # Known types parse as their typed models above (so a malformed known block
+        # still fails cleanly); this fallback only catches the rest.
+        if v in _KNOWN_ANTHROPIC_BLOCK_TYPES:
+            raise ValueError("known block type handled by its typed model")
+        return v
+
 
 AnthropicContentBlock = Union[
     AnthropicTextBlock,
     AnthropicImageBlock,
     AnthropicToolUseBlock,
     AnthropicToolResultBlock,
+    AnthropicUnknownBlock,
 ]
 
 
@@ -1583,6 +1612,40 @@ class AnthropicMessage(BaseModel):
     role: Literal["user", "assistant"]
     content: Union[str, list[AnthropicContentBlock]]
 
+    @model_validator(mode = "before")
+    @classmethod
+    def _normalize_content(cls, data):
+        # Role-aware leniency that never silently drops real user input:
+        #  - assistant: a resumed tool-only turn's null content -> "" (str|list would
+        #    400 on null; "" keeps the converter's `for block in content` safe).
+        #    Unknown blocks (thinking / future types) validate via
+        #    AnthropicUnknownBlock and are dropped by the converter.
+        #  - user: keep strict. Null user content stays None so str|list rejects it
+        #    (400) rather than forwarding an empty prompt; and reject block types the
+        #    converter cannot translate, since it silently skips unknown user blocks
+        #    -- a user turn made only of them would validate yet send no content
+        #    (silent data loss).
+        if not isinstance(data, dict):
+            return data
+        content = data.get("content")
+        if data.get("role") == "assistant":
+            # Coerce only an explicit null (resumed tool-only turn). A missing
+            # content key stays malformed so the required-field check still 400s.
+            if "content" in data and content is None:
+                return {**data, "content": ""}
+            return data
+        if isinstance(content, list):
+            for block in content:
+                btype = (
+                    block.get("type") if isinstance(block, dict) else getattr(block, "type", None)
+                )
+                # Guard the value: a non-string type is unsupported too, and a
+                # membership test on an unhashable value would raise TypeError
+                # (escaping as a 500 instead of a clean 400).
+                if not isinstance(btype, str) or btype not in _KNOWN_ANTHROPIC_BLOCK_TYPES:
+                    raise ValueError(f"unsupported content block type {btype!r} in a user message")
+        return data
+
 
 class AnthropicTool(BaseModel):
     # Client tools have input_schema; server tools may only have type/name.
@@ -2492,9 +2555,21 @@ class GalleryVideo(BaseModel):
 
 
 class VideoGenerateResponse(BaseModel):
-    """The persisted gallery record for one generation call."""
+    """Acknowledgement that a generation was accepted and started.
 
-    video: GalleryVideo = Field(..., description = "Saved record for the generated clip")
+    Generation runs as a background job (a clip takes minutes, and secure mode's
+    tunnel caps the origin response window near 100 seconds, so the POST cannot
+    span it). The saved gallery record arrives via GET /video/generate-progress
+    when its phase reaches "completed"."""
+
+    status: Literal["started"] = Field(
+        "started", description = "Discriminator: the generation job was started"
+    )
+    video: Optional[GalleryVideo] = Field(
+        None,
+        description = "Always null (kept for response-shape compatibility); the saved "
+        "record is delivered by generate-progress on completion",
+    )
 
 
 class VideoGalleryListResponse(BaseModel):
@@ -2505,13 +2580,23 @@ class VideoGalleryListResponse(BaseModel):
 
 
 class VideoGenerateProgressResponse(BaseModel):
-    """Live progress for an in-flight video generation."""
+    """Live progress for an in-flight video generation, plus the terminal outcome
+    of the background job POST /video/generate started."""
 
     active: bool = Field(False, description = "Whether a generation is running")
-    phase: Optional[str] = Field(None, description = "Current phase: denoise | export | null")
+    phase: Optional[str] = Field(
+        None,
+        description = "Current phase: queued | denoise | export | completed | failed | null",
+    )
     step: int = Field(0, description = "Denoising steps completed so far")
     total: int = Field(0, description = "Total denoising steps for this run")
     eta_seconds: Optional[float] = Field(None, description = "Estimated seconds remaining")
+    video: Optional[GalleryVideo] = Field(
+        None, description = "Saved gallery record when phase is 'completed'"
+    )
+    error: Optional[str] = Field(
+        None, description = "Client-safe failure detail when phase is 'failed'"
+    )
 
 
 class VideoLoadProgressResponse(BaseModel):

studio/backend/routes/models.py

@@ -36,6 +36,11 @@ class CachedModelRepo(BaseModel):
     # weights). The picker must not treat a partial base repo as a usable download, or an
     # On Device click routes to a fresh multi-GB re-download instead of the complete GGUF.
     partial: Optional[bool] = None
+    # True for a diffusion-tagged repo with NO top-level model_index.json: a single-file
+    # checkpoint that needs from_single_file + a filename. The task-scoped pickers must not
+    # offer it as a pipeline load (from_pretrained on it fails after the GPU handoff)
+    # unless the curated catalog carries its artifact.
+    single_file: Optional[bool] = None
 
 
 class CachedModelsResponse(BaseModel):
@@ -3380,6 +3385,30 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)):
         return {"cached": []}
 
 
+def _repo_has_pipeline_index(repo_info) -> bool:
+    """Whether the cached snapshot carries a ROOT model_index.json, i.e. is loadable
+    as a full diffusers pipeline (from_pretrained reads only the repo root). A nested
+    subdir/model_index.json does not count: loading the repo root still fails, so the
+    row must keep its single_file flag. CachedFileInfo.file_name is the basename, so
+    a name match alone would also claim nested copies -- scope by file_path when the
+    scan provides it."""
+    try:
+        for rev in repo_info.revisions:
+            snapshot = getattr(rev, "snapshot_path", None)
+            for f in rev.files:
+                name = str(getattr(f, "file_name", "") or "")
+                path = getattr(f, "file_path", None)
+                if path is not None and snapshot is not None:
+                    p = Path(path)
+                    if p.name == "model_index.json" and p.parent == Path(snapshot):
+                        return True
+                elif name == "model_index.json":
+                    return True
+    except Exception:
+        pass
+    return False
+
+
 def _repo_is_diffusers(repo_info) -> bool:
     """True for an image-diffusion repo, so the chat picker hides it (it renders
     images, not chat) and the Images picker claims it — mirroring how cached
@@ -3390,13 +3419,8 @@ def _repo_is_diffusers(repo_info) -> bool:
     Qwen-Image or a z-image .safetensors) ship none. For those, fall back to the
     repo id resolving to a known diffusion family — the same resolver the Images
     backend loads from — so they don't surface as loadable chat models."""
-    try:
-        for rev in repo_info.revisions:
-            for f in rev.files:
-                if f.file_name == "model_index.json" or f.file_name.endswith("/model_index.json"):
-                    return True
-    except Exception:
-        pass
+    if _repo_has_pipeline_index(repo_info):
+        return True
     try:
         from core.inference.diffusion_families import detect_family
         if detect_family(getattr(repo_info, "repo_id", "") or "") is not None:
@@ -3502,6 +3526,11 @@ async def list_cached_models(
                         }
                         if is_partial:
                             row["partial"] = True
+                        # Flag diffusion repos with no pipeline index: loadable only via
+                        # from_single_file with a checkpoint filename, so the pickers must
+                        # not offer them as pipeline loads unless the catalog carries them.
+                        if row["task"] is not None and not _repo_has_pipeline_index(repo_info):
+                            row["single_file"] = True
                         # Keep the newest timestamp across duplicate caches;
                         # attach only when known so absent rows sort as oldest.
                         lm = max(last_modified, (existing or {}).get("last_modified", 0.0))

studio/backend/routes/training.py

@@ -134,6 +134,21 @@ async def get_visible_hardware_utilization(current_subject: str = Depends(get_cu
     return get_visible_gpu_utilization()
 
 
+def _background_video_generation_active() -> bool:
+    """Whether a video clip is generating on the video backend's worker thread.
+
+    POST /video/generate returns at once and generates in the background, so an
+    in-flight clip is invisible to the keep-warm in-flight request count the
+    API-key training guards consult; ask the backend directly. Best-effort: a
+    probe failure must never block a training start."""
+    try:
+        from core.inference.video import get_video_backend
+        return bool(get_video_backend().generate_progress().get("active"))
+    except Exception as e:  # noqa: BLE001
+        logger.warning("Could not check video generation state for training guard: %s", e)
+        return False
+
+
 @router.post("/start")
 async def start_training(
     request: TrainingStartRequest,
@@ -156,7 +171,10 @@ async def start_training(
         # session is not yet special-cased.)
         if via_api_key is True:
             from core.inference.llama_keepwarm import other_inference_request_count
-            if other_inference_request_count(current_request_counted = False) > 0:
+            if (
+                other_inference_request_count(current_request_counted = False) > 0
+                or _background_video_generation_active()
+            ):
                 raise HTTPException(
                     status_code = 409,
                     detail = (
@@ -506,8 +524,18 @@ async def start_training(
                 logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e)
 
         # The hook runs only once start guards pass -> VRAM freed iff training starts.
-        success = backend.start_training(
-            job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs
+        # Offloaded to a worker thread: the hook's diffusion/video unload() waits on the
+        # engines' generation locks until an in-flight denoise step reaches its cancel
+        # callback (and the export subprocess teardown can take seconds), which would
+        # otherwise block the event loop and freeze every concurrent status/cancel/UI
+        # request -- the same reason start_diffusion_training runs
+        # _free_gpu_for_diffusion_training via asyncio.to_thread. Overlapping starts are
+        # serialized by the backend's own start-in-progress guard.
+        success = await asyncio.to_thread(
+            backend.start_training,
+            job_id = job_id,
+            before_spawn = _free_vram_for_training,
+            **training_kwargs,
         )
 
         if not success:
@@ -1213,6 +1241,29 @@ def _preflight_gated_base(base_model: str, hf_token: Optional[str]) -> None:
         return
 
 
+def _resolve_diffusion_data_dir(raw: str) -> Path:
+    """Resolve a diffusion-training ``data_dir``. The upload/labeling routes create and
+    manage image datasets directly under ``datasets_root()`` and the UI passes the bare
+    folder name back as ``data_dir``, but the generic :func:`resolve_dataset_path`
+    searches the LLM uploads and recipe dataset roots FIRST -- so an unrelated upload
+    file or recipe folder sharing that name would shadow the just-uploaded image
+    dataset (preflight 400 "not a directory", or training the wrong data). Prefer the
+    image dataset root for a bare single-component name that exists there; everything
+    else (explicit "uploads/..." / "recipes/..." prefixes, absolute paths, missing
+    names) resolves exactly as before."""
+    from utils.paths import datasets_root
+
+    value = str(raw or "").strip()
+    if value and "\x00" not in value:
+        p = Path(value)
+        # Single component and not ".." -> joining under datasets_root() cannot escape it.
+        if not p.is_absolute() and len(p.parts) == 1 and p.parts[0] != "..":
+            direct = datasets_root() / value
+            if direct.is_dir():
+                return direct
+    return resolve_dataset_path(raw)
+
+
 @router.post("/diffusion/start", response_model = DiffusionTrainingStartResponse)
 async def start_diffusion_training(
     body: DiffusionTrainingStartRequest,
@@ -1228,7 +1279,10 @@ async def start_diffusion_training(
     # a diffusion start cannot silently drop an active API inference request.
     if via_api_key is True:
         from core.inference.llama_keepwarm import other_inference_request_count
-        if other_inference_request_count(current_request_counted = False) > 0:
+        if (
+            other_inference_request_count(current_request_counted = False) > 0
+            or _background_video_generation_active()
+        ):
             raise HTTPException(
                 status_code = 409,
                 detail = (
@@ -1259,8 +1313,8 @@ async def start_diffusion_training(
     # trainer subprocess otherwise resolves them relative to its own cwd.
     config = body.model_dump()
     try:
-        from utils.paths import resolve_dataset_path, resolve_output_dir
-        config["data_dir"] = str(resolve_dataset_path(config["data_dir"]))
+        from utils.paths import resolve_output_dir
+        config["data_dir"] = str(_resolve_diffusion_data_dir(config["data_dir"]))
         config["output_dir"] = str(resolve_output_dir(config["output_dir"]))
     except ValueError as e:
         raise HTTPException(status_code = 400, detail = str(e))
@@ -1600,6 +1654,25 @@ async def upload_diffusion_dataset(
                 status_code = 400,
                 detail = f"Unsupported file '{f.filename}'. Allowed: {exts}",
             )
+        # Reject an EXACT duplicate name within THIS batch (two cat.png dragged from
+        # different folders, or an API client repeating a part). The same-name exemption
+        # below exists for SEPARATE repeat uploads, where re-sending a name is a
+        # deliberate overwrite of the file on disk; inside one batch the two parts are
+        # distinct files staged to the same destination on EVERY filesystem, so the later
+        # tmp.replace(dest) in the commit loop would silently discard the earlier one
+        # while `uploaded` still counts both. Exact match only: a case VARIANT pair
+        # (pic.png vs Pic.png) stays exempt like the stem guard documents -- one file /
+        # an overwrite on case-insensitive filesystems, two files on Linux.
+        fname_cf = filename.casefold()
+        if filename in names:
+            raise HTTPException(
+                status_code = 400,
+                detail = (
+                    f"Duplicate file '{filename}' appears more than once in this upload. "
+                    "Files sharing a name would overwrite each other; rename one before "
+                    "uploading."
+                ),
+            )
         # Reject a second IMAGE that shares this one's stem but differs by extension (sample.png
         # vs sample.jpg): both resolve to the same <stem>.txt caption sidecar (the kohya/diffusers
         # convention the reader, editor, and delete paths all use), so keeping both would silently
@@ -1616,7 +1689,6 @@ async def upload_diffusion_dataset(
             # one caption. Casefolding the name guard too keeps a same-name case variant
             # (sample.png vs Sample.png, one file / an overwrite on those filesystems) exempt.
             stem_cf = stem.casefold()
-            fname_cf = filename.casefold()
             clash = next(
                 (
                     p.name

studio/backend/routes/video.py

@@ -8,15 +8,16 @@ these routes mirror the /images/* routes one-for-one: the same validate-before-e
 load ordering, the same GPU arbiter handoff (VIDEO owner in place of DIFFUSION),
 the same error boundary mapping backend exceptions to HTTP, and the same gallery
 CRUD shape. The backend runs in-process and is synchronous, so the blocking
-load/generate/unload calls are offloaded with asyncio.to_thread to keep the event
-loop free. This module is the single error boundary: backend methods raise, we
-map to HTTP here.
+calls are offloaded with asyncio.to_thread to keep the event loop free; the slow
+operations (load AND generate) run as background jobs whose begin_* calls return
+at once, with progress + terminal outcome polled from their *-progress routes.
+This module is the single error boundary: backend methods raise, we map to HTTP
+here.
 """
 
 from __future__ import annotations
 
 import asyncio
-import time
 
 from fastapi import APIRouter, Depends, HTTPException, Response
 from pydantic import ValidationError
@@ -149,14 +150,18 @@ async def video_load_progress(current_subject: str = Depends(get_current_subject
 async def generate_video(
     request: VideoGenerateRequest, current_subject: str = Depends(get_current_subject)
 ):
-    from core.inference import video_gallery
+    """Start a generation job and return at once (the begin_load pattern): a clip
+    takes minutes, and secure mode's tunnel caps the origin response window near
+    100 seconds, so the response must not span the generation. The worker runs the
+    generate + gallery-persist pipeline; the terminal outcome (completed with the
+    saved record / failed with a client-safe error) arrives via generate-progress."""
     from core.inference.video import get_video_backend
-    from core.inference.video_families import VIDEO_CANCELLED_MSG, VIDEO_NOT_LOADED_MSG
+    from core.inference.video_families import VIDEO_GENERATION_BUSY_MSG, VIDEO_NOT_LOADED_MSG
 
     backend = get_video_backend()
     try:
-        result = await asyncio.to_thread(
-            backend.generate,
+        await asyncio.to_thread(
+            backend.begin_generate,
             prompt = request.prompt,
             negative_prompt = request.negative_prompt,
             width = request.width,
@@ -169,53 +174,19 @@ async def generate_video(
             seed = request.seed,
         )
     except ValueError as exc:
-        # Bad client input (a workflow the loaded family doesn't support) -- a 400 with
-        # the reason, not a generic 500.
+        # Bad client input -- a 400 with the reason, not a generic 500.
         raise HTTPException(status_code = 400, detail = str(exc))
     except RuntimeError as exc:
-        # Only "no model loaded" / user-cancelled are client-state (409). Match the
-        # sentinels exactly, not as a substring, so an execution failure that merely
-        # contains "cancelled" can't misroute to 409 and leak that output.
+        # Only "no model loaded" / "already generating" are client-state (409).
+        # Match the sentinels exactly, not as a substring, so an unrelated failure
+        # can't misroute to 409 and leak its message.
         msg = str(exc)
-        if msg in (VIDEO_NOT_LOADED_MSG, VIDEO_CANCELLED_MSG):
+        if msg in (VIDEO_NOT_LOADED_MSG, VIDEO_GENERATION_BUSY_MSG):
             raise HTTPException(status_code = 409, detail = msg)
         logger.error("video.generate_failed: %s", exc, exc_info = True)
         raise HTTPException(status_code = 500, detail = "Video generation failed.")
-    except Exception as exc:
-        logger.error("video.generate_failed: %s", exc, exc_info = True)
-        raise HTTPException(status_code = 500, detail = "Video generation failed.")
 
-    # Persist the clip with its full recipe as the JSON sidecar the gallery reads back.
-    created_at = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
-
-    def _persist() -> dict:
-        return video_gallery.save(
-            result["mp4_bytes"],
-            {
-                "prompt": request.prompt,
-                "negative_prompt": request.negative_prompt,
-                "width": result["width"],
-                "height": result["height"],
-                "num_frames": result["num_frames"],
-                "fps": result["fps"],
-                "duration_s": result["duration_s"],
-                "steps": result["steps"],
-                "guidance": result["guidance"],
-                "guidance_2": request.guidance_2,
-                "seed": result["seed"],
-                "has_audio": result["has_audio"],
-                "model": result["repo_id"],
-                "created_at": created_at,
-            },
-        )
-
-    try:
-        record = await asyncio.to_thread(_persist)
-    except Exception as exc:  # noqa: BLE001
-        logger.error("video.persist_failed: %s", exc)
-        raise HTTPException(status_code = 500, detail = "Failed to save the generated video.")
-
-    return VideoGenerateResponse(video = GalleryVideo(**record))
+    return VideoGenerateResponse()
 
 
 @router.get("/video/generate-progress", response_model = VideoGenerateProgressResponse)

studio/backend/tests/test_anthropic_messages.py

@@ -1770,3 +1770,187 @@ class TestAnthropicMessagesToolRouting:
 
         _drive(anthropic_messages(payload, request = None, current_subject = "t"))
         assert backend.calls[0][0] == "plain"
+
+
+def test_resumed_session_thinking_and_null_content_do_not_400():
+    # A resumed session replays assistant turns with `thinking` (and sometimes null)
+    # content. Those must be accepted (thinking dropped by the converter), not 400ed.
+    from pydantic import ValidationError
+
+    req = AnthropicMessagesRequest(
+        model = "x",
+        max_tokens = 16,
+        messages = [
+            {"role": "user", "content": "hi"},
+            {
+                "role": "assistant",
+                "content": [
+                    {"type": "thinking", "thinking": "secret reasoning", "signature": "s"},
+                    {"type": "text", "text": "the answer"},
+                    {"type": "tool_use", "id": "t1", "name": "f", "input": {}},
+                ],
+            },
+            {"role": "assistant", "content": None},  # tool-only turn serialized as null
+        ],
+    )
+    # Known blocks still parse as their typed models; only the unknown one is loose.
+    assert type(req.messages[1].content[0]).__name__ == "AnthropicUnknownBlock"
+    assert type(req.messages[1].content[1]).__name__ == "AnthropicTextBlock"
+    assert req.messages[2].content == ""  # null coerced
+
+    openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages])
+    assistant = next(m for m in openai if m["role"] == "assistant" and m.get("content"))
+    assert assistant["content"] == "the answer"
+    assert "secret reasoning" not in json.dumps(openai)  # thinking never forwarded
+
+    # A malformed KNOWN block still fails cleanly instead of being swallowed.
+    with pytest.raises(ValidationError):
+        AnthropicMessagesRequest(
+            model = "x",
+            max_tokens = 16,
+            messages = [{"role": "assistant", "content": [{"type": "tool_use", "name": "f"}]}],
+        )
+
+
+def test_user_null_content_rejected():
+    # The null->"" leniency is assistant-only; a null user content must be rejected
+    # at the boundary, not coerced into an empty prompt and forwarded to the model.
+    from pydantic import ValidationError
+    with pytest.raises(ValidationError):
+        AnthropicMessagesRequest(
+            model = "x",
+            max_tokens = 16,
+            messages = [{"role": "user", "content": None}],
+        )
+
+
+def test_user_unknown_block_rejected_not_silently_dropped():
+    # The converter skips user blocks it cannot translate, so a user turn whose only
+    # block is unknown would validate yet forward no content. Reject at the boundary
+    # to avoid that silent data loss (the assistant fallback is unaffected).
+    from pydantic import ValidationError
+    with pytest.raises(ValidationError):
+        AnthropicMessagesRequest(
+            model = "x",
+            max_tokens = 16,
+            messages = [
+                {"role": "user", "content": [{"type": "document", "source": {}}]},
+            ],
+        )
+
+
+def test_user_translatable_blocks_still_accepted():
+    # text / image / tool_result are translatable, so a real user message built from
+    # them must still pass; the unknown-block guard only trips on other types.
+    req = AnthropicMessagesRequest(
+        model = "x",
+        max_tokens = 16,
+        messages = [
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "What is this?"},
+                    {
+                        "type": "image",
+                        "source": {"type": "base64", "media_type": "image/png", "data": "AA"},
+                    },
+                    {"type": "tool_result", "tool_use_id": "t1", "content": "ok"},
+                ],
+            }
+        ],
+    )
+    assert [type(b).__name__ for b in req.messages[0].content] == [
+        "AnthropicTextBlock",
+        "AnthropicImageBlock",
+        "AnthropicToolResultBlock",
+    ]
+
+    openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages])
+    assert any(m["role"] == "tool" and m["tool_call_id"] == "t1" for m in openai)
+
+
+def test_user_malformed_known_block_still_rejected():
+    # The guard only allow-lists a user block's *type*; the union still validates its
+    # shape, so a known-but-malformed block (tool_result without tool_use_id) fails.
+    from pydantic import ValidationError
+    with pytest.raises(ValidationError):
+        AnthropicMessagesRequest(
+            model = "x",
+            max_tokens = 16,
+            messages = [
+                {"role": "user", "content": [{"type": "tool_result", "content": "x"}]},
+            ],
+        )
+
+
+def test_user_content_block_non_string_type_rejected_cleanly():
+    # A user block whose `type` is a non-string (unhashable list / dict, or a stray
+    # int) must fail as a clean validation error, not raise TypeError from the
+    # frozenset membership test and escape as a 500.
+    from pydantic import ValidationError
+    for bad_type in ([], {}, 5):
+        with pytest.raises(ValidationError):
+            AnthropicMessagesRequest(
+                model = "x",
+                max_tokens = 16,
+                messages = [{"role": "user", "content": [{"type": bad_type}]}],
+            )
+
+
+def test_assistant_missing_content_key_still_rejected():
+    # The null -> "" leniency is only for an EXPLICIT null. An assistant message that
+    # omits content entirely stays malformed and must fail required-field validation.
+    from pydantic import ValidationError
+
+    with pytest.raises(ValidationError):
+        AnthropicMessagesRequest(
+            model = "x",
+            max_tokens = 16,
+            messages = [{"role": "assistant"}],
+        )
+    # An explicit null is still accepted and coerced (regression guard).
+    req = AnthropicMessagesRequest(
+        model = "x",
+        max_tokens = 16,
+        messages = [
+            {"role": "user", "content": "hi"},
+            {"role": "assistant", "content": None},
+        ],
+    )
+    assert req.messages[1].content == ""
+
+
+def test_resumed_null_assistant_between_users_coalesced_on_messages_route(monkeypatch):
+    # user -> assistant(null) -> user is now accepted: the null assistant turn coerces
+    # to "" and is dropped. The route must then coalesce the two remaining user turns
+    # so a strict GGUF chat template does not 400 on non-alternating roles.
+    backend = _mock_backend(monkeypatch, context_length = 2048)
+
+    class _Req:
+        state = SimpleNamespace()
+        url = SimpleNamespace(path = "/v1/messages")
+        method = "POST"
+
+        async def is_disconnected(self):
+            return False
+
+    payload = AnthropicMessagesRequest(
+        model = "x",
+        max_tokens = 16,
+        messages = [
+            {"role": "user", "content": "first question"},
+            {"role": "assistant", "content": None},
+            {"role": "user", "content": "please continue"},
+        ],
+    )
+
+    response = _drive(anthropic_messages(payload, request = _Req(), current_subject = "t"))
+    assert response.status_code == 200
+
+    [(_path, kwargs)] = backend.calls
+    user_turns = [m for m in kwargs["messages"] if m.get("role") == "user"]
+    assert len(user_turns) == 1  # the two user turns were merged, not left adjacent
+    merged = user_turns[0]["content"]
+    if isinstance(merged, list):
+        merged = " ".join(p.get("text", "") for p in merged if isinstance(p, dict))
+    assert "first question" in merged and "please continue" in merged

studio/backend/tests/test_cached_gguf_routes.py

@@ -1076,3 +1076,70 @@ def test_cached_repo_partial_scopes_probe_to_snapshot_dir(monkeypatch):
 
     monkeypatch.setattr(scan, "is_snapshot_partial", _boom)
     assert models_route._cached_repo_partial("Org/Repo", snapshot_dir) is False
+
+
+def test_repo_has_pipeline_index_requires_root_model_index(tmp_path):
+    # Only a ROOT model_index.json makes a repo pipeline-loadable: from_pretrained
+    # reads the repo root, so a nested subdir/model_index.json must NOT clear the
+    # single_file flag. CachedFileInfo.file_name is the basename, so the helper has
+    # to scope by file_path/snapshot_path -- a name-only match would claim both.
+    snap = tmp_path / "snapshots" / "abc"
+    nested = SimpleNamespace(
+        file_name = "model_index.json",
+        file_path = snap / "prior" / "model_index.json",
+    )
+    repo_nested = SimpleNamespace(
+        repo_id = "unsloth/nested-index",
+        revisions = [SimpleNamespace(files = [nested], snapshot_path = snap)],
+    )
+    assert models_route._repo_has_pipeline_index(repo_nested) is False
+
+    root = SimpleNamespace(
+        file_name = "model_index.json",
+        file_path = snap / "model_index.json",
+    )
+    repo_root = SimpleNamespace(
+        repo_id = "unsloth/root-index",
+        revisions = [SimpleNamespace(files = [root], snapshot_path = snap)],
+    )
+    assert models_route._repo_has_pipeline_index(repo_root) is True
+
+
+def test_list_cached_models_flags_single_file_diffusion_repos(monkeypatch, tmp_path):
+    # A diffusion-tagged repo with NO top-level model_index.json is a single-file
+    # checkpoint: the task pickers must not offer it as a pipeline load (from_pretrained
+    # fails on it), so the row carries single_file=True. A full pipeline repo (has
+    # model_index.json) and a chat repo (task None) carry no flag.
+    single = _repo(
+        "unsloth/Qwen-Image-fp8-single",
+        [_file("qwen-image-fp8.safetensors", 10_000)],
+        tmp_path / "models--unsloth--Qwen-Image-fp8-single",
+    )
+    pipeline = _repo(
+        "unsloth/Qwen-Image-pipeline",
+        [_file("model_index.json", 10), _file("transformer/model.safetensors", 10_000)],
+        tmp_path / "models--unsloth--Qwen-Image-pipeline",
+    )
+    chat = _repo(
+        "Org/ChatRepo",
+        [_file("model.safetensors", 10_000)],
+        tmp_path / "models--Org--ChatRepo",
+    )
+
+    monkeypatch.setattr(
+        models_route,
+        "_cached_repo_task",
+        lambda repo_info: ("text-to-image" if "Qwen-Image" in repo_info.repo_id else None),
+    )
+    monkeypatch.setattr(
+        models_route,
+        "_all_hf_cache_scans",
+        lambda: [SimpleNamespace(repos = [single, pipeline, chat])],
+    )
+
+    result = asyncio.run(models_route.list_cached_models(current_subject = "test-user"))
+
+    rows = {r["repo_id"]: r for r in result["cached"]}
+    assert rows["unsloth/Qwen-Image-fp8-single"].get("single_file") is True
+    assert "single_file" not in rows["unsloth/Qwen-Image-pipeline"]
+    assert "single_file" not in rows["Org/ChatRepo"]

studio/backend/tests/test_diffusion_cache.py

@@ -945,6 +945,8 @@ def test_arming_skips_partial_captured_inner(monkeypatch):
 
 
 def test_arming_covers_every_cache_hook_family(monkeypatch):
+    # FBCache is the image cache today, but the hook-name table already covers the
+    # MagCache layout too (same fn_ref shape), so a future mode arms for free.
     _stub_torch_compile(monkeypatch)
     names = (
         "mag_cache_leader_block_hook",
@@ -1007,3 +1009,69 @@ def test_apply_step_cache_arms_compiled_blocks_on_toggle(monkeypatch):
     assert engaged == TC_MAGCACHE
     assert hook.fn_ref.original_forward is not orig
     assert hook._unsloth_orig_inner is orig
+
+
+def test_toggle_disable_restores_inners_before_disable(monkeypatch):
+    # remove_hook splices fn_ref.original_forward back into module.forward, so the
+    # compiled wrapper must be swapped out BEFORE disable_cache runs.
+    _stub_diffusers(monkeypatch)
+    order = []
+
+    class _T(_ToggleTransformer):
+        def disable_cache(self):
+            super().disable_cache()
+            order.append("disable")
+
+        def modules(self):
+            order.append("restore-walk")
+            return []
+
+    t = _T()
+    maybe_toggle_step_cache(_pipe(t), steps = 28)
+    mode = maybe_toggle_step_cache(_pipe(t), steps = 8)
+    assert mode is None and t.disables == 1
+    assert order[-2:] == ["restore-walk", "disable"]
+
+
+def test_enable_failure_restores_inners_before_partial_disable(monkeypatch):
+    # enable_cache can fail after hooking (and arming) some blocks; the partial-hook
+    # cleanup must un-arm them before disable_cache splices original_forward back.
+    _stub_diffusers(monkeypatch)
+    order = []
+
+    class _T(_ToggleTransformer):
+        def enable_cache(self, config):
+            raise RuntimeError("block signature not recognised")
+
+        def disable_cache(self):
+            super().disable_cache()
+            order.append("disable")
+
+        def modules(self):
+            order.append("restore-walk")
+            return []
+
+    t = _T()
+    assert apply_step_cache(_pipe(t), mode = "fbcache") is None
+    assert order == ["restore-walk", "disable"]
+
+
+# ── stale child-registry cache invalidation (mid-session enable) ────────────────────
+
+
+def test_enable_invalidates_stale_child_registry_cache(monkeypatch):
+    # diffusers 0.39 caches the child-registry list on first cache_context use; an
+    # UNCACHED generation already populates it (empty), so a later toggle-time
+    # enable_cache would install hooks the context never reaches ("No context is set").
+    _stub_diffusers(monkeypatch)
+    t = _MixinTransformer()
+    t._diffusers_hook = types.SimpleNamespace(_child_registries_cache = ["stale"])
+    assert apply_step_cache(_pipe(t), mode = "fbcache") == TC_FBCACHE
+    assert t._diffusers_hook._child_registries_cache is None
+
+
+def test_invalidate_child_registry_cache_tolerates_absence():
+    _invalidate_child_registry_cache(types.SimpleNamespace())  # no registry: no-op
+    reg = types.SimpleNamespace(_child_registries_cache = None)
+    _invalidate_child_registry_cache(types.SimpleNamespace(_diffusers_hook = reg))
+    assert reg._child_registries_cache is None

studio/backend/tests/test_diffusion_dataset_api.py

@@ -392,6 +392,31 @@ def test_upload_same_stem_collision_within_one_batch(client, ds_root):
     assert "Duplicate image name" in r.json()["detail"]
 
 
+def test_upload_rejects_exact_duplicate_name_within_one_batch(client, ds_root):
+    # Two parts with the SAME name in ONE multipart batch are distinct files (dragged from
+    # different folders, or an API client repeating a part); the staged commit would let the
+    # later tmp.replace(dest) silently discard the earlier one while `uploaded` still counts
+    # both. The batch must be rejected whole. Re-sending a name in a SEPARATE upload stays a
+    # deliberate overwrite (test_upload_allows_exact_name_overwrite_and_caption_sidecar).
+    r = _upload(
+        client,
+        "styleset",
+        [("sample.png", _png_bytes((10, 20, 30))), ("sample.png", _png_bytes((90, 90, 90)))],
+    )
+    assert r.status_code == 400
+    assert "more than once" in r.json()["detail"]
+    assert not (ds_root / "styleset" / "sample.png").exists()  # all-or-nothing
+    # Caption files collide at one destination the same way.
+    r = _upload(client, "styleset", [("sample.txt", b"a"), ("sample.txt", b"b")])
+    assert r.status_code == 400
+    assert "more than once" in r.json()["detail"]
+    # A case VARIANT pair (Cat.png vs cat.png) stays exempt, matching the stem-guard
+    # contract: it is one file / an overwrite on case-insensitive filesystems and two
+    # files on Linux, not silent same-destination data loss.
+    r = _upload(client, "styleset", [("Cat.png", _png_bytes()), ("cat.png", _png_bytes())])
+    assert r.status_code == 200
+
+
 def test_upload_allows_exact_name_overwrite_and_caption_sidecar(client, ds_root):
     # Re-uploading the EXACT same name (stem AND extension) is an allowed overwrite, and a .txt
     # caption for the same stem is the intended kohya flow -- neither is a same-stem image collision.

studio/backend/tests/test_diffusion_precision.py

@@ -647,3 +647,81 @@ def test_quantize_explicit_fp8_dynamic_refused_for_ltx2(monkeypatch):
         is None
     )
     assert calls == []
+
+
+# ── zero-output-row guard (per-row fp8 NaN protection) ───────────────────────────
+
+
+class _FakeAmaxVec:
+    def __init__(self, vals):
+        self._vals = vals
+
+    def __eq__(self, other):  # noqa: PLW0642 -- tensor-style elementwise compare
+        return _FakeAmaxVec([v == other for v in self._vals])
+
+    def any(self):
+        return _FakeScalar(any(self._vals))
+
+
+class _FakeScalar:
+    def __init__(self, v):
+        self._v = v
+
+    def item(self):
+        return self._v
+
+
+class _FakeWeight:
+    """Tensor-shaped stand-in supporting the exact chain the guard runs:
+    ``weight.abs().amax(dim = -1) == 0 -> .any().item()``."""
+
+    ndim = 2
+
+    def __init__(self, rows):
+        self._rows = rows
+
+    def abs(self):
+        return _FakeWeight([[abs(v) for v in r] for r in self._rows])
+
+    def amax(self, dim = -1):
+        return _FakeAmaxVec([max(r) for r in self._rows])
+
+
+def test_weight_zero_output_row_detection():
+    # A dead output row NaNs torchao's per-row fp8 (scale 0 -> 0/0); SDXL's
+    # text_encoder_2 (OpenCLIP bigG) really ships one in layers.2.self_attn.out_proj --
+    # measured: every fp8_dynamic SDXL render was black until the row is kept dense.
+    zero_row = types.SimpleNamespace(weight = _FakeWeight([[0.1, 0.2], [0.0, 0.0]]))
+    dense = types.SimpleNamespace(weight = _FakeWeight([[0.1, 0.2], [0.3, 0.0]]))
+    assert dp._weight_has_zero_output_row(zero_row) is True
+    assert dp._weight_has_zero_output_row(dense) is False
+    # Non-2D / absent weights are not the per-row scheme's input: never flagged.
+    w3 = _FakeWeight([[1.0]])
+    w3.ndim = 3
+    assert dp._weight_has_zero_output_row(types.SimpleNamespace(weight = w3)) is False
+    assert dp._weight_has_zero_output_row(types.SimpleNamespace()) is False
+
+    # An unreadable weight falls through to quantize_'s own handling.
+    class _Boom:
+        @property
+        def weight(self):
+            raise RuntimeError("meta tensor")
+
+    assert dp._weight_has_zero_output_row(_Boom()) is False
+
+
+def test_fp8_dynamic_filter_skips_zero_row_linear(monkeypatch):
+    # The fp8_dynamic caster must leave a zero-output-row Linear dense while the rest
+    # of the encoder still quantises (a family-wide deny would forfeit the whole win).
+    _stub_torch(monkeypatch)
+    captured: dict = {}
+    _stub_transformer_quant(monkeypatch, captured)
+    enc = types.SimpleNamespace(_keep_in_fp32_modules = [])
+
+    dp._cast_fp8_dynamic(enc, _target())
+
+    ff = captured["filter_fn"]
+    dead = types.SimpleNamespace(weight = _FakeWeight([[0.5, 0.5], [0.0, 0.0]]))
+    live = types.SimpleNamespace(weight = _FakeWeight([[0.5, 0.5], [0.5, 0.5]]))
+    assert ff(dead, "text_model.encoder.layers.2.self_attn.out_proj") is False
+    assert ff(live, "text_model.encoder.layers.2.mlp.fc1") is True

studio/backend/tests/test_diffusion_training.py

@@ -655,6 +655,30 @@ def test_route_start_rejects_uncontained_paths(client):
     assert r.status_code == 400
 
 
+def test_route_start_resolves_bare_name_under_image_dataset_root(client, monkeypatch, tmp_path):
+    # The upload/labeling routes manage image datasets directly under datasets_root() and
+    # the UI passes the bare folder name back as data_dir. The generic resolve_dataset_path
+    # searches the LLM uploads and recipe roots FIRST, so an unrelated upload file or recipe
+    # folder sharing the name would shadow the just-uploaded image dataset (preflight 400
+    # "not a directory", or training the wrong data). The route must prefer the image
+    # dataset root for a bare name that exists there.
+    import utils.paths as up
+
+    ds_root = tmp_path / "assets" / "datasets"
+    img_ds = ds_root / "my-photos"
+    img_ds.mkdir(parents = True)
+    (img_ds / "a.png").write_bytes(b"x")
+    # Shadowing entries the generic resolver would pick first.
+    (ds_root / "uploads").mkdir()
+    (ds_root / "uploads" / "my-photos").write_text("an LLM dataset upload, not a folder")
+    (ds_root / "recipes" / "my-photos").mkdir(parents = True)
+    monkeypatch.setattr(up, "datasets_root", lambda: ds_root)
+
+    r = client.post("/api/train/diffusion/start", json = {**_BODY, "data_dir": "my-photos"})
+    assert r.status_code == 200, r.text
+    assert client._fake.started_with["data_dir"] == str(img_ds)
+
+
 def test_route_start_blocked_by_active_llm_training(client, monkeypatch):
     import routes.training as tr
 

studio/backend/tests/test_inference_dispatcher_resilience.py

@@ -112,7 +112,7 @@ def test_route_llama_streaming_async_clients_disable_proxy_env():
             continue
         calls.append(node)
 
-    assert len(calls) == 4
+    assert len(calls) == 5
     for call in calls:
         assert any(
             kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False

studio/backend/tests/test_install_resolve_prebuilt.py

@@ -55,6 +55,27 @@ def _host(**kw):
     return ilp.HostInfo(**base)
 
 
+def test_force_cpu_clears_all_gpu_attributes_including_intel():
+    # --cpu-fallback is the "select the CPU prebuilt even when a GPU is present"
+    # escape hatch. It must drop EVERY GPU attribute, including has_intel_gpu, or
+    # the planner still prepends the Vulkan asset on an Intel-GPU host.
+    host = _host(
+        is_linux = True,
+        is_x86_64 = True,
+        has_usable_nvidia = True,
+        has_physical_nvidia = True,
+        has_rocm = True,
+        rocm_gfx_target = "gfx1100",
+        has_intel_gpu = True,
+    )
+    forced = ilp._apply_host_overrides(host, force_cpu = True)
+    assert forced.has_usable_nvidia is False
+    assert forced.has_physical_nvidia is False
+    assert forced.has_rocm is False
+    assert forced.rocm_gfx_target is None
+    assert forced.has_intel_gpu is False
+
+
 def test_macos_upstream_pin_only_for_explicit_pre26_upstream():
     pre26 = _host(
         system = "Darwin",
@@ -313,3 +334,152 @@ def test_sm103_host_drops_cuda128_windows_build():
     )
     kept_b200 = ilp._drop_blackwell_incapable_windows_cuda(b200, [cuda128, cuda129])
     assert [a.name for a in kept_b200] == [cuda128.name, cuda129.name]
+
+
+def _upstream_release(tag, asset_names):
+    return {
+        "tag_name": tag,
+        "assets": [
+            {"name": n, "browser_download_url": f"https://example/{n}"} for n in asset_names
+        ],
+    }
+
+
+def test_direct_upstream_arm64_intel_prefers_vulkan():
+    # Auto-detected Intel GPU on Linux arm64 -> Vulkan prebuilt first, CPU
+    # second (mirrors the x86_64 branch; ggml-org ships the arm64 Vulkan asset).
+    host = _host(is_linux = True, is_arm64 = True, machine = "aarch64", has_intel_gpu = True)
+    rel = _upstream_release(
+        "b9925",
+        ["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"],
+    )
+    plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
+    kinds = [a.install_kind for a in plan.attempts]
+    assert kinds[0] == "linux-vulkan", kinds
+    assert "linux-arm64" in kinds
+    assert plan.attempts[0].name == "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz"
+
+
+def test_direct_upstream_intel_with_hidden_nvidia_is_cpu_only():
+    # A host with a physical NVIDIA hidden via CUDA_VISIBLE_DEVICES (physical
+    # True, usable False) + an Intel iGPU must NOT get the Vulkan archive even
+    # when planning directly against upstream: Vulkan ignores CUDA_VISIBLE_DEVICES
+    # and could grab the reserved card. It falls through to the CPU asset.
+    host = _host(
+        is_linux = True,
+        is_x86_64 = True,
+        has_intel_gpu = True,
+        has_physical_nvidia = True,
+        has_usable_nvidia = False,
+    )
+    rel = _upstream_release(
+        "b9925",
+        ["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"],
+    )
+    plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
+    assert [a.install_kind for a in plan.attempts] == ["linux-cpu"]
+
+
+def test_direct_upstream_arm64_without_intel_is_cpu_only():
+    host = _host(is_linux = True, is_arm64 = True, machine = "aarch64")
+    rel = _upstream_release(
+        "b9925",
+        ["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"],
+    )
+    plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
+    assert [a.install_kind for a in plan.attempts] == ["linux-arm64"]
+
+
+def test_direct_upstream_x86_intel_prefers_vulkan():
+    host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
+    rel = _upstream_release(
+        "b9925",
+        ["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"],
+    )
+    plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
+    kinds = [a.install_kind for a in plan.attempts]
+    assert kinds[0] == "linux-vulkan", kinds
+    assert "linux-cpu" in kinds
+
+
+def test_linux_vulkan_health_glob_matches_bare_cpu_lib():
+    # The widened glob must cover both arch-suffixed (x64) and bare (arm64) CPU
+    # libs so a valid Vulkan install is not re-flagged unhealthy every check.
+    choice = ilp.AssetChoice(
+        repo = UPSTREAM,
+        tag = "b9925",
+        name = "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz",
+        url = "https://example/x",
+        source_label = "upstream",
+        install_kind = "linux-vulkan",
+    )
+    groups = ilp.runtime_payload_health_groups(choice)
+    assert ["libggml-cpu*.so*"] in groups
+    assert ["libggml-cpu-*.so*"] not in groups
+
+
+def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin():
+    # Routing fork -> upstream also drops the fork release pin, which is in a
+    # different tag namespace and would make the upstream resolver miss.
+    host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
+    routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = False)
+    assert repo == UPSTREAM
+    assert tag == ""
+    assert routed.has_intel_gpu is True
+
+
+def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin():
+    # A pin set WITH an explicit upstream repo is already on upstream -> kept.
+    host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
+    _routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, UPSTREAM, "b9596", force_cpu = False)
+    assert repo == UPSTREAM
+    assert tag == "b9596"
+
+
+def test_route_to_vulkan_prebuilt_cpu_fallback_wins():
+    # --cpu-fallback suppresses Vulkan routing even for an Intel host.
+    host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
+    routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = True)
+    assert repo == FORK
+    assert tag == "b9596-mix-abc"
+    assert routed is host
+
+
+def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted():
+    # A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1):
+    # physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or
+    # Vulkan (which ignores CUDA_VISIBLE_DEVICES) could grab the reserved GPU.
+    host = _host(
+        is_linux = True,
+        is_x86_64 = True,
+        has_intel_gpu = True,
+        has_physical_nvidia = True,
+        has_usable_nvidia = False,
+    )
+    _routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
+    assert repo == FORK
+
+
+def test_route_to_vulkan_prebuilt_rocm_host_not_rerouted():
+    # An Intel iGPU alongside a usable ROCm GPU stays on its ROCm/fork path.
+    host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True, has_rocm = True)
+    _routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
+    assert repo == FORK
+
+
+def test_route_to_vulkan_prebuilt_non_intel_unchanged():
+    host = _host(is_linux = True, is_x86_64 = True)
+    routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
+    assert repo == FORK
+    assert routed is host
+
+
+def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys):
+    # The --resolve-prebuilt probe must agree with the install path: an
+    # auto-detected Intel host resolves against upstream (Vulkan), not the fork.
+    monkeypatch.setattr(
+        ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
+    )
+    seen, out = _run_resolve_capture_host(monkeypatch, capsys)
+    assert seen["repo"] == UPSTREAM
+    assert out["repo"] == UPSTREAM

studio/backend/tests/test_llama_cpp_stream_cancel.py

@@ -0,0 +1,81 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
+
+import contextlib
+import os
+import sys
+import threading
+
+import httpx
+import pytest
+
+_backend = os.path.join(os.path.dirname(__file__), "..")
+sys.path.insert(0, _backend)
+
+from core.inference.llama_cpp import LlamaCppBackend, _LlamaStreamCancelled
+
+
+def _backend_stub() -> LlamaCppBackend:
+    backend = LlamaCppBackend.__new__(LlamaCppBackend)
+    backend._process = object()
+    backend._healthy = True
+    backend._port = 48848
+    backend._effective_context_length = 4096
+    backend._supports_reasoning = False
+    backend._reasoning_always_on = False
+    backend._reasoning_style = "enable_thinking"
+    backend._supports_preserve_thinking = False
+    return backend
+
+
+def test_stream_cancel_uses_internal_exception_not_generator_exit():
+    class FakeResponse:
+        status_code = 200
+
+        def close(self):
+            pass
+
+    class FakeStream:
+        def __enter__(self):
+            return FakeResponse()
+
+        def __exit__(self, *_args):
+            return False
+
+    class FakeClient:
+        def stream(self, *_args, **_kwargs):
+            return FakeStream()
+
+    cancel_event = threading.Event()
+
+    with pytest.raises(Exception) as exc_info:
+        with LlamaCppBackend._stream_with_retry(
+            FakeClient(),
+            "http://llama.test/v1/chat/completions",
+            {},
+            cancel_event,
+        ):
+            cancel_event.set()
+            raise httpx.ReadError("client closed")
+
+    assert exc_info.type is _LlamaStreamCancelled
+    assert not issubclass(exc_info.type, GeneratorExit)
+
+
+def test_generate_chat_completion_swallows_internal_stream_cancel(monkeypatch):
+    backend = _backend_stub()
+
+    @contextlib.contextmanager
+    def fake_open_stream(*_args, **_kwargs):
+        raise _LlamaStreamCancelled
+
+    monkeypatch.setattr(backend, "_open_stream", fake_open_stream)
+
+    chunks = list(
+        backend.generate_chat_completion(
+            [{"role": "user", "content": "hi"}],
+            cancel_event = threading.Event(),
+        )
+    )
+
+    assert chunks == []

studio/backend/tests/test_llama_cpp_update.py

@@ -448,6 +448,48 @@ def test_start_update_happy_path(monkeypatch, tmp_path):
     assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5"
 
 
+def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
+    # A Vulkan install (marker asset carries 'vulkan') must re-assert
+    # UNSLOTH_FORCE_VULKAN on update, or detect_host on a GPU box re-routes to
+    # CUDA/ROCm and silently replaces the Vulkan build.
+    install_dir = tmp_path / "llama.cpp"
+    binary = _write_install(
+        install_dir,
+        "b9493",
+        repo = "ggml-org/llama.cpp",
+        asset = "llama-b9493-bin-ubuntu-vulkan-x64.tar.gz",
+    )
+    monkeypatch.setattr(upd, "_find_binary", lambda: binary)
+    monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
+    monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
+
+    def _on_start(cmd):
+        _write_install(
+            install_dir,
+            "b9518",
+            repo = "ggml-org/llama.cpp",
+            asset = "llama-b9518-bin-ubuntu-vulkan-x64.tar.gz",
+        )
+
+    popen_kwargs: dict = {}
+    _patch_installer_popen(
+        monkeypatch,
+        lines = ["installed\n"],
+        on_start = _on_start,
+        captured_kwargs = popen_kwargs,
+    )
+
+    assert upd.start_update()["started"] is True
+    deadline = time.time() + 10
+    while time.time() < deadline:
+        job = upd.get_update_status()["job"]
+        if job["state"] in ("success", "error"):
+            break
+        time.sleep(0.05)
+    assert job["state"] == "success", job
+    assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1"
+
+
 def test_start_update_reports_full_release_tag(monkeypatch, tmp_path):
     install_dir = tmp_path / "llama.cpp"
     binary = _write_install(install_dir, "b9595")

studio/backend/tests/test_llama_cpp_vulkan_probe.py

@@ -0,0 +1,193 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""Vulkan free-VRAM reader regression tests on a synthetic probe output.
+
+Covers the post-probe handling in
+``LlamaCppBackend._get_gpu_free_memory_vulkan``:
+
+  * integrated GPUs (probe reports is_igpu=1) leave a flat per-device host
+    margin matching llama.cpp's --fit-target, so context auto-sizing can't
+    over-commit shared RAM, and report total 0 (shared RAM is not a budget),
+  * discrete GPUs (is_igpu=0) keep their free untouched and pass their real
+    total through so the fit can reserve absolute headroom,
+  * an inherited ``GGML_VK_VISIBLE_DEVICES`` is passed through to ggml unchanged
+    (ggml applies it), not stripped or filtered in Python -- the probe reports
+    ggml's compact ordinal, which load_model pins with ``--device Vulkan<i>``.
+
+The ggml Vulkan library is never loaded: subprocess.run is mocked to emit
+the tab-separated lines the real ``_vulkan_probe.py`` would print.
+"""
+
+from __future__ import annotations
+
+import subprocess
+import sys
+import types as _types
+from pathlib import Path
+from unittest import mock
+
+import pytest
+
+_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
+if _BACKEND_DIR not in sys.path:
+    sys.path.insert(0, _BACKEND_DIR)
+
+import importlib as _importlib  # noqa: E402
+
+
+def _maybe_stub(name: str, builder):
+    try:
+        _importlib.import_module(name)
+    except ImportError:
+        sys.modules[name] = builder()
+
+
+def _build_loggers_stub():
+    m = _types.ModuleType("loggers")
+    m.get_logger = lambda name: __import__("logging").getLogger(name)
+    return m
+
+
+_maybe_stub("loggers", _build_loggers_stub)
+_maybe_stub("structlog", lambda: _types.ModuleType("structlog"))
+
+from core.inference import llama_cpp as _llama_mod  # noqa: E402
+from core.inference.llama_cpp import (  # noqa: E402
+    LlamaCppBackend,
+    _llama_lib_dir,
+    _vulkan_lib_filename,
+)
+
+MIB = 1024 * 1024
+GIB = 1024 * MIB
+
+
+def _make_vulkan_install(tmp_path: Path) -> str:
+    """A binary whose sibling dir holds the Vulkan ggml lib, so the
+    reader's ``is_vulkan_backend`` sibling-file check passes."""
+    bindir = tmp_path / "build" / "bin"
+    bindir.mkdir(parents = True)
+    binary = bindir / ("llama-server.exe" if sys.platform == "win32" else "llama-server")
+    binary.write_bytes(b"stub")
+    (bindir / _vulkan_lib_filename()).write_bytes(b"stub")
+    return str(binary)
+
+
+def _mock_probe(rows: list[str], captured_env: dict | None = None):
+    """Patch subprocess.run so the _vulkan_probe.py call returns ``rows``
+    (already tab-formatted), recording the env it was launched with."""
+    real_run = subprocess.run
+
+    def fake_run(cmd, *args, **kwargs):
+        if isinstance(cmd, list) and any("_vulkan_probe" in str(c) for c in cmd):
+            if captured_env is not None:
+                captured_env.clear()
+                captured_env.update(kwargs.get("env") or {})
+            return subprocess.CompletedProcess(
+                args = cmd, returncode = 0, stdout = "\n".join(rows), stderr = ""
+            )
+        return real_run(cmd, *args, **kwargs)
+
+    return mock.patch("subprocess.run", side_effect = fake_run)
+
+
+def _row(
+    idx: int,
+    free_bytes: int,
+    is_igpu: int,
+    total_bytes: int = 0,
+) -> str:
+    return f"{idx}\t{free_bytes}\t{is_igpu}\t{total_bytes}"
+
+
+def test_integrated_gpu_leaves_host_margin(tmp_path):
+    binary = _make_vulkan_install(tmp_path)
+    # iGPU with 30 GiB free; reserve a flat 1024 MiB (llama.cpp --fit-target).
+    # total stays 0: shared system RAM is not a VRAM budget for the fit.
+    rows = [_row(0, 30 * GIB, is_igpu = 1, total_bytes = 32 * GIB)]
+    with _mock_probe(rows):
+        gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
+    assert gpus == [(0, 30 * 1024 - 1024, 0)], gpus
+
+
+def test_discrete_gpu_free_is_untouched_and_total_passed_through(tmp_path):
+    binary = _make_vulkan_install(tmp_path)
+    # 6 GiB free on a partially occupied 24 GiB card: free is untouched and the
+    # real total flows through so the fit reserves absolute headroom (CUDA/ROCm
+    # parity) instead of the looser free*frac budget.
+    rows = [_row(0, 6 * GIB, is_igpu = 0, total_bytes = 24 * GIB)]
+    with _mock_probe(rows):
+        gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
+    assert gpus == [(0, 6 * 1024, 24 * 1024)], gpus
+
+
+def test_large_discrete_gpu_is_untouched(tmp_path):
+    binary = _make_vulkan_install(tmp_path)
+    # A 48 GiB discrete card stays untouched regardless of size; only the
+    # iGPU flag triggers the host margin, never a VRAM/RAM ratio.
+    rows = [_row(0, 47 * GIB, is_igpu = 0, total_bytes = 48 * GIB)]
+    with _mock_probe(rows):
+        gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
+    assert gpus == [(0, 47 * 1024, 48 * 1024)], gpus
+
+
+def test_inherited_visible_devices_mask_is_passed_through_to_probe(tmp_path, monkeypatch):
+    # The mask is NOT stripped or filtered in Python: ggml parses it in raw
+    # physical-device space while this probe reports the compact post-filter
+    # ordinal, so mixing spaces would be wrong. It is passed through unchanged
+    # so ggml applies it to the same device list the launch will enumerate.
+    binary = _make_vulkan_install(tmp_path)
+    monkeypatch.setenv("GGML_VK_VISIBLE_DEVICES", "1")
+    captured: dict = {}
+    rows = [_row(0, 23 * GIB, is_igpu = 0, total_bytes = 24 * GIB)]
+    with _mock_probe(rows, captured_env = captured):
+        LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
+    assert captured.get("GGML_VK_VISIBLE_DEVICES") == "1", captured
+
+
+def test_vulkan_pin_args_uses_device_names_not_env_mask():
+    # Pin by compact device name via --device (the space the probe reports and
+    # the registry names), never by writing a compact ordinal into the raw
+    # GGML_VK_VISIBLE_DEVICES index space.
+    assert LlamaCppBackend._vulkan_pin_args([0]) == ["--device", "Vulkan0"]
+    assert LlamaCppBackend._vulkan_pin_args([1, 2]) == ["--device", "Vulkan1,Vulkan2"]
+    assert LlamaCppBackend._vulkan_pin_args(None) == []
+    assert LlamaCppBackend._vulkan_pin_args([]) == []
+
+
+def test_vulkan_only_build_is_detected(tmp_path):
+    binary = _make_vulkan_install(tmp_path)
+    assert LlamaCppBackend._is_vulkan_backend(binary) is True
+
+
+def test_multi_backend_build_is_not_vulkan_only(tmp_path):
+    # A custom build that ships CUDA (or HIP) alongside Vulkan must NOT be
+    # treated as Vulkan-only, or its CUDA GPU would be probed/pinned as a Vulkan
+    # device; defer to the CUDA/HIP path instead.
+    binary = _make_vulkan_install(tmp_path)
+    cuda = "ggml-cuda.dll" if sys.platform == "win32" else "libggml-cuda.so"
+    (_llama_lib_dir(binary) / cuda).write_bytes(b"stub")
+    assert LlamaCppBackend._is_vulkan_backend(binary) is False
+
+
+@pytest.mark.skipif(sys.platform == "win32", reason = "shell wrapper fallback is POSIX")
+def test_shell_wrapper_entrypoint_resolves_to_real_lib_dir(tmp_path):
+    # create_exec_entrypoint falls back to a #!/bin/sh wrapper at the install root
+    # when it cannot symlink; _find_llama_server_binary returns that root entrypoint,
+    # so _llama_lib_dir must follow the wrapper's exec target to build/bin -- else
+    # _is_vulkan_backend misses libggml-vulkan.so and the Vulkan probe/pin silently
+    # never engage on a valid Vulkan install.
+    import os
+
+    binary = _make_vulkan_install(tmp_path)  # tmp_path/build/bin/llama-server + vulkan lib
+    bindir = Path(binary).parent
+    wrapper = tmp_path / "llama-server"
+    wrapper.write_text('#!/bin/sh\nexec "$(dirname "$0")/build/bin/llama-server" "$@"\n')
+    os.chmod(wrapper, 0o755)
+    assert _llama_lib_dir(str(wrapper)) == bindir
+    assert LlamaCppBackend._is_vulkan_backend(str(wrapper)) is True
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))

studio/backend/tests/test_llama_route_timeouts.py

@@ -203,3 +203,87 @@ def test_preheader_send_cleanup_on_disconnect_and_cancel():
 
     asyncio.run(_run(False))
     asyncio.run(_run(True))
+
+
+def test_stream_stall_timeout_callable_re_resolved_each_read():
+    # The OpenAI passthrough passes a callable so the stall bound can switch to
+    # the short post-terminal grace mid-stream; it must be re-resolved per read,
+    # not captured once at generator start.
+    async def _run():
+        response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}}))
+        values = iter([100.0, 2.0])
+        seen = []
+
+        class _Request:
+            async def is_disconnected(self):
+                return False
+
+        class _Items:
+            def __init__(self):
+                self.count = 0
+
+            async def __anext__(self):
+                self.count += 1
+                if self.count > 3:
+                    raise StopAsyncIteration
+                return "data: {}"
+
+        async for _ in inf_mod._aiter_llama_stream_items(
+            _Items(),
+            cancel_event = threading.Event(),
+            request = _Request(),
+            response = response,
+            first_token_deadline = time.monotonic() + 1,
+            post_first_item_read_timeout_s = lambda: next(values, 5.0),
+        ):
+            seen.append(response.request.extensions["timeout"].get("read"))
+
+        assert len(seen) == 3
+        # The callable is resolved right after the first item (arming the
+        # post-first window) and again before each later read, consuming
+        # successive values.
+        assert seen[0] == 100.0
+        assert 1.0 <= seen[1] <= 2.0
+        assert 4.0 <= seen[2] <= 5.0
+
+    asyncio.run(_run())
+
+
+def test_stream_stall_timeout_disabled_clears_read_timeout():
+    # UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT=0 disables the stall guard, so
+    # the callable returns None. Once a chunk has arrived the leftover
+    # first-token read timeout must be cleared, else a long post-first-chunk gap
+    # trips a stale deadline the operator asked to turn off.
+    async def _run():
+        response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}}))
+        seen = []
+
+        class _Request:
+            async def is_disconnected(self):
+                return False
+
+        class _Items:
+            def __init__(self):
+                self.count = 0
+
+            async def __anext__(self):
+                self.count += 1
+                if self.count > 2:
+                    raise StopAsyncIteration
+                return "data: {}"
+
+        async for _ in inf_mod._aiter_llama_stream_items(
+            _Items(),
+            cancel_event = threading.Event(),
+            request = _Request(),
+            response = response,
+            first_token_deadline = time.monotonic() + 5,
+            post_first_item_read_timeout_s = lambda: None,
+        ):
+            seen.append(response.request.extensions["timeout"].get("read"))
+
+        # The first-token path armed a finite read timeout; after the first chunk
+        # with the guard disabled, it is cleared to None on every subsequent read.
+        assert seen == [None, None], seen
+
+    asyncio.run(_run())

studio/backend/tests/test_middleware.py

@@ -204,6 +204,57 @@ class TestMaxBodyMiddleware:
                 default_request_body_limit_bytes()
             ), path
 
+    def test_diffusion_dataset_trailing_slash_gets_upload_cap(self, main_module):
+        # The trailing-slash variant reaches the middleware BEFORE the router's
+        # redirect_slashes 307, so it must resolve to the same passthrough + upload cap
+        # as the canonical path or a large upload 413s on the default /api/train cap.
+        # JSON sub-routes keep extra components after normalization, so they stay capped.
+        from utils.upload_limits import (
+            default_request_body_limit_bytes,
+            upload_request_limit_bytes,
+        )
+
+        slashed = "/api/train/diffusion/dataset/"
+        assert main_module._get_upload_passthrough_request_max_bytes(slashed) == (
+            upload_request_limit_bytes()
+        )
+        # End-to-end through the middleware: a body over the default cap but under the
+        # upload cap passes through on both the canonical and the slashed path.
+        app = _make_protected_app(
+            128,
+            main_module,
+            upload_passthrough_max_bytes_getter = lambda _p: 1024,
+            upload_passthrough_exact_paths = ("/api/train/diffusion/dataset",),
+        )
+
+        @app.post("/api/train/diffusion/dataset")
+        async def upload(request: Request):
+            body = await request.body()
+            return {"total": len(body)}
+
+        c = TestClient(app)
+        for path in ("/api/train/diffusion/dataset", "/api/train/diffusion/dataset/"):
+            r = c.post(
+                path,
+                content = b"x" * 512,
+                headers = {"content-type": "application/octet-stream"},
+            )
+            assert r.status_code == 200, path
+            assert r.json()["total"] == 512, path
+        # A slashed JSON sub-route is still NOT passthrough: over-cap body is rejected.
+        r = c.post(
+            "/api/train/diffusion/dataset/import-example/",
+            content = b"x" * 512,
+            headers = {"content-type": "application/octet-stream"},
+        )
+        assert r.status_code == 413
+        assert (
+            main_module._get_upload_passthrough_request_max_bytes(
+                "/api/train/diffusion/dataset/import-example/"
+            )
+            == default_request_body_limit_bytes()
+        )
+
     def test_v1_surface_is_body_protected(self, main_module):
         # /images/generations is mounted at both /api/inference and /v1; the /v1 alias (and every
         # other /v1 POST route) must be body-capped via the /v1 blanket prefix, or an unbounded

studio/backend/tests/test_mtp_drafter_companion.py

@@ -23,7 +23,11 @@ if _BACKEND_DIR not in sys.path:
 
 from hub.utils.download_manifest import ExpectedFile
 from hub.utils.gguf import is_mtp_drafter_path
-from hub.utils.gguf_plan import build_gguf_variant_plans, plan_from_expected_files
+from hub.utils.gguf_plan import (
+    build_gguf_variant_plans,
+    plan_from_expected_files,
+    preferred_mtp_sibling,
+)
 from utils.models.model_config import (
     _is_mtp_drafter,
     detect_gguf_model,
@@ -37,6 +41,8 @@ from utils.models.model_config import (
 DRAFTER_CASES = [
     ("mtp-gemma-4-12b-it.gguf", True),
     ("MTP/gemma-4-12b-it-Q8_0-MTP.gguf", True),
+    # New-scheme MTP/ copies carry the mtp- basename prefix too.
+    ("MTP/mtp-gemma-4-E4B-it-BF16.gguf", True),
     ("foo/MTP/bar.gguf", True),
     ("gemma-4-12b-it-Q8_0.gguf", False),
     # Baked-in Qwen MTP repos: the head is inside the main GGUF, the file
@@ -274,3 +280,178 @@ def test_detect_gguf_model_rejects_mtp_subdir_copy(tmp_path):
     assert detect_gguf_model(str(copy)) is None
     # Selecting the MTP dir itself must not surface the copies as models.
     assert detect_gguf_model(str(sub)) is None
+
+
+# ── Root drafter wins over new-scheme MTP/ copies ────────────────────
+# The MTP/ copies were renamed to share the mtp- basename prefix (e.g.
+# MTP/mtp-gemma-4-E4B-it-BF16.gguf). Auto-fetch/load must still resolve the
+# small repo-root drafter, not a sort-first MTP/ copy (uppercase precedes
+# lowercase, so the subdir path would otherwise win).
+
+NEW_SCHEME_SIBLINGS = [
+    _sib("gemma-4-12b-it-Q4_K_M.gguf", 4_000, "main-q4"),
+    _sib("gemma-4-12b-it-Q8_0.gguf", 8_000, "main-q8"),
+    _sib("mtp-gemma-4-12b-it.gguf", 100, "drafter"),
+    _sib("MTP/mtp-gemma-4-12b-it-Q8_0.gguf", 100, "mtp-sub-q8"),
+    _sib("MTP/mtp-gemma-4-12b-it-BF16.gguf", 200, "mtp-sub-bf16"),
+    _sib("mmproj-F16.gguf", 500, "mmproj"),
+]
+
+
+def test_preferred_mtp_sibling_prefers_root_over_new_scheme_copies():
+    picked = preferred_mtp_sibling(NEW_SCHEME_SIBLINGS)
+    assert picked is not None and picked.rfilename == "mtp-gemma-4-12b-it.gguf"
+
+
+def test_variant_plans_new_scheme_uses_root_drafter():
+    plans = build_gguf_variant_plans(NEW_SCHEME_SIBLINGS)
+    assert set(plans) == {"q4_k_m", "q8_0"}
+    for plan in plans.values():
+        assert "mtp-gemma-4-12b-it.gguf" in plan.target_filenames
+        assert not any("MTP/" in name for name in plan.target_filenames)
+        assert "drafter" in plan.companion_hashes
+    # Download size = main + mmproj + root drafter (not the 200-byte BF16 copy).
+    assert plans["q4_k_m"].download_size_bytes == 4_600
+
+
+def test_download_mtp_prefers_root_over_new_scheme_copies(monkeypatch):
+    # _pick_mtp is nested; capture it via the companion-download seam.
+    from core.inference.llama_cpp import LlamaCppBackend
+
+    monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)  # online: skip reuse probe
+    captured = {}
+
+    def _fake_companion(
+        *,
+        hf_repo,
+        hf_token,
+        pick,
+        label,
+        cancel_event = None,
+    ):
+        captured["pick"] = pick
+        return None
+
+    b = LlamaCppBackend()
+    b._download_companion_gguf = _fake_companion
+    b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
+
+    repo_files = [
+        "MTP/mtp-gemma-4-E4B-it-BF16.gguf",
+        "MTP/mtp-gemma-4-E4B-it-Q4_0.gguf",
+        "MTP/mtp-gemma-4-E4B-it-Q8_0.gguf",
+        "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf",
+        "mmproj-F16.gguf",
+        "mtp-gemma-4-E4B-it.gguf",
+    ]
+    assert captured["pick"](repo_files) == "mtp-gemma-4-E4B-it.gguf"
+
+
+# ── Reuse an on-disk drafter offline; fetch fresh online ─────────────
+
+
+def _seed_snapshot(tmp_path, names):
+    snap = tmp_path / "snap"
+    for rel in names:
+        f = snap / rel
+        f.parent.mkdir(parents = True, exist_ok = True)
+        f.write_bytes(b"x")
+    return snap
+
+
+def test_download_mtp_reuses_cached_root_drafter_offline(tmp_path, monkeypatch):
+    import utils.models.model_config as mc
+    from core.inference.llama_cpp import LlamaCppBackend
+
+    monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+    snap = _seed_snapshot(
+        tmp_path,
+        [
+            "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf",
+            "mtp-gemma-4-E4B-it.gguf",
+            "MTP/mtp-gemma-4-E4B-it-BF16.gguf",
+            "mmproj-F16.gguf",
+        ],
+    )
+    monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap])
+
+    got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
+    assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf"
+
+
+def test_download_mtp_reuses_cached_subdir_copy_when_no_root_offline(tmp_path, monkeypatch):
+    # Pre-fix build may have fetched only the MTP/ copy; reuse it offline.
+    import utils.models.model_config as mc
+    from core.inference.llama_cpp import LlamaCppBackend
+
+    monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+    snap = _seed_snapshot(
+        tmp_path,
+        [
+            "gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf",
+            "MTP/mtp-gemma-4-E4B-it-BF16.gguf",
+        ],
+    )
+    monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap])
+
+    got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
+    assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it-BF16.gguf"
+
+
+def test_download_mtp_prefers_root_across_snapshots_offline(tmp_path, monkeypatch):
+    # A newer partial snapshot holds only the MTP/ copy; an older one has the
+    # root. Must still return the small root, not the large subdir copy.
+    import utils.models.model_config as mc
+    from core.inference.llama_cpp import LlamaCppBackend
+
+    monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+    snap_partial = _seed_snapshot(tmp_path / "new", ["MTP/mtp-gemma-4-E4B-it-BF16.gguf"])
+    snap_full = _seed_snapshot(tmp_path / "old", ["mtp-gemma-4-E4B-it.gguf"])
+    monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap_partial, snap_full])
+
+    got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
+    assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf"
+
+
+def test_download_mtp_reuse_follows_snapshot_order_offline(tmp_path, monkeypatch):
+    # Two snapshots both hold a root drafter; newest-first order must win so a
+    # fresh main GGUF is not paired with a stale drafter revision.
+    import utils.models.model_config as mc
+    from core.inference.llama_cpp import LlamaCppBackend
+
+    monkeypatch.setenv("HF_HUB_OFFLINE", "1")
+    newest = _seed_snapshot(tmp_path / "newest", ["mtp-gemma-4-E4B-it.gguf"])
+    oldest = _seed_snapshot(tmp_path / "oldest", ["mtp-gemma-4-E4B-it.gguf"])
+    monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [newest, oldest])
+
+    got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
+    assert got is not None and Path(got).parent.parent.name == "newest"
+
+
+def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch):
+    # Online, do not reuse a cached copy: go to the download path so a changed
+    # drafter is refetched (hf_hub_download checks the current revision).
+    import utils.models.model_config as mc
+    from core.inference.llama_cpp import LlamaCppBackend
+
+    monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
+    snap = _seed_snapshot(tmp_path, ["mtp-gemma-4-E4B-it.gguf"])
+    monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap])
+
+    reached = {}
+
+    def _fake_companion(
+        *,
+        hf_repo,
+        hf_token,
+        pick,
+        label,
+        cancel_event = None,
+    ):
+        reached["hit"] = True
+        return None
+
+    b = LlamaCppBackend()
+    b._download_companion_gguf = _fake_companion
+    assert b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") is None
+    assert reached.get("hit") is True

studio/backend/tests/test_passthrough_healing.py

@@ -515,6 +515,7 @@ class ScriptedClient:
         _url,
         json = None,
         timeout = None,
+        headers = None,
     ):
         self.posts.append(json)
         return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)])

studio/backend/tests/test_training_start_offload.py

@@ -0,0 +1,115 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+"""/api/train/start must run backend.start_training off the event loop.
+
+start_training() runs the _free_vram_for_training before_spawn hook inline, and that
+hook's diffusion/video unload() blocks on the engines' generation locks until an
+in-flight denoise step reaches its cancel callback (seconds to tens of seconds for
+video). Executed inline in the async route it would freeze every concurrent
+status/cancel/UI request -- the same reason start_diffusion_training offloads
+_free_gpu_for_diffusion_training via asyncio.to_thread. The backend guards the
+overlapping-starts window this offload opens with a compare-and-set flag.
+"""
+
+import asyncio
+import threading
+
+import routes.training as tr
+from models import TrainingStartRequest
+
+
+class _FakeBackend:
+    def __init__(self, result = True):
+        self._result = result
+        self.start_thread = None
+        self.hook = None
+        self.current_job_id = None
+
+    def is_training_active(self):
+        return False
+
+    def start_training(
+        self,
+        job_id,
+        *,
+        before_spawn = None,
+        **kwargs,
+    ):
+        # The real backend runs before_spawn synchronously inside this call, so the
+        # thread this method runs on is the thread the blocking VRAM hook runs on.
+        self.start_thread = threading.current_thread()
+        self.hook = before_spawn
+        self.current_job_id = job_id
+        return self._result
+
+
+def _request() -> TrainingStartRequest:
+    return TrainingStartRequest(
+        model_name = "unsloth/tiny-model",
+        training_type = "LoRA/QLoRA",
+        format_type = "alpaca",
+        hf_dataset = "org/data",
+        # Skip the YAML trust_remote_code lookup (needs the model catalog on disk).
+        trust_remote_code = True,
+    )
+
+
+def test_start_route_offloads_blocking_start(monkeypatch):
+    fake = _FakeBackend()
+    monkeypatch.setattr(tr, "get_training_backend", lambda: fake)
+    monkeypatch.setattr(tr, "_diffusion_training_active", lambda: False)
+
+    async def _run():
+        return threading.current_thread(), await tr.start_training(
+            request = _request(), current_subject = "test-user", via_api_key = False
+        )
+
+    loop_thread, resp = asyncio.run(_run())
+
+    assert resp.status == "queued", resp
+    # The VRAM-freeing hook was wired in and the blocking call left the loop thread.
+    assert fake.hook is not None
+    assert fake.start_thread is not None
+    assert fake.start_thread is not loop_thread
+
+
+def test_backend_start_guard_blocks_overlapping_starts():
+    # With the route offloaded to worker threads, two overlapping /train/start requests
+    # can reach TrainingBackend.start_training concurrently; the compare-and-set
+    # _start_in_progress flag must let exactly one of them spawn.
+    from core.training.training import TrainingBackend
+
+    backend = TrainingBackend()
+    first_entered = threading.Event()
+    release_first = threading.Event()
+    results = {}
+
+    def _slow_impl(
+        job_id,
+        *,
+        before_spawn = None,
+        **kwargs,
+    ):
+        first_entered.set()
+        release_first.wait(timeout = 5.0)
+        return True
+
+    backend._start_training_impl = _slow_impl
+
+    def _first():
+        results["first"] = backend.start_training("job-a")
+
+    t = threading.Thread(target = _first, daemon = True)
+    t.start()
+    assert first_entered.wait(timeout = 5.0)
+    # Second start while the first is still inside the impl: refused by the guard,
+    # without ever entering the impl.
+    results["second"] = backend.start_training("job-b")
+    release_first.set()
+    t.join(timeout = 5.0)
+
+    assert results["first"] is True
+    assert results["second"] is False
+    # The flag is cleared once the winning start returns, so a later start may proceed.
+    assert backend._start_in_progress is False

studio/backend/tests/test_video_routes.py

@@ -12,6 +12,9 @@ video_gallery code.
 
 from __future__ import annotations
 
+import threading
+import time
+
 import pytest
 from fastapi import FastAPI
 from fastapi.testclient import TestClient
@@ -20,7 +23,11 @@ import core.inference.gpu_arbiter as gpu_arbiter
 import core.inference.video as video_module
 import core.inference.video_gallery as gallery_module
 from auth.authentication import get_current_subject
-from core.inference.video_families import VIDEO_CANCELLED_MSG, VIDEO_NOT_LOADED_MSG
+from core.inference.video_families import (
+    VIDEO_CANCELLED_MSG,
+    VIDEO_GENERATION_BUSY_MSG,
+    VIDEO_NOT_LOADED_MSG,
+)
 from routes.video import router as video_router
 
 
@@ -60,14 +67,29 @@ def _unloaded_status():
     }
 
 
-class _FakeBackend:
+class _FakeBackend(video_module.VideoBackend):
+    """Overrides the heavy load/generate/status surface but INHERITS the real
+    begin_generate / _run_generate / generate_progress / cancel_generate job
+    machinery, so the asynchronous generate contract (immediate accept, busy
+    guard, terminal completed/failed state, cancel) is exercised for real."""
+
     def __init__(self) -> None:
-        self.loaded = False
+        super().__init__()
         self.last_load_kwargs: dict = {}
         # Repo ids of in-flight (not yet committed) loads; empty tuple = none. The unload
         # route reads this to keep VIDEO ownership while a concurrent load is still loading.
         self.loading: tuple = ()
 
+    # The real backend keys "loaded" off its committed pipeline state (_state); map
+    # the fake's flag onto it so the inherited begin_generate sees the same thing.
+    @property
+    def loaded(self) -> bool:
+        return self._state is not None
+
+    @loaded.setter
+    def loaded(self, value: bool) -> None:
+        self._state = object() if value else None
+
     def loading_repo_ids(self) -> tuple:
         return tuple(self.loading)
 
@@ -131,6 +153,7 @@ class _FakeBackend:
         *,
         prompt,
         seed = None,
+        cancel_event = None,
         **kwargs,
     ):
         if not self.loaded:
@@ -149,12 +172,6 @@ class _FakeBackend:
             "guidance": 4.0 if kwargs.get("guidance") is None else kwargs.get("guidance"),
         }
 
-    def generate_progress(self):
-        return {"active": False}
-
-    def cancel_generate(self):
-        return False
-
     def unload(self):
         self.loaded = False
         return _unloaded_status()
@@ -204,6 +221,33 @@ def client(monkeypatch, tmp_path):
     return TestClient(app)
 
 
+def _wait_terminal(client, timeout = 5.0) -> dict:
+    """Poll generate-progress until the background job records a terminal phase.
+    Generation is asynchronous now (the POST returns as soon as the job starts),
+    so its outcome is only observable here."""
+    deadline = time.monotonic() + timeout
+    progress: dict = {}
+    while time.monotonic() < deadline:
+        progress = client.get("/api/inference/video/generate-progress").json()
+        if progress.get("phase") in ("completed", "failed"):
+            return progress
+        time.sleep(0.01)
+    raise AssertionError(f"generation never reached a terminal state: {progress}")
+
+
+def _generate_and_wait(client, payload) -> dict:
+    """Start a generation, assert the immediate accepted response, and return the
+    saved gallery record the completed progress state carries."""
+    resp = client.post("/api/inference/video/generate", json = payload)
+    assert resp.status_code == 200
+    body = resp.json()
+    assert body["status"] == "started" and body["video"] is None
+    progress = _wait_terminal(client)
+    assert progress["phase"] == "completed", progress
+    assert progress["active"] is False and progress["error"] is None
+    return progress["video"]
+
+
 def test_load_happy_path_and_arbiter_acquired(client, monkeypatch):
     # Force the device to cuda so the load takes the GPU arbiter, and record the acquire.
     import types
@@ -278,11 +322,8 @@ def test_load_threads_transformer_quant_and_guidance_2(client):
     kwargs = video_module.get_video_backend().last_load_kwargs
     assert kwargs.get("transformer_quant") == "fp8"
 
-    gen = client.post(
-        "/api/inference/video/generate",
-        json = {"prompt": "a sloth", "guidance": 5.0, "guidance_2": 3.0},
-    )
-    assert gen.status_code == 200
+    video = _generate_and_wait(client, {"prompt": "a sloth", "guidance": 5.0, "guidance_2": 3.0})
+    assert video["guidance"] == 5.0 and video["guidance_2"] == 3.0
 
 
 def test_load_rejects_bad_transformer_quant_422(client):
@@ -338,16 +379,14 @@ def test_load_progress_route(client):
     assert ready.json()["phase"] == "ready"
 
 
-def test_generate_happy_path_persists_and_returns_record(client):
+def test_generate_happy_path_persists_and_reports_record(client):
     client.post(
         "/api/inference/video/load",
         json = {"model_path": "unsloth/LTX-2.3-GGUF", "gguf_filename": "q.gguf"},
     )
-    gen = client.post(
-        "/api/inference/video/generate", json = {"prompt": "a sloth surfing", "seed": 7}
-    )
-    assert gen.status_code == 200
-    video = gen.json()["video"]
+    # The POST returns at once ("started"); the saved record arrives through the
+    # generate-progress terminal state (asserted inside the helper).
+    video = _generate_and_wait(client, {"prompt": "a sloth surfing", "seed": 7})
     assert video["seed"] == 7 and video["prompt"] == "a sloth surfing" and video["id"]
     assert video["has_audio"] is True
     assert video["model"] == "unsloth/LTX-2.3-GGUF"
@@ -370,7 +409,9 @@ def test_generate_without_load_returns_409(client):
     assert resp.json()["detail"] == VIDEO_NOT_LOADED_MSG
 
 
-def test_generate_cancelled_returns_409(client, monkeypatch):
+def test_generate_cancelled_reports_failed_with_sentinel(client, monkeypatch):
+    # A cancel mid-run surfaces as the job's terminal failed state carrying the exact
+    # sentinel (the frontend suppresses the toast on it), not as an HTTP error.
     backend = video_module.get_video_backend()
     backend.loaded = True
 
@@ -379,13 +420,16 @@ def test_generate_cancelled_returns_409(client, monkeypatch):
 
     monkeypatch.setattr(backend, "generate", _cancel)
     resp = client.post("/api/inference/video/generate", json = {"prompt": "p"})
-    assert resp.status_code == 409
-    assert resp.json()["detail"] == VIDEO_CANCELLED_MSG
+    assert resp.status_code == 200
+    progress = _wait_terminal(client)
+    assert progress["phase"] == "failed"
+    assert progress["error"] == VIDEO_CANCELLED_MSG
+    assert progress["active"] is False
 
 
-def test_generate_pipeline_error_returns_sanitized_500(client, monkeypatch):
-    # A loaded model that fails mid-pipeline (CUDA OOM) is a server failure: 500 with a
-    # generic message, not a 409 echoing the raw exception.
+def test_generate_pipeline_error_reports_sanitized_failure(client, monkeypatch):
+    # A loaded model that fails mid-pipeline (CUDA OOM) is a server failure: the job's
+    # terminal state carries a generic message, never the raw exception.
     backend = video_module.get_video_backend()
     backend.loaded = True
 
@@ -394,12 +438,15 @@ def test_generate_pipeline_error_returns_sanitized_500(client, monkeypatch):
 
     monkeypatch.setattr(backend, "generate", _oom)
     resp = client.post("/api/inference/video/generate", json = {"prompt": "p"})
-    assert resp.status_code == 500
-    assert resp.json()["detail"] == "Video generation failed."
-    assert "CUDA" not in resp.json()["detail"]
+    assert resp.status_code == 200
+    progress = _wait_terminal(client)
+    assert progress["phase"] == "failed"
+    assert progress["error"] == "Video generation failed."
+    assert "CUDA" not in progress["error"]
 
 
-def test_generate_value_error_returns_400(client, monkeypatch):
+def test_generate_value_error_reports_reason(client, monkeypatch):
+    # Bad client input is feedback: the terminal failed state carries the reason.
     backend = video_module.get_video_backend()
     backend.loaded = True
 
@@ -408,14 +455,48 @@ def test_generate_value_error_returns_400(client, monkeypatch):
 
     monkeypatch.setattr(backend, "generate", _bad)
     resp = client.post("/api/inference/video/generate", json = {"prompt": "p"})
-    assert resp.status_code == 400
-    assert "not supported" in resp.json()["detail"]
+    assert resp.status_code == 200
+    progress = _wait_terminal(client)
+    assert progress["phase"] == "failed"
+    assert "not supported" in progress["error"]
+
+
+def test_generate_concurrent_second_returns_409(client, monkeypatch):
+    # While a job is running, a second generate is refused synchronously with the busy
+    # sentinel; the first job still completes and persists once released.
+    backend = video_module.get_video_backend()
+    backend.loaded = True
+    release = threading.Event()
+    real_generate = _FakeBackend.generate
+
+    def _slow(**kwargs):
+        assert release.wait(5)
+        return real_generate(backend, **kwargs)
+
+    monkeypatch.setattr(backend, "generate", _slow)
+    first = client.post("/api/inference/video/generate", json = {"prompt": "a", "seed": 1})
+    assert first.status_code == 200 and first.json()["status"] == "started"
+
+    second = client.post("/api/inference/video/generate", json = {"prompt": "b"})
+    assert second.status_code == 409
+    assert second.json()["detail"] == VIDEO_GENERATION_BUSY_MSG
+
+    running = client.get("/api/inference/video/generate-progress").json()
+    assert running["active"] is True
+
+    release.set()
+    progress = _wait_terminal(client)
+    assert progress["phase"] == "completed" and progress["video"]["seed"] == 1
+    # With the job finished, a new generate is accepted again.
+    assert _generate_and_wait(client, {"prompt": "c", "seed": 2})["seed"] == 2
 
 
 def test_generate_progress_route(client):
     resp = client.get("/api/inference/video/generate-progress")
     assert resp.status_code == 200
-    assert resp.json()["active"] is False
+    body = resp.json()
+    assert body["active"] is False
+    assert body["phase"] is None and body["video"] is None and body["error"] is None
 
 
 def test_cancel_generation_route(client):
@@ -424,6 +505,31 @@ def test_cancel_generation_route(client):
     assert resp.json()["cancelled"] is False
 
 
+def test_cancel_running_job(client, monkeypatch):
+    # Cancel still works against the background job: begin_generate registers the
+    # cancel event before the worker starts, so the cancel route reports True at
+    # once and the job lands in the failed(cancelled) terminal state.
+    backend = video_module.get_video_backend()
+    backend.loaded = True
+
+    def _wait_for_cancel(*, cancel_event = None, **kwargs):
+        assert cancel_event is not None and cancel_event.wait(5)
+        raise RuntimeError(VIDEO_CANCELLED_MSG)
+
+    monkeypatch.setattr(backend, "generate", _wait_for_cancel)
+    resp = client.post("/api/inference/video/generate", json = {"prompt": "p"})
+    assert resp.status_code == 200
+
+    cancelled = client.post("/api/inference/video/generate/cancel")
+    assert cancelled.status_code == 200 and cancelled.json()["cancelled"] is True
+
+    progress = _wait_terminal(client)
+    assert progress["phase"] == "failed"
+    assert progress["error"] == VIDEO_CANCELLED_MSG
+    # Nothing was persisted for the cancelled run.
+    assert client.get("/api/inference/video/gallery").json()["videos"] == []
+
+
 def test_file_endpoint_404_for_bad_id(client):
     resp = client.get("/api/inference/video/gallery/does-not-exist/file")
     assert resp.status_code == 404
@@ -434,8 +540,8 @@ def test_delete_and_clear(client):
         "/api/inference/video/load",
         json = {"model_path": "unsloth/LTX-2.3-GGUF", "gguf_filename": "q.gguf"},
     )
-    first = client.post("/api/inference/video/generate", json = {"prompt": "a"}).json()["video"]
-    second = client.post("/api/inference/video/generate", json = {"prompt": "b"}).json()["video"]
+    first = _generate_and_wait(client, {"prompt": "a"})
+    second = _generate_and_wait(client, {"prompt": "b"})
     assert len(client.get("/api/inference/video/gallery").json()["videos"]) == 2
 
     # Delete one, then confirm it 404s and the other remains.
@@ -456,7 +562,7 @@ def test_gallery_pagination(client):
         json = {"model_path": "unsloth/LTX-2.3-GGUF", "gguf_filename": "q.gguf"},
     )
     for i in range(5):
-        client.post("/api/inference/video/generate", json = {"prompt": f"clip {i}", "seed": i})
+        _generate_and_wait(client, {"prompt": f"clip {i}", "seed": i})
     page1 = client.get("/api/inference/video/gallery?limit=2&offset=0").json()
     assert len(page1["videos"]) == 2 and page1["has_more"] is True
     last = client.get("/api/inference/video/gallery?limit=2&offset=4").json()
@@ -572,7 +678,7 @@ def test_export_endpoint_validation(client, monkeypatch):
         "/api/inference/video/load",
         json = {"model_path": "unsloth/LTX-2.3-GGUF", "gguf_filename": "q.gguf"},
     )
-    video = client.post("/api/inference/video/generate", json = {"prompt": "a"}).json()["video"]
+    video = _generate_and_wait(client, {"prompt": "a"})
     resp = client.get(f"/api/inference/video/gallery/{video['id']}/export?format=webm")
     assert resp.status_code == 501
     assert "PyAV" in resp.json()["detail"]

studio/backend/utils/llama_cpp_update.py

@@ -514,6 +514,12 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path
         logger.info("llama update: installing", cmd = " ".join(cmd))
         # Stream progress lines into job["progress"].
         env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5")
+        # Preserve a Vulkan install across updates: detect_host on a CUDA/ROCm
+        # box would otherwise re-route and silently replace the Vulkan build.
+        # Re-assert it via the same env flag setup uses (mirrors
+        # _rocm_install_args).
+        if asset and "vulkan" in asset.lower():
+            env["UNSLOTH_FORCE_VULKAN"] = "1"
         proc = subprocess.Popen(
             cmd,
             stdout = subprocess.PIPE,

studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx

@@ -2273,8 +2273,12 @@ export function HubModelPicker({
             // Gate on a curated ARTIFACT (artifactForRepoId, what loadSpecFor resolves), not a
             // group-key match: a base / uncurated-quant sibling (Qwen/Qwen-Image-2512) matches
             // the group by key but has no loadable artifact and dead-ends at the trust gate.
+            // An unsloth repo must also be a full pipeline (not single_file): the selection
+            // fall-through loads uncataloged rows as kind "pipeline", and from_pretrained on
+            // a single-file checkpoint repo (no model_index.json) fails after the handoff.
+            // Curated single-file artifacts stay: loadSpecFor carries their filename.
             (!task ||
-              isUnslothRepoId(c.repo_id) ||
+              (isUnslothRepoId(c.repo_id) && !c.single_file) ||
               (catalog ? artifactForRepoId(c.repo_id, catalog) !== null : false)),
         ),
         downloadedSort,

studio/frontend/src/features/chat/api/chat-adapter.ts

@@ -2,6 +2,10 @@
 // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
 
 import { getAuthToken } from "@/features/auth";
+import {
+  loadRememberedLoadSettings,
+  rememberedLoadSettingsKey,
+} from "@/components/assistant-ui/model-selector/remembered-load-settings";
 import { projectHasSources } from "@/features/rag/api/rag-api";
 import { apiUrl } from "@/lib/api-base";
 import { parseParamCountB } from "@/lib/model-size";
@@ -63,11 +67,17 @@ import {
   listStoredChatThreads,
   updateStoredChatThread,
 } from "../utils/chat-history-storage";
+import {
+  readLastLocalModelLoad,
+  recordLastLocalModelLoad,
+  type LastLocalModelKind,
+} from "../utils/last-local-model-load";
 import { getImageInputUnavailableReason } from "../utils/image-input-support";
 import {
   hasClosedThinkTag,
   parseAssistantContent,
 } from "../utils/parse-assistant-content";
+import { resolveLoadMaxSeqLength } from "../presets/preset-policy";
 import {
   generateAudio,
   listCachedGguf,
@@ -1309,6 +1319,30 @@ const BIG_ENDIAN_GGUF_FILENAME_RE = /(^|[-_])be(?:[._-]|$)/gi;
 const GGUF_KNOWN_QUANT_RE =
   /(UD-)?(MXFP[0-9]+(?:_[A-Z0-9]+)*|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?|TQ[0-9]+_[0-9]+|Q[0-9]+_K_[A-Z]+|Q[0-9]+_[0-9]+|Q[0-9]+_K|BF16|F16|F32)/i;
 
+type AutoLoadCandidate = {
+  id: string;
+  kind: LastLocalModelKind;
+  ggufVariant: string | null;
+  maxSeqLength: number;
+  successLabel: string;
+};
+
+function autoLoadCandidateKey(
+  kind: LastLocalModelKind,
+  id: string,
+  ggufVariant?: string | null,
+): string {
+  return `${kind}:${id.toLowerCase()}:${(ggufVariant ?? "").toLowerCase()}`;
+}
+
+function findCachedRepo<T extends { repo_id: string }>(
+  repos: T[],
+  id: string,
+): T | undefined {
+  const normalized = id.toLowerCase();
+  return repos.find((repo) => repo.repo_id.toLowerCase() === normalized);
+}
+
 function hasBigEndianGgufMarker(filename: string, quant?: string | null): boolean {
   const normalized = filename.replace(/\\/g, "/").toLowerCase();
   const separatorIndex = normalized.lastIndexOf("/");
@@ -1357,14 +1391,18 @@ async function autoLoadSmallestModel(): Promise<{
   const hfToken = store.hfToken || null;
   const trustRemoteCode = store.params.trustRemoteCode ?? false;
   const specSettings = resolveSpeculativeSettingsForLoad();
+  const lastLoaded = readLastLocalModelLoad();
   const toastId = toast("Loading a model…", {
-    description: "Auto-selecting the smallest downloaded model.",
+    description: lastLoaded
+      ? "Loading last used model."
+      : "Auto-selecting the smallest downloaded model.",
     duration: 5000,
     closeButton: true,
   });
   let blockedByTrustRemoteCode = false;
   let hadNonTrustFailure = false;
   let loadAttempts = 0;
+  const skippedAutoLoadCandidates = new Set<string>();
 
   async function canAutoLoad(payload: {
     model_path: string;
@@ -1389,12 +1427,224 @@ async function autoLoadSmallestModel(): Promise<{
     }
     return true;
   }
+
+  async function loadAutoLoadCandidate(
+    candidate: AutoLoadCandidate,
+  ): Promise<boolean> {
+    if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) {
+      return false;
+    }
+    const currentStore = useChatRuntimeStore.getState();
+    const remembered = loadRememberedLoadSettings(
+      rememberedLoadSettingsKey({
+        id: candidate.id,
+        ggufVariant: candidate.ggufVariant,
+      }),
+    );
+    const effectiveMaxSeqLength = resolveLoadMaxSeqLength({
+      modelId: candidate.id,
+      ggufVariant: candidate.ggufVariant,
+      isGguf: candidate.kind === "gguf",
+      customContextLength: remembered?.contextLength ?? null,
+      ggufContextLength: null,
+      currentCheckpoint: currentStore.params.checkpoint,
+      activeGgufVariant: currentStore.activeGgufVariant,
+      maxSeqLength: candidate.maxSeqLength,
+      presetSource: currentStore.activePresetSource,
+    });
+    const effectiveSpeculativeType =
+      remembered?.speculativeType ?? specSettings.speculativeType;
+    const effectiveSpecDraftNMax =
+      remembered?.specDraftNMax ?? specSettings.specDraftNMax;
+    if (
+      !(await canAutoLoad({
+        model_path: candidate.id,
+        max_seq_length: effectiveMaxSeqLength,
+        is_lora: false,
+        gguf_variant: candidate.ggufVariant,
+      }))
+    ) {
+      skippedAutoLoadCandidates.add(
+        autoLoadCandidateKey(candidate.kind, candidate.id, candidate.ggufVariant),
+      );
+      return false;
+    }
+    loadAttempts += 1;
+    const loadResp = await loadModel({
+      model_path: candidate.id,
+      hf_token: hfToken,
+      max_seq_length: effectiveMaxSeqLength,
+      load_in_4bit: true,
+      is_lora: false,
+      gguf_variant: candidate.ggufVariant,
+      trust_remote_code: trustRemoteCode,
+      cache_type_kv: remembered?.kvCacheDtype ?? null,
+      speculative_type: effectiveSpeculativeType,
+      spec_draft_n_max: effectiveSpecDraftNMax,
+      tensor_parallel: remembered?.tensorParallel ?? false,
+    });
+    saveSpeculativeType(effectiveSpeculativeType);
+    useChatRuntimeStore
+      .getState()
+      .setCheckpoint(candidate.id, candidate.ggufVariant ?? undefined);
+    const store = useChatRuntimeStore.getState();
+    store.setModelRequiresTrustRemoteCode(
+      loadResp.requires_trust_remote_code ?? false,
+    );
+    store.setParams({
+      ...store.params,
+      maxTokens:
+        candidate.kind === "gguf"
+          ? loadResp.context_length ?? 131072
+          : effectiveMaxSeqLength,
+    });
+    const autoModel: ChatModelSummary = {
+      id: candidate.id,
+      name: loadResp.display_name ?? candidate.id,
+      isVision: loadResp.is_vision ?? false,
+      isLora: loadResp.is_lora ?? false,
+      isGguf: loadResp.is_gguf ?? candidate.kind === "gguf",
+      isAudio: loadResp.is_audio ?? false,
+      audioType: loadResp.audio_type ?? null,
+      hasAudioInput: loadResp.has_audio_input ?? false,
+    };
+    if (!store.models.some((m) => m.id === candidate.id)) {
+      store.setModels([...store.models, autoModel]);
+    }
+    if (candidate.kind === "gguf") {
+      useChatRuntimeStore.setState({
+        ggufContextLength: loadResp.context_length ?? 131072,
+        ggufMaxContextLength:
+          loadResp.max_context_length ?? loadResp.context_length ?? 131072,
+        ggufNativeContextLength: loadResp.native_context_length ?? null,
+        supportsReasoning: loadResp.supports_reasoning ?? false,
+        reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
+        reasoningEnabled: loadResp.supports_reasoning ?? false,
+        ...reasoningCapsFromLoad(loadResp),
+        supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false,
+        supportsTools: loadResp.supports_tools ?? false,
+        ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false),
+        kvCacheDtype: loadResp.cache_type_kv ?? null,
+        loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
+        tensorParallel: loadResp.tensor_parallel ?? false,
+        loadedTensorParallel: loadResp.tensor_parallel ?? false,
+        defaultChatTemplate: loadResp.chat_template ?? null,
+        chatTemplateOverride: null,
+        loadedChatTemplateOverride: null,
+        loadedIsMultimodal: isMultimodalResponse(loadResp),
+        loadedIsDiffusion: loadResp.is_diffusion ?? false,
+        ...resolveLoadedSpeculativeSettings(loadResp),
+      });
+    } else {
+      useChatRuntimeStore.setState({
+        supportsReasoning: loadResp.supports_reasoning ?? false,
+        reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
+        reasoningEnabled: loadResp.supports_reasoning ?? false,
+        ...reasoningCapsFromLoad(loadResp),
+        supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false,
+        supportsTools: loadResp.supports_tools ?? false,
+        ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false),
+        kvCacheDtype: loadResp.cache_type_kv ?? null,
+        loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
+        tensorParallel: loadResp.tensor_parallel ?? false,
+        loadedTensorParallel: loadResp.tensor_parallel ?? false,
+        defaultChatTemplate: loadResp.chat_template ?? null,
+        chatTemplateOverride: null,
+        loadedChatTemplateOverride: null,
+        ...resolveLoadedSpeculativeSettings(loadResp),
+        loadedIsMultimodal: isMultimodalResponse(loadResp),
+        loadedIsDiffusion: loadResp.is_diffusion ?? false,
+      });
+    }
+    if (!(loadResp.is_lora ?? false)) {
+      recordLastLocalModelLoad({
+        id: candidate.id,
+        kind: candidate.kind,
+        ggufVariant: candidate.ggufVariant,
+      });
+    }
+    toast.success(candidate.successLabel, { id: toastId });
+    return true;
+  }
   try {
     const [ggufRepos, modelRepos] = await Promise.all([
       listCachedGguf().catch(() => []),
       listCachedModels().catch(() => []),
     ]);
 
+    if (lastLoaded) {
+      if (lastLoaded.kind === "gguf") {
+        const repo = findCachedRepo(ggufRepos, lastLoaded.id);
+        if (repo && lastLoaded.ggufVariant) {
+          try {
+            const variants = await listGgufVariants(repo.repo_id);
+            const variant = variants.variants.find(
+              (entry) =>
+                entry.downloaded &&
+                entry.quant?.toLowerCase() ===
+                  lastLoaded.ggufVariant?.toLowerCase() &&
+                isAutoLoadableGgufVariant(entry),
+            );
+            if (variant) {
+              toast("Loading last used model…", {
+                id: toastId,
+                description: `${repo.repo_id} (${variant.quant})`,
+                duration: 5000,
+              });
+              if (
+                await loadAutoLoadCandidate({
+                  id: repo.repo_id,
+                  kind: "gguf",
+                  ggufVariant: variant.quant,
+                  maxSeqLength: 0,
+                  successLabel: `Loaded ${repo.repo_id} (${variant.quant})`,
+                })
+              ) {
+                return { loaded: true, blockedByTrustRemoteCode: false };
+              }
+            }
+          } catch {
+            hadNonTrustFailure = true;
+            skippedAutoLoadCandidates.add(
+              autoLoadCandidateKey("gguf", repo.repo_id, lastLoaded.ggufVariant),
+            );
+          }
+        }
+      } else {
+        const repo = findCachedRepo(modelRepos, lastLoaded.id);
+        if (repo) {
+          try {
+            toast("Loading last used model…", {
+              id: toastId,
+              description: repo.repo_id,
+              duration: 5000,
+            });
+            if (
+              await loadAutoLoadCandidate({
+                id: repo.repo_id,
+                kind: "model",
+                ggufVariant: null,
+                maxSeqLength: store.params.maxSeqLength,
+                successLabel: `Loaded ${repo.repo_id}`,
+              })
+            ) {
+              return { loaded: true, blockedByTrustRemoteCode: false };
+            }
+          } catch {
+            hadNonTrustFailure = true;
+            skippedAutoLoadCandidates.add(
+              autoLoadCandidateKey("model", repo.repo_id),
+            );
+          }
+        }
+      }
+      toast("Loading a model…", {
+        id: toastId,
+        description: "Auto-selecting the smallest downloaded model.",
+        duration: 5000,
+      });
+    }
+
     // GGUF first: smallest-total-size repo, then its smallest variant.
     if (ggufRepos.length > 0) {
       const sorted = [...ggufRepos].sort((a, b) => a.size_bytes - b.size_bytes);
@@ -1408,82 +1658,23 @@ async function autoLoadSmallestModel(): Promise<{
           if (downloaded.length > 0) {
             const variant = downloaded[0];
             if (
-              !(await canAutoLoad({
-                model_path: repo.repo_id,
-                max_seq_length: 0,
-                is_lora: false,
-                gguf_variant: variant.quant,
-              }))
+              skippedAutoLoadCandidates.has(
+                autoLoadCandidateKey("gguf", repo.repo_id, variant.quant),
+              )
             ) {
               continue;
             }
-            loadAttempts += 1;
-            const loadResp = await loadModel({
-              model_path: repo.repo_id,
-              hf_token: hfToken,
-              max_seq_length: 0,
-              load_in_4bit: true,
-              is_lora: false,
-              gguf_variant: variant.quant,
-              trust_remote_code: trustRemoteCode,
-              speculative_type: specSettings.speculativeType,
-              spec_draft_n_max: specSettings.specDraftNMax,
-            });
-            saveSpeculativeType(specSettings.speculativeType);
-            useChatRuntimeStore
-              .getState()
-              .setCheckpoint(repo.repo_id, variant.quant);
-            const store = useChatRuntimeStore.getState();
-            store.setModelRequiresTrustRemoteCode(
-              loadResp.requires_trust_remote_code ?? false,
-            );
-            store.setParams({
-              ...store.params,
-              maxTokens: loadResp.context_length ?? 131072,
-            });
-            // Add to store so the selector shows the name.
-            const autoModel: ChatModelSummary = {
-              id: repo.repo_id,
-              name: loadResp.display_name ?? repo.repo_id,
-              isVision: loadResp.is_vision ?? false,
-              isLora: loadResp.is_lora ?? false,
-              isGguf: loadResp.is_gguf ?? false,
-              isAudio: loadResp.is_audio ?? false,
-              audioType: loadResp.audio_type ?? null,
-              hasAudioInput: loadResp.has_audio_input ?? false,
-            };
-            const existingModels = store.models;
-            if (!existingModels.some((m) => m.id === repo.repo_id)) {
-              store.setModels([...existingModels, autoModel]);
+            if (
+              await loadAutoLoadCandidate({
+                id: repo.repo_id,
+                kind: "gguf",
+                ggufVariant: variant.quant,
+                maxSeqLength: 0,
+                successLabel: `Loaded ${repo.repo_id} (${variant.quant})`,
+              })
+            ) {
+              return { loaded: true, blockedByTrustRemoteCode: false };
             }
-            useChatRuntimeStore.setState({
-              ggufContextLength: loadResp.context_length ?? 131072,
-              ggufMaxContextLength:
-                loadResp.max_context_length ??
-                loadResp.context_length ??
-                131072,
-              supportsReasoning: loadResp.supports_reasoning ?? false,
-              reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
-              reasoningEnabled: loadResp.supports_reasoning ?? false,
-              ...reasoningCapsFromLoad(loadResp),
-              supportsPreserveThinking:
-                loadResp.supports_preserve_thinking ?? false,
-              supportsTools: loadResp.supports_tools ?? false,
-              ...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false),
-              kvCacheDtype: loadResp.cache_type_kv ?? null,
-              loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
-              tensorParallel: loadResp.tensor_parallel ?? false,
-              loadedTensorParallel: loadResp.tensor_parallel ?? false,
-              defaultChatTemplate: loadResp.chat_template ?? null,
-              chatTemplateOverride: null,
-              loadedChatTemplateOverride: null,
-              loadedIsMultimodal: isMultimodalResponse(loadResp),
-              ...resolveLoadedSpeculativeSettings(loadResp),
-            });
-            toast.success(`Loaded ${repo.repo_id} (${variant.quant})`, {
-              id: toastId,
-            });
-            return { loaded: true, blockedByTrustRemoteCode: false };
           }
         } catch {
           hadNonTrustFailure = true;
@@ -1501,64 +1692,23 @@ async function autoLoadSmallestModel(): Promise<{
         if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break;
         try {
           if (
-            !(await canAutoLoad({
-              model_path: repo.repo_id,
-              max_seq_length: 4096,
-              is_lora: false,
-              gguf_variant: null,
-            }))
+            skippedAutoLoadCandidates.has(
+              autoLoadCandidateKey("model", repo.repo_id),
+            )
           ) {
             continue;
           }
-          loadAttempts += 1;
-          const sfLoadResp = await loadModel({
-            model_path: repo.repo_id,
-            hf_token: hfToken,
-            max_seq_length: 4096,
-            load_in_4bit: true,
-            is_lora: false,
-            gguf_variant: null,
-            trust_remote_code: trustRemoteCode,
-            speculative_type: specSettings.speculativeType,
-            spec_draft_n_max: specSettings.specDraftNMax,
-          });
-          saveSpeculativeType(specSettings.speculativeType);
-          useChatRuntimeStore.getState().setCheckpoint(repo.repo_id);
-          const store = useChatRuntimeStore.getState();
-          store.setModelRequiresTrustRemoteCode(
-            sfLoadResp.requires_trust_remote_code ?? false,
-          );
-          store.setParams({ ...store.params, maxTokens: 4096 });
-          useChatRuntimeStore.setState({
-            supportsReasoning: sfLoadResp.supports_reasoning ?? false,
-            reasoningAlwaysOn: sfLoadResp.reasoning_always_on ?? false,
-            reasoningEnabled: sfLoadResp.supports_reasoning ?? false,
-            ...reasoningCapsFromLoad(sfLoadResp),
-            supportsPreserveThinking:
-              sfLoadResp.supports_preserve_thinking ?? false,
-            supportsTools: sfLoadResp.supports_tools ?? false,
-            // Parity with the GGUF branch above.
-            ...resolveToolsEnabledOnLoad(sfLoadResp.supports_tools ?? false),
-            defaultChatTemplate: sfLoadResp.chat_template ?? null,
-            chatTemplateOverride: null,
-            loadedChatTemplateOverride: null,
-            ...resolveLoadedSpeculativeSettings(sfLoadResp),
-          });
-          const sfModel: ChatModelSummary = {
-            id: repo.repo_id,
-            name: sfLoadResp.display_name ?? repo.repo_id,
-            isVision: sfLoadResp.is_vision ?? false,
-            isLora: sfLoadResp.is_lora ?? false,
-            isGguf: sfLoadResp.is_gguf ?? false,
-          };
-          if (!store.models.some((m) => m.id === repo.repo_id)) {
-            store.setModels([...store.models, sfModel]);
+          if (
+            await loadAutoLoadCandidate({
+              id: repo.repo_id,
+              kind: "model",
+              ggufVariant: null,
+              maxSeqLength: 4096,
+              successLabel: `Loaded ${repo.repo_id}`,
+            })
+          ) {
+            return { loaded: true, blockedByTrustRemoteCode: false };
           }
-          useChatRuntimeStore.setState({
-            loadedIsMultimodal: isMultimodalResponse(sfLoadResp),
-          });
-          toast.success(`Loaded ${repo.repo_id}`, { id: toastId });
-          return { loaded: true, blockedByTrustRemoteCode: false };
         } catch {
           hadNonTrustFailure = true;
           continue;
@@ -1650,6 +1800,11 @@ async function autoLoadSmallestModel(): Promise<{
         loadedIsMultimodal: isMultimodalResponse(loadResp),
         ...resolveLoadedSpeculativeSettings(loadResp),
       });
+      recordLastLocalModelLoad({
+        id: "unsloth/Qwen3.5-4B-MTP-GGUF",
+        kind: "gguf",
+        ggufVariant: "UD-Q4_K_XL",
+      });
       toast.success("Loaded Qwen3.5-4B-MTP (UD-Q4_K_XL)", { id: toastId });
       return { loaded: true, blockedByTrustRemoteCode: false };
     } catch {

studio/frontend/src/features/chat/api/chat-api.ts

@@ -326,6 +326,10 @@ export interface CachedModelRepo {
   /** True when the snapshot is incomplete (a cancelled/partial download). Such a
    * repo must not count as downloaded, or a click re-downloads the full weights. */
   partial?: boolean;
+  /** True for a diffusion repo with no model_index.json: a single-file checkpoint that
+   * loads only via from_single_file + a checkpoint filename. Task pickers must not offer
+   * it as a pipeline load unless the curated catalog carries its artifact. */
+  single_file?: boolean;
 }
 
 export async function listCachedModels(

studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts

@@ -44,6 +44,7 @@ import {
   mergeBackendRecommendedInference,
   resolveLoadMaxSeqLength,
 } from "../presets/preset-policy";
+import { recordLastLocalModelLoad } from "../utils/last-local-model-load";
 import {
   isMultimodalResponse,
 } from "../types/api";
@@ -818,6 +819,23 @@ export function useChatModelRuntime() {
               }
             }
             await refresh({ signal: abortCtrl.signal });
+            if (
+              !isLora &&
+              !(loadResponse.is_lora ?? false) &&
+              !nativePathToken &&
+              !isLocalModelPath(modelId) &&
+              !isExternalModelId(modelId)
+            ) {
+              if (loadResponse.is_gguf || isGguf || ggufVariant) {
+                recordLastLocalModelLoad({
+                  id: modelId,
+                  kind: "gguf",
+                  ggufVariant: ggufVariant ?? null,
+                });
+              } else {
+                recordLastLocalModelLoad({ id: modelId, kind: "model" });
+              }
+            }
             // A successful load owns the shared (pick-unscoped) settings fields,
             // so any surviving stage is stale: the just-loaded pick itself, or a
             // pick queued for a different model mid-load whose knobs this load

studio/frontend/src/features/chat/utils/last-local-model-load.ts

@@ -0,0 +1,86 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+
+export type LastLocalModelKind = "gguf" | "model";
+
+export type LastLocalModelLoad = {
+  id: string;
+  kind: LastLocalModelKind;
+  ggufVariant: string | null;
+  loadedAt: number;
+};
+
+const STORAGE_KEY = "unsloth.last-local-model-load.v1";
+
+function storage(): Storage | null {
+  try {
+    return typeof localStorage === "undefined" ? null : localStorage;
+  } catch {
+    return null;
+  }
+}
+
+function isLastLocalModelKind(value: unknown): value is LastLocalModelKind {
+  return value === "gguf" || value === "model";
+}
+
+export function readLastLocalModelLoad(): LastLocalModelLoad | null {
+  try {
+    const raw = storage()?.getItem(STORAGE_KEY);
+    if (!raw) {
+      return null;
+    }
+    const parsed = JSON.parse(raw) as Partial<LastLocalModelLoad>;
+    if (
+      typeof parsed.id !== "string" ||
+      !parsed.id.trim() ||
+      !isLastLocalModelKind(parsed.kind) ||
+      typeof parsed.loadedAt !== "number"
+    ) {
+      return null;
+    }
+    if (
+      parsed.kind === "gguf" &&
+      (typeof parsed.ggufVariant !== "string" || !parsed.ggufVariant.trim())
+    ) {
+      return null;
+    }
+    return {
+      id: parsed.id,
+      kind: parsed.kind,
+      ggufVariant:
+        typeof parsed.ggufVariant === "string" ? parsed.ggufVariant : null,
+      loadedAt: parsed.loadedAt,
+    };
+  } catch {
+    return null;
+  }
+}
+
+export function recordLastLocalModelLoad(input: {
+  id: string;
+  kind: LastLocalModelKind;
+  ggufVariant?: string | null;
+}): void {
+  const id = input.id.trim();
+  if (!id) {
+    return;
+  }
+  const ggufVariant = input.ggufVariant?.trim() || null;
+  if (input.kind === "gguf" && !ggufVariant) {
+    return;
+  }
+  try {
+    storage()?.setItem(
+      STORAGE_KEY,
+      JSON.stringify({
+        id,
+        kind: input.kind,
+        ggufVariant: input.kind === "gguf" ? ggufVariant : null,
+        loadedAt: Date.now(),
+      } satisfies LastLocalModelLoad),
+    );
+  } catch {
+    // Ignore disabled storage / quota errors; auto-load falls back to size order.
+  }
+}

studio/frontend/src/features/images/images-page.tsx

@@ -321,6 +321,10 @@ function formatTimestamp(epochSeconds: number): string {
 // Bar label for an in-flight generation: step count plus an ETA once it's known
 // (formatEta returns "" for non-positive, so the last step shows just the step).
 function genStepLabel(p: DiffusionGenerateProgress): string {
+  // Text encoding (and any first-run warmup) happens before the first scheduler
+  // tick, so step 0 means "working, not denoising yet" -- label it that way
+  // instead of sitting on "Step 0/N".
+  if (p.step === 0) return "Preparing (text encoding + warmup)…";
   const base = `Step ${p.step}/${p.total_steps}`;
   const eta = p.eta_seconds != null ? formatEta(p.eta_seconds) : "";
   return eta ? `${base} · ~${eta}` : base;
@@ -1042,6 +1046,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
   // Live per-step progress (step / total + ETA) polled during generation.
   const [genStep, setGenStep] = useState<DiffusionGenerateProgress | null>(null);
   const genPollTimer = useRef<ReturnType<typeof setInterval> | null>(null);
+  // visibilitychange handler active while a generation poll runs: background tabs clamp
+  // setInterval to >=1s (and can suspend it outright after ~5 min), so returning to the
+  // tab fires one immediate poll instead of waiting for a throttled tick.
+  const genVisibilityListener = useRef<(() => void) | null>(null);
   const [status, setStatus] = useState<DiffusionStatus | null>(null);
   // Controlled so the body-portaled overlays force-close when this page is mounted
   // but off-tab (a hidden/inert parent can't contain a body portal): the model
@@ -1473,6 +1481,10 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
     return () => {
       if (pollTimer.current) clearTimeout(pollTimer.current);
       if (genPollTimer.current) clearInterval(genPollTimer.current);
+      if (genVisibilityListener.current) {
+        document.removeEventListener("visibilitychange", genVisibilityListener.current);
+        genVisibilityListener.current = null;
+      }
       dismissLoadToast();
     };
   }, [refreshStatus, dismissLoadToast, pollLoadProgress]);
@@ -1648,6 +1660,29 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
         });
         return;
       }
+      // A direct local single-file .safetensors pick (custom folder / on-device file)
+      // must load via from_single_file: the pipeline route rejects a bare file (no
+      // model_index.json) and only after evicting the resident model. Split into
+      // (parent dir, basename) exactly like the local GGUF branch above.
+      if (meta.source === "local" && id.toLowerCase().endsWith(".safetensors")) {
+        const norm = id.replace(/\\/g, "/");
+        const slash = norm.lastIndexOf("/");
+        const filename = slash >= 0 ? norm.slice(slash + 1) : norm;
+        const dir = slash >= 0 ? norm.slice(0, slash) : ".";
+        const prevQuant = quant;
+        quantRevert.current = { prev: prevQuant };
+        setQuant(filename);
+        const dsf = defaultsFor(id);
+        setSteps(dsf.steps);
+        setGuidance(dsf.guidance);
+        void handleLoad(dir, { kind: "single_file", filename }).then((started) => {
+          if (!started) {
+            setQuant(prevQuant);
+            quantRevert.current = null;
+          }
+        });
+        return;
+      }
       // Otherwise treat it as a full diffusers repo (safetensors / bnb-4bit). The backend
       // infers the family + base repo from the id and gates loads to unsloth/* repos or
       // on-device paths, so only attempt those; other Hub orgs can't be assembled here.
@@ -1827,12 +1862,26 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
     const runs = Number.isFinite(count) && count >= 1 ? Math.floor(count) : 1;
     if (runs !== count) setCount(runs);
 
+    // An explicit seed near the 2**53-1 backend cap can overflow once the per-run
+    // offset (base + i*batchSize) and the engine's in-batch +j offsets are added,
+    // 422ing a later run AFTER earlier images already generated. Fail before any
+    // GPU work. Subtraction keeps the comparison exact where the sum would round.
+    if (baseSeed > Number.MAX_SAFE_INTEGER - (runs * batchSize - 1)) {
+      toast.error("Seed too large for this run count and batch size; use a smaller seed");
+      return;
+    }
+
     setBusy("generating");
     setGenDone(0);
     setGenStep(null);
     // Poll the backend's per-step progress across the whole run (all sequential
-    // generations), so the bar tracks the live denoising steps.
-    genPollTimer.current = setInterval(async () => {
+    // generations), so the bar tracks the live denoising steps. A named poll body
+    // (guarded against overlap) also serves the visibilitychange listener: a
+    // background tab's throttled interval catches up the moment the tab is visible.
+    let pollInFlight = false;
+    const pollGenerateOnce = async () => {
+      if (pollInFlight) return;
+      pollInFlight = true;
       try {
         const p = await getGenerateProgress();
         // Skip the state update (and re-render) when nothing the bar shows moved.
@@ -1843,8 +1892,17 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
         });
       } catch {
         // transient; keep polling
+      } finally {
+        pollInFlight = false;
       }
-    }, 300);
+    };
+    if (genVisibilityListener.current)
+      document.removeEventListener("visibilitychange", genVisibilityListener.current);
+    genVisibilityListener.current = () => {
+      if (document.visibilityState === "visible") void pollGenerateOnce();
+    };
+    document.addEventListener("visibilitychange", genVisibilityListener.current);
+    genPollTimer.current = setInterval(() => void pollGenerateOnce(), 300);
     try {
       for (let i = 0; i < runs; i++) {
         // The page truly unmounted mid-run (app close / chat-only eject): stop
@@ -1903,16 +1961,25 @@ export function ImagesPage({ active = true }: { active?: boolean }) {
         res.images.forEach((image) => void ensureSrc(image));
         setGenDone(i + 1);
       }
+      // A generation can change server-side status: Speed=Auto compiles the
+      // transformer on the 3rd LoRA-free run (supports_lora flips to false), so
+      // without a refresh the LoRA picker stays enabled and the next LoRA run
+      // fails on the backend. Cheap status GET; also picks up any other drift.
+      if (isMounted.current) void refreshStatus();
     } catch (err) {
       toast.error(err instanceof Error ? err.message : "Image generation failed");
     } finally {
       if (genPollTimer.current) clearInterval(genPollTimer.current);
       genPollTimer.current = null;
+      if (genVisibilityListener.current) {
+        document.removeEventListener("visibilitychange", genVisibilityListener.current);
+        genVisibilityListener.current = null;
+      }
       setBusy(null);
       setGenDone(null);
       setGenStep(null);
     }
-  }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, workflow, initImage, maskImage, strength, extendPct, extendSides, upscaleFactor, upscaleStrength, referenceImages, loras, controlnetCapable, controlnetId, controlImage, controlType, controlStrength, ensureSrc]);
+  }, [prompt, negativePrompt, width, height, steps, guidance, seed, batchSize, count, workflow, initImage, maskImage, strength, extendPct, extendSides, upscaleFactor, upscaleStrength, referenceImages, loras, controlnetCapable, controlnetId, controlImage, controlType, controlStrength, ensureSrc, refreshStatus]);
 
   // Keep the active workflow valid for the loaded model: an edit-only model (Qwen-Image-
   // Edit) has no Create/Transform tabs, a base model has no Edit tab. Snap to the first

studio/frontend/src/features/images/train/diffusion-train-panel.tsx

@@ -612,10 +612,21 @@ export function DiffusionTrainPanel({
       toast.error("Name the adapter (this becomes its folder under Studio outputs).");
       return;
     }
-    if (selectedDataset && selectedDataset.caption_count === 0 && !instancePrompt.trim()) {
+    // Require a trigger prompt whenever ANY image lacks a caption, not only when none
+    // have one: without an instance_prompt the backend discovery silently skips every
+    // uncaptioned image, so a partially captioned dataset would train on a subset.
+    if (
+      selectedDataset &&
+      selectedDataset.caption_count < selectedDataset.image_count &&
+      !instancePrompt.trim()
+    ) {
       toast.error(
-        "These images have no captions - add a trigger prompt so the trainer knows " +
-          "what to learn (it becomes the caption for every image).",
+        selectedDataset.caption_count === 0
+          ? "These images have no captions - add a trigger prompt so the trainer knows " +
+              "what to learn (it becomes the caption for every image)."
+          : `Only ${selectedDataset.caption_count} of ${selectedDataset.image_count} images ` +
+              "have captions - the rest would be silently skipped. Add a trigger prompt " +
+              "(it becomes their caption) or caption every image.",
       );
       return;
     }

studio/frontend/src/features/video/api.ts

@@ -60,11 +60,16 @@ export interface VideoStatus {
 
 export interface VideoGenerateProgress {
   active: boolean;
-  // "denoise" | "export" | null.
+  // "queued" | "denoise" | "export" | "completed" | "failed" | null. The terminal
+  // phases carry the outcome of the background job POST /video/generate started.
   phase?: string | null;
   step: number;
   total: number;
   eta_seconds?: number | null;
+  // Saved gallery record when phase is "completed".
+  video?: GalleryVideo | null;
+  // Client-safe failure detail when phase is "failed".
+  error?: string | null;
 }
 
 export interface VideoLoadProgress {
@@ -149,8 +154,12 @@ export interface GalleryVideo {
   created_at: string;
 }
 
+// Acknowledgement that the generation job started; the saved record arrives via
+// getVideoGenerateProgress when its phase reaches "completed".
 export interface VideoGenerateResponse {
-  video: GalleryVideo;
+  status: "started";
+  // Always null (kept for response-shape compatibility).
+  video?: GalleryVideo | null;
 }
 
 async function parseJson<T>(response: Response): Promise<T> {
@@ -182,6 +191,9 @@ export async function loadVideoModel(body: VideoLoadRequest): Promise<VideoStatu
   );
 }
 
+/** Start a generation job. Returns as soon as the backend accepts it (the clip takes
+ *  minutes, and secure mode's tunnel caps responses near 100s, so the POST cannot span
+ *  the generation); poll getVideoGenerateProgress for completion. */
 export async function generateVideo(
   body: VideoGenerateRequest,
 ): Promise<VideoGenerateResponse> {

studio/frontend/src/features/video/video-page.tsx

@@ -194,6 +194,10 @@ function clipMeta(video: GalleryVideo): string {
 // "Encoding video…" during export) plus an ETA once known.
 function genStepLabel(p: VideoGenerateProgress): string {
   if (p.phase === "export") return "Encoding video…";
+  // Text encoding and the first-step warmup run inside the pipeline before the first
+  // scheduler tick, so step 0 means "working, not denoising yet" -- up to a minute at
+  // 720p. Label that phase honestly instead of sitting on "Denoising step 0/N".
+  if (p.step === 0) return "Preparing (text encoding + warmup)…";
   const base = p.total > 0 ? `Denoising step ${p.step}/${p.total}` : "Denoising…";
   const eta = p.eta_seconds != null ? formatEta(p.eta_seconds) : "";
   return eta ? `${base} · ~${eta}` : base;
@@ -541,6 +545,10 @@ export function VideoPage({ active = true }: { active?: boolean }) {
   // Live per-step progress (phase / step / total + ETA) polled during generation.
   const [genStep, setGenStep] = useState<VideoGenerateProgress | null>(null);
   const genPollTimer = useRef<ReturnType<typeof setInterval> | null>(null);
+  // visibilitychange handler active while a generation poll runs: background tabs clamp
+  // setInterval to >=1s (and can suspend it outright after ~5 min), so returning to the
+  // tab fires one immediate poll instead of waiting for a throttled tick.
+  const genVisibilityListener = useRef<(() => void) | null>(null);
   const [status, setStatus] = useState<VideoStatus | null>(null);
   // Controlled so the body-portaled overlays force-close when this page is mounted but
   // off-tab (a hidden/inert parent can't contain a body portal): the model selector.
@@ -911,6 +919,10 @@ export function VideoPage({ active = true }: { active?: boolean }) {
     return () => {
       if (pollTimer.current) clearTimeout(pollTimer.current);
       if (genPollTimer.current) clearInterval(genPollTimer.current);
+      if (genVisibilityListener.current) {
+        document.removeEventListener("visibilitychange", genVisibilityListener.current);
+        genVisibilityListener.current = null;
+      }
       dismissLoadToast();
     };
   }, [refreshStatus, dismissLoadToast, pollLoadProgress]);
@@ -1042,6 +1054,29 @@ export function VideoPage({ active = true }: { active?: boolean }) {
         });
         return;
       }
+      // A direct local single-file .safetensors pick must load via from_single_file:
+      // the pipeline route rejects a bare file (no model_index.json) and only after
+      // evicting the resident model. Split into (parent dir, basename) exactly like
+      // the local GGUF branch above.
+      if (meta.source === "local" && id.toLowerCase().endsWith(".safetensors")) {
+        const norm = id.replace(/\\/g, "/");
+        const slash = norm.lastIndexOf("/");
+        const filename = slash >= 0 ? norm.slice(slash + 1) : norm;
+        const dir = slash >= 0 ? norm.slice(0, slash) : ".";
+        const prevQuant = quant;
+        quantRevert.current = { prev: prevQuant };
+        setQuant(filename);
+        const dsf = defaultsFor(id);
+        setSteps(dsf.steps);
+        setGuidance(dsf.guidance);
+        void handleLoad(dir, { kind: "single_file", filename }).then((started) => {
+          if (!started) {
+            setQuant(prevQuant);
+            quantRevert.current = null;
+          }
+        });
+        return;
+      }
       // Otherwise treat it as a full diffusers repo. The backend gates loads to unsloth/*
       // repos, the family bases, or on-device paths, so only attempt those.
       if (meta.source !== "local" && !id.toLowerCase().startsWith("unsloth/")) {
@@ -1108,11 +1143,68 @@ export function VideoPage({ active = true }: { active?: boolean }) {
 
     setBusy("generating");
     setGenStep(null);
-    // Poll the backend's per-step progress so the bar tracks the live denoising steps and
-    // the encode phase.
-    genPollTimer.current = setInterval(async () => {
+    // The POST only STARTS the job and returns at once (a clip takes minutes, and
+    // secure mode's tunnel caps responses near 100s, so completion cannot ride the
+    // POST). A synchronous rejection (no model / already generating / bad input)
+    // still surfaces here; everything after acceptance arrives via the poll.
+    try {
+      await generateVideo({
+        prompt: prompt.trim(),
+        // Only send a negative prompt when guidance uses it, so the recipe doesn't record
+        // one the model ignored.
+        negative_prompt: guidance > 0 ? negativePrompt.trim() || undefined : undefined,
+        width: w,
+        height: h,
+        num_frames: numFrames,
+        fps,
+        steps,
+        guidance,
+        seed: resolvedSeed,
+      });
+    } catch (err) {
+      if (!isMounted.current) return;
+      toast.error(err instanceof Error ? err.message : "Video generation failed");
+      setBusy(null);
+      setGenStep(null);
+      return;
+    }
+    // Poll the backend's per-step progress so the bar tracks the live denoising steps
+    // and the encode phase, and drive completion off the terminal phase: "completed"
+    // carries the saved gallery record, "failed" the client-safe error. A named poll
+    // body (guarded against overlap) also serves the visibilitychange listener: a
+    // background tab's throttled interval catches up the moment the tab is visible.
+    let pollInFlight = false;
+    const stopGenPoll = () => {
+      if (genPollTimer.current) clearInterval(genPollTimer.current);
+      genPollTimer.current = null;
+      if (genVisibilityListener.current) {
+        document.removeEventListener("visibilitychange", genVisibilityListener.current);
+        genVisibilityListener.current = null;
+      }
+    };
+    const pollGenerateOnce = async () => {
+      if (pollInFlight) return;
+      pollInFlight = true;
       try {
         const p = await getVideoGenerateProgress();
+        if (p.phase === "completed" || p.phase === "failed") {
+          stopGenPoll();
+          if (!isMounted.current) return;
+          setBusy(null);
+          setGenStep(null);
+          if (p.phase === "completed" && p.video) {
+            // Prepend the new clip (newest first) and load its blob.
+            const clip = p.video;
+            setVideos((prev) => [clip, ...prev.filter((v) => v.id !== clip.id)]);
+            setSelectedId(clip.id);
+            void ensureSrc(clip);
+          } else if (p.phase === "failed") {
+            const msg = p.error || "Video generation failed";
+            // The user's own Cancel surfaces as the backend's cancelled sentinel; not an error.
+            if (!msg.toLowerCase().includes("cancelled")) toast.error(msg);
+          }
+          return;
+        }
         setGenStep((prev) => {
           if (!p.active) return null;
           if (
@@ -1126,37 +1218,17 @@ export function VideoPage({ active = true }: { active?: boolean }) {
         });
       } catch {
         // transient; keep polling
+      } finally {
+        pollInFlight = false;
       }
-    }, 300);
-    try {
-      const res = await generateVideo({
-        prompt: prompt.trim(),
-        // Only send a negative prompt when guidance uses it, so the recipe doesn't record
-        // one the model ignored.
-        negative_prompt: guidance > 0 ? negativePrompt.trim() || undefined : undefined,
-        width: w,
-        height: h,
-        num_frames: numFrames,
-        fps,
-        steps,
-        guidance,
-        seed: resolvedSeed,
-      });
-      if (!isMounted.current) return;
-      // Prepend the new clip (newest first) and load its blob.
-      setVideos((prev) => [res.video, ...prev.filter((v) => v.id !== res.video.id)]);
-      setSelectedId(res.video.id);
-      void ensureSrc(res.video);
-    } catch (err) {
-      const msg = err instanceof Error ? err.message : "Video generation failed";
-      // The user's own Cancel comes back as the backend's 409 sentinel; not an error.
-      if (!msg.toLowerCase().includes("cancelled")) toast.error(msg);
-    } finally {
-      if (genPollTimer.current) clearInterval(genPollTimer.current);
-      genPollTimer.current = null;
-      setBusy(null);
-      setGenStep(null);
-    }
+    };
+    if (genVisibilityListener.current)
+      document.removeEventListener("visibilitychange", genVisibilityListener.current);
+    genVisibilityListener.current = () => {
+      if (document.visibilityState === "visible") void pollGenerateOnce();
+    };
+    document.addEventListener("visibilitychange", genVisibilityListener.current);
+    genPollTimer.current = setInterval(() => void pollGenerateOnce(), 300);
   }, [
     prompt,
     negativePrompt,

studio/install_llama_prebuilt.py

@@ -10,6 +10,7 @@ import argparse
 import atexit
 import errno
 import fnmatch
+import glob
 import hashlib
 import json
 import os
@@ -265,6 +266,7 @@ class HostInfo:
     has_physical_nvidia: bool
     has_usable_nvidia: bool
     has_rocm: bool = False
+    has_intel_gpu: bool = False
     rocm_gfx_target: str | None = None
     # (major, minor) from platform.mac_ver(); None off macOS or if unparseable.
     # Skips a macos prebuilt whose minimum-OS exceeds this host.
@@ -1284,162 +1286,6 @@ def synthetic_checksums_for_release(
     )
 
 
-def parse_direct_linux_release_bundle(
-    repo: str, release: dict[str, Any]
-) -> PublishedReleaseBundle | None:
-    release_tag = release.get("tag_name")
-    if not isinstance(release_tag, str) or not release_tag:
-        return None
-
-    assets = release_asset_map(release)
-    artifacts: list[PublishedLlamaArtifact] = []
-    inferred_labels: list[str] = []
-
-    linux_asset_re = re.compile(
-        r"^app-(?P<label>.+)-(?P<target>linux-x64(?:-cpu)?|linux-x64-cuda\d+-(?:older|newer|portable))\.tar\.gz$"
-    )
-    for asset_name in sorted(assets):
-        match = linux_asset_re.fullmatch(asset_name)
-        if not match:
-            continue
-        inferred_labels.append(match.group("label"))
-        target = match.group("target")
-        if target in {"linux-x64", "linux-x64-cpu"}:
-            artifacts.append(
-                PublishedLlamaArtifact(
-                    asset_name = asset_name,
-                    install_kind = "linux-cpu",
-                    runtime_line = None,
-                    coverage_class = None,
-                    supported_sms = [],
-                    min_sm = None,
-                    max_sm = None,
-                    bundle_profile = None,
-                    rank = 1000,
-                )
-            )
-            continue
-
-        bundle_profile = target.removeprefix("linux-x64-")
-        profile = _resolve_linux_bundle_profile(bundle_profile)
-        if profile is None:
-            continue
-        artifacts.append(
-            PublishedLlamaArtifact(
-                asset_name = asset_name,
-                install_kind = "linux-cuda",
-                runtime_line = str(profile["runtime_line"]),
-                coverage_class = str(profile["coverage_class"]),
-                supported_sms = [str(value) for value in profile["supported_sms"]],
-                min_sm = int(profile["min_sm"]),
-                max_sm = int(profile["max_sm"]),
-                bundle_profile = bundle_profile,
-                rank = int(profile["rank"]),
-            )
-        )
-
-    if not artifacts:
-        return None
-
-    upstream_tag = (
-        release_tag
-        if is_release_tag_like(release_tag)
-        else inferred_labels[0]
-        if len(set(inferred_labels)) == 1 and inferred_labels
-        else release_tag
-    )
-    selection_log = [
-        f"published_release: repo={repo}",
-        f"published_release: tag={release_tag}",
-        f"published_release: upstream_tag={upstream_tag}",
-        "published_release: direct_asset_scan=linux",
-    ]
-    return PublishedReleaseBundle(
-        repo = repo,
-        release_tag = release_tag,
-        upstream_tag = upstream_tag,
-        assets = assets,
-        manifest_asset_name = DEFAULT_PUBLISHED_MANIFEST_ASSET,
-        artifacts = artifacts,
-        selection_log = selection_log,
-    )
-
-
-def direct_linux_release_plan(
-    release: dict[str, Any], host: HostInfo, repo: str, requested_tag: str
-) -> InstallReleasePlan | None:
-    bundle = parse_direct_linux_release_bundle(repo, release)
-    if bundle is None:
-        return None
-    if not direct_release_matches_request(
-        release_tag = bundle.release_tag,
-        llama_tag = bundle.upstream_tag,
-        requested_tag = requested_tag,
-    ):
-        return None
-
-    attempts: list[AssetChoice] = []
-    if host.has_usable_nvidia:
-        # Prefer the cudart major Studio loads at runtime (torch's bundled
-        # libcudart), not the newest on disk. Otherwise a stray cuda13
-        # runtime outranks the torch cuda12 the binary links against.
-        torch_preference = detect_torch_cuda_runtime_preference(host)
-        selection = linux_cuda_choice_from_release(
-            host,
-            bundle,
-            preferred_runtime_line = torch_preference.runtime_line,
-            selection_preamble = torch_preference.selection_log,
-        )
-        if selection is not None:
-            attempts.extend(selection.attempts)
-    elif not host.has_rocm:
-        # A ROCm-only host gets no CPU asset: leaving attempts empty lets the
-        # raise below trigger a HIP source build instead of shipping a CPU
-        # binary on a GPU host (this ggml-org path has no per-gfx ROCm asset).
-        cpu_choice = published_asset_choice_for_kind(bundle, "linux-cpu")
-        if cpu_choice is not None:
-            attempts.append(cpu_choice)
-    # NVIDIA hosts whose CUDA selection produced nothing fall through to the
-    # raise below (mirroring the ROCm policy above): the caller then walks
-    # back to an older release that still ships a usable CUDA line instead of
-    # silently installing a CPU binary on a GPU host. Today's walk-back only
-    # works because partial releases ship no CPU bundle; this keeps it working
-    # if a future partial release does.
-    if not attempts:
-        raise PrebuiltFallback("no compatible Linux prebuilt asset was found")
-    approved_checksums = synthetic_checksums_for_release(
-        repo,
-        bundle.release_tag,
-        bundle.upstream_tag,
-    )
-    resolved_upstream_tag = bundle.upstream_tag
-    if DEFAULT_PUBLISHED_SHA256_ASSET in bundle.assets and not is_release_tag_like(
-        bundle.upstream_tag
-    ):
-        approved_checksums = load_approved_release_checksums(repo, bundle.release_tag)
-        # Require exact source provenance for branch/pull/commit releases.
-        # Mirrors validated_checksums_for_bundle so incomplete metadata fails
-        # closed instead of degrading to the legacy branch-as-tag source
-        # hydration path this PR eliminates.
-        if (
-            not approved_checksums.source_commit
-            or exact_source_archive_hash(approved_checksums) is None
-            or source_clone_url_from_checksums(approved_checksums) is None
-        ):
-            raise PrebuiltFallback(
-                f"approved checksum asset {DEFAULT_PUBLISHED_SHA256_ASSET} for "
-                f"{repo}@{bundle.release_tag} did not contain exact source provenance"
-            )
-        attempts = apply_approved_hashes(attempts, approved_checksums)
-    return InstallReleasePlan(
-        requested_tag = requested_tag,
-        llama_tag = resolved_upstream_tag,
-        release_tag = bundle.release_tag,
-        attempts = attempts,
-        approved_checksums = approved_checksums,
-    )
-
-
 def direct_upstream_release_plan(
     release: dict[str, Any], host: HostInfo, repo: str, requested_tag: str
 ) -> InstallReleasePlan | None:
@@ -1482,6 +1328,24 @@ def direct_upstream_release_plan(
                         install_kind = "windows-hip",
                     )
                 )
+        # Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. Gate
+        # on no PHYSICAL NVIDIA (not just no usable one): a host that hid NVIDIA
+        # via CUDA_VISIBLE_DEVICES must not reach Vulkan, which ignores that mask
+        # and could enumerate the reserved card. Falls through to CPU below.
+        elif host.has_intel_gpu and not host.has_physical_nvidia:
+            vulkan_asset = f"llama-{release_tag}-bin-win-vulkan-x64.zip"
+            vulkan_url = assets.get(vulkan_asset)
+            if vulkan_url:
+                attempts.append(
+                    AssetChoice(
+                        repo = repo,
+                        tag = release_tag,
+                        name = vulkan_asset,
+                        url = vulkan_url,
+                        source_label = "upstream",
+                        install_kind = "windows-vulkan",
+                    )
+                )
         cpu_asset = f"llama-{release_tag}-bin-win-cpu-x64.zip"
         cpu_url = assets.get(cpu_asset)
         if cpu_url:
@@ -1545,6 +1409,23 @@ def direct_upstream_release_plan(
         # ROCm hosts are excluded: this ggml-org path ships no per-gfx ROCm
         # asset, so they fall through to the empty-attempts raise (HIP source
         # build) rather than silently getting a CPU binary on a GPU host.
+        # Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. The
+        # elif already excludes usable NVIDIA and ROCm; also require no PHYSICAL
+        # NVIDIA so a CUDA-hidden card isn't reached through Vulkan (CPU below).
+        if host.has_intel_gpu and not host.has_physical_nvidia:
+            vulkan_asset = f"llama-{release_tag}-bin-ubuntu-vulkan-x64.tar.gz"
+            vulkan_url = assets.get(vulkan_asset)
+            if vulkan_url:
+                attempts.append(
+                    AssetChoice(
+                        repo = repo,
+                        tag = release_tag,
+                        name = vulkan_asset,
+                        url = vulkan_url,
+                        source_label = "upstream",
+                        install_kind = "linux-vulkan",
+                    )
+                )
         asset_name = f"llama-{release_tag}-bin-ubuntu-x64.tar.gz"
         asset_url = assets.get(asset_name)
         if asset_url:
@@ -1564,6 +1445,23 @@ def direct_upstream_release_plan(
         # selector returned 0 attempts and the installer fell back to a
         # source build on every Linux ARM64 host (DGX Spark, Ampere
         # Altra, GitHub-hosted ubuntu-24.04-arm runners, etc.).
+        # Intel (or other non-NVIDIA/non-AMD) GPU: prefer the Vulkan prebuilt,
+        # mirroring the x86_64 branch. Upstream ships bin-ubuntu-vulkan-arm64.
+        # No physical NVIDIA: don't reach a CUDA-hidden card through Vulkan.
+        if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm:
+            vulkan_asset = f"llama-{release_tag}-bin-ubuntu-vulkan-arm64.tar.gz"
+            vulkan_url = assets.get(vulkan_asset)
+            if vulkan_url:
+                attempts.append(
+                    AssetChoice(
+                        repo = repo,
+                        tag = release_tag,
+                        name = vulkan_asset,
+                        url = vulkan_url,
+                        source_label = "upstream",
+                        install_kind = "linux-vulkan",
+                    )
+                )
         asset_name = f"llama-{release_tag}-bin-ubuntu-arm64.tar.gz"
         asset_url = assets.get(asset_name)
         if asset_url:
@@ -3075,6 +2973,40 @@ def detect_host() -> HostInfo:
         # Note: amdhip64.dll presence alone is NOT treated as GPU evidence
         # since the HIP SDK can be installed without an AMD GPU.
 
+    # Detect an Intel GPU; gates the Vulkan prebuilt. Linux reads the DRM sysfs
+    # vendor id (0x8086); Windows queries the WMI video controller list. Only
+    # probed with no usable NVIDIA and no ROCm (matching the Vulkan branches),
+    # keeping the probe (notably the Windows powershell call) off that path.
+    has_intel_gpu = False
+    if not has_usable_nvidia and not has_rocm:
+        if is_linux:
+            for _vendor_file in glob.glob("/sys/class/drm/card*/device/vendor"):
+                try:
+                    with open(_vendor_file) as _vf:
+                        if _vf.read().strip().lower() == "0x8086":
+                            has_intel_gpu = True
+                            break
+                except OSError:
+                    continue
+        elif is_windows:
+            _ps = shutil.which("powershell") or shutil.which("pwsh")
+            if _ps:
+                try:
+                    _result = run_capture(
+                        [
+                            _ps,
+                            "-NoProfile",
+                            "-Command",
+                            "Get-CimInstance Win32_VideoController | "
+                            "Select-Object -ExpandProperty Name",
+                        ],
+                        timeout = 15,
+                    )
+                    if _result.returncode == 0 and "intel" in _result.stdout.lower():
+                        has_intel_gpu = True
+                except Exception:
+                    pass
+
     return HostInfo(
         system = system,
         machine = machine,
@@ -3090,6 +3022,7 @@ def detect_host() -> HostInfo:
         has_physical_nvidia = has_physical_nvidia,
         has_usable_nvidia = has_usable_nvidia,
         has_rocm = has_rocm,
+        has_intel_gpu = has_intel_gpu,
         rocm_gfx_target = rocm_gfx_target,
         macos_version = macos_version,
     )
@@ -3126,6 +3059,7 @@ def _apply_host_overrides(
             has_physical_nvidia = False,
             has_rocm = False,
             rocm_gfx_target = None,
+            has_intel_gpu = False,
         )
     gfx = _normalize_forwarded_gfx(override_rocm_gfx)
     if gfx:
@@ -3866,6 +3800,23 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
                 "falling back to source build with HIP support"
             )
 
+        # Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. No
+        # physical NVIDIA (not just no usable one): a CUDA-hidden card must not
+        # be reached through Vulkan, which ignores CUDA_VISIBLE_DEVICES.
+        if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm:
+            vulkan_name = f"llama-{llama_tag}-bin-ubuntu-vulkan-x64.tar.gz"
+            if vulkan_name in upstream_assets:
+                log(f"Intel GPU detected -- using upstream Vulkan prebuilt {vulkan_name}")
+                return AssetChoice(
+                    repo = UPSTREAM_REPO,
+                    tag = llama_tag,
+                    name = vulkan_name,
+                    url = upstream_assets[vulkan_name],
+                    source_label = "upstream",
+                    install_kind = "linux-vulkan",
+                )
+            log("Intel GPU detected but no Vulkan prebuilt found -- falling back to CPU")
+
         upstream_name = f"llama-{llama_tag}-bin-ubuntu-x64.tar.gz"
         if upstream_name not in upstream_assets:
             raise PrebuiltFallback("upstream Linux CPU asset was not found")
@@ -3908,6 +3859,24 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
                 )
             log("AMD ROCm detected on Windows but no HIP prebuilt found -- falling back to CPU")
 
+        # Intel (or other non-NVIDIA/non-AMD) GPU on Windows: use Vulkan. No
+        # physical NVIDIA so a CUDA-hidden card isn't reached through Vulkan.
+        if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm:
+            vulkan_name = f"llama-{llama_tag}-bin-win-vulkan-x64.zip"
+            if vulkan_name in upstream_assets:
+                log(
+                    f"Intel GPU detected on Windows -- using upstream Vulkan prebuilt {vulkan_name}"
+                )
+                return AssetChoice(
+                    repo = UPSTREAM_REPO,
+                    tag = llama_tag,
+                    name = vulkan_name,
+                    url = upstream_assets[vulkan_name],
+                    source_label = "upstream",
+                    install_kind = "windows-vulkan",
+                )
+            log("Intel GPU detected on Windows but no Vulkan prebuilt found -- falling back to CPU")
+
         upstream_name = f"llama-{llama_tag}-bin-win-cpu-x64.zip"
         if upstream_name not in upstream_assets:
             raise PrebuiltFallback("upstream Windows CPU asset was not found")
@@ -4503,6 +4472,7 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]:
         "linux-arm64-cuda",
         "linux-rocm",
         "linux-arm64",
+        "linux-vulkan",
     }:
         return ["llama-server", "llama-quantize", "llama-diffusion-gemma-visual-server", "lib*.so*"]
     if choice.install_kind in {"macos-arm64", "macos-x64"}:
@@ -4516,6 +4486,7 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]:
         "windows-cpu",
         "windows-cuda",
         "windows-hip",
+        "windows-vulkan",
         "windows-rocm",
         "windows-arm64",
     }:
@@ -5731,8 +5702,10 @@ def validate_server(
             "linux-cuda",
             "linux-arm64-cuda",
             "linux-rocm",
+            "linux-vulkan",
             "windows-cuda",
             "windows-hip",
+            "windows-vulkan",
             "windows-rocm",
             "macos-arm64",
         }
@@ -6354,6 +6327,20 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]:
             ["libmtmd.so*"],
             ["libggml-hip.so*"],
         ]
+    if choice.install_kind == "linux-vulkan":
+        return [
+            ["libllama-common.so*"],
+            ["libllama.so*"],
+            ["libggml.so*"],
+            ["libggml-base.so*"],
+            # Match the sibling globs (linux-cuda/-rocm): x64 bundles ship
+            # arch-suffixed libggml-cpu-<variant>.so, arm64 may ship a bare
+            # libggml-cpu.so; the '-' form missed the latter and re-flagged
+            # the install unhealthy on every check.
+            ["libggml-cpu*.so*"],
+            ["libmtmd.so*"],
+            ["libggml-vulkan.so*"],
+        ]
     if choice.install_kind in {"windows-cpu", "windows-arm64"}:
         return [["llama.dll"]]
     if choice.install_kind == "windows-cuda":
@@ -6373,6 +6360,8 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]:
         return groups
     if choice.install_kind in {"windows-hip", "windows-rocm"}:
         return [["llama.dll"], ["*hip*.dll"]]
+    if choice.install_kind == "windows-vulkan":
+        return [["llama.dll"], ["ggml-vulkan.dll"]]
     return []
 
 
@@ -6654,6 +6643,89 @@ def validate_prebuilt_attempts(
     raise PrebuiltFallback("no prebuilt bundle passed validation")
 
 
+def force_vulkan_requested() -> bool:
+    """Whether UNSLOTH_FORCE_VULKAN opts this host into the Vulkan llama.cpp
+    prebuilt instead of its detected CUDA/ROCm backend (e.g. so an AMD user can
+    run the Vulkan build for inference). Scoped to the llama.cpp backend; the
+    torch/training stack installs separately and still sees the real GPU.
+    """
+    return os.environ.get("UNSLOTH_FORCE_VULKAN", "").strip().lower() in (
+        "1",
+        "true",
+        "yes",
+    )
+
+
+def _vulkan_only_host(host: HostInfo) -> HostInfo:
+    """Rewrite ``host`` so the asset selectors take their Vulkan branch.
+
+    That branch fires on ``has_intel_gpu and not nvidia and not rocm``, so clear
+    the CUDA/ROCm flags and raise the integrated-GPU flag. The synthetic flag
+    never leaves install planning -- it only routes the llama.cpp prebuilt
+    choice, not the torch/training stack.
+    """
+    return dataclasses_replace(
+        host,
+        has_usable_nvidia = False,
+        has_physical_nvidia = False,
+        has_rocm = False,
+        has_intel_gpu = True,
+    )
+
+
+def _route_to_vulkan_prebuilt(
+    host: HostInfo, published_repo: str, published_release_tag: str, *, force_cpu: bool
+) -> tuple[HostInfo, str, str]:
+    """Point a Vulkan-capable host at the upstream ggml-org Vulkan prebuilt.
+
+    The unsloth published repo ships only CUDA/ROCm/CPU assets, so Vulkan comes
+    from UPSTREAM_REPO. Two triggers route here, both suppressed under
+    --cpu-fallback (the explicit "give me CPU" last resort wins):
+      * UNSLOTH_FORCE_VULKAN forces Vulkan over the detected CUDA/ROCm backend;
+      * an auto-detected Intel GPU with NO physical NVIDIA/ROCm -- the purpose
+        of the has_intel_gpu probe, since the fork manifest ships no Vulkan asset.
+    Applied by BOTH the install path and the --resolve-prebuilt probe so the
+    "is a prebuilt available" answer matches what actually gets installed.
+
+    Returns the (possibly rewritten) host, repo, and release tag.
+    """
+    forced = force_vulkan_requested()
+    # Gate auto-routing on no PHYSICAL NVIDIA, not merely no usable one: a mixed
+    # NVIDIA+Intel host that hides NVIDIA with CUDA_VISIBLE_DEVICES=""/-1 keeps
+    # has_physical_nvidia=True while has_usable_nvidia goes False. Vulkan ignores
+    # CUDA_VISIBLE_DEVICES, so auto-routing such a host would let it grab the
+    # reserved NVIDIA GPU. An explicit UNSLOTH_FORCE_VULKAN still overrides.
+    auto_intel = host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm
+    if force_cpu or not (forced or auto_intel):
+        return host, published_repo, published_release_tag
+    if host.is_macos:
+        if forced:
+            log(
+                "UNSLOTH_FORCE_VULKAN is set but ignored on macOS "
+                "(Metal is used; there is no Vulkan prebuilt)"
+            )
+        return host, published_repo, published_release_tag
+    if forced:
+        log(
+            "UNSLOTH_FORCE_VULKAN is set; installing the upstream Vulkan "
+            "llama.cpp prebuilt instead of the detected GPU backend"
+        )
+        # Forcing may override a detected NVIDIA/ROCm host, so normalize it to
+        # Vulkan-only; an auto-detected Intel host already is.
+        host = _vulkan_only_host(host)
+    else:
+        log("Intel GPU detected; installing the upstream Vulkan llama.cpp prebuilt")
+    # Swapping the fork for upstream invalidates a fork release pin: the two use
+    # different tag namespaces (fork b9596-mix-<sha> vs upstream b9596), so a
+    # pinned fork tag would make the upstream resolver query a nonexistent
+    # release and fall back to source. Drop it and let the upstream resolver
+    # pick by the requested llama tag. A pin already on an explicit upstream repo
+    # (repo unchanged here) is preserved.
+    if published_repo != UPSTREAM_REPO:
+        published_release_tag = ""
+    return host, UPSTREAM_REPO, published_release_tag
+
+
 def diffusion_visual_server_backfill_needed(
     install_dir: Path, host: HostInfo, choice: AssetChoice
 ) -> bool:
@@ -6696,6 +6768,9 @@ def install_prebuilt(
         override_rocm_gfx = override_rocm_gfx,
         force_cpu = force_cpu,
     )
+    host, published_repo, published_release_tag = _route_to_vulkan_prebuilt(
+        host, published_repo, published_release_tag, force_cpu = force_cpu
+    )
     choice: AssetChoice | None = None
     try:
         with install_lock(install_lock_path(install_dir)):
@@ -6708,7 +6783,9 @@ def install_prebuilt(
                     f"no existing llama.cpp install detected at {install_dir}; performing fresh prebuilt install"
                 )
             # Single resolver: every fork host selects from the release manifest;
-            # an explicit ggml-org override selects by asset filename instead.
+            # an explicit ggml-org override selects by asset filename instead. A
+            # forced-Vulkan host already has published_repo pointed at
+            # UPSTREAM_REPO above, so the resolver takes the Vulkan asset branch.
             requested_tag, release_plans = resolve_simple_install_release_plans(
                 llama_tag,
                 host,
@@ -6994,10 +7071,14 @@ def main() -> int:
             override_rocm_gfx = args.rocm_gfx,
             force_cpu = args.cpu_fallback,
         )
-        repo = args.published_repo
+        # Same Vulkan routing the install path applies, so the probe's answer
+        # matches what would install (an Intel/forced-Vulkan host -> upstream).
+        host, repo, release_tag = _route_to_vulkan_prebuilt(
+            host, args.published_repo, args.published_release_tag or "", force_cpu = args.cpu_fallback
+        )
         try:
             _requested, plans = resolve_simple_install_release_plans(
-                args.resolve_prebuilt, host, repo, args.published_release_tag or ""
+                args.resolve_prebuilt, host, repo, release_tag
             )
             choice = plans[0].attempts[0] if plans and plans[0].attempts else None
             if choice is None:

studio/setup.ps1

@@ -2621,7 +2621,18 @@ function Fast-Install {
     param([Parameter(ValueFromRemainingArguments=$true)]$Args_)
     if ($UseUv) {
         $VenvPy = (Get-Command python).Source
-        $result = & uv pip install --python $VenvPy @Args_ 2>&1
+        # An explicit --index-url must win. Inherited uv index env vars otherwise
+        # override it and pull CPU torch over the CUDA/ROCm build (#6898), so drop
+        # them only for index-pinned installs; mirrors still apply elsewhere.
+        $saved = @{}
+        if (@($Args_) -contains '--index-url') {
+            foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
+                $saved[$n] = [Environment]::GetEnvironmentVariable($n)
+                Remove-Item "Env:$n" -ErrorAction SilentlyContinue
+            }
+        }
+        try { $result = & uv pip install --python $VenvPy @Args_ 2>&1 }
+        finally { foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } } }
         if ($LASTEXITCODE -eq 0) { return }
     }
     & python -m pip install @Args_ 2>&1

tests/python/test_fast_sentence_transformer_embedding_parity.py

@@ -0,0 +1,122 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team.
+"""Regression guard for issue #6881: FastSentenceTransformer must preprocess text
+like a stock SentenceTransformer for decoder embedding models. ST 5.x infers a
+"message" modality for chat-template models (e.g. Qwen/Qwen3-Embedding), so building
+via `Transformer(model_name, ...)` chat-wraps inputs and degrades embeddings;
+`_create_transformer_module` uses `Transformer.load(...)` instead.
+
+Layers: test_transformer_load_signature_supports_unsloth_kwargs (fast, runs when ST
+is importable) and test_fast_sentence_transformer_matches_stock_st (end-to-end parity,
+opt-in via UNSLOTH_EMBEDDING_PARITY_MODEL so default CI is unaffected).
+"""
+
+from __future__ import annotations
+
+import inspect
+import os
+
+import pytest
+
+
+def test_transformer_load_signature_supports_unsloth_kwargs():
+    """Forwards-compat tripwire: a Hub-capable Transformer.load must accept the kwargs
+    the #6881 fix passes. Legacy ST 3.x/4.x expose load(input_path); the code falls back
+    to Transformer(...) there, so mirror that gate and skip."""
+    models = pytest.importorskip("sentence_transformers.models")
+    load = getattr(models.Transformer, "load", None)
+    assert callable(load), (
+        "sentence_transformers Transformer.load is missing; the #6881 fix in "
+        "unsloth.models.sentence_transformer._create_transformer_module depends on it."
+    )
+    params = inspect.signature(load).parameters
+    accepts_var_kw = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
+    # Mirror _create_transformer_module's hub_capable gate.
+    hub_capable = accepts_var_kw or any(k in params for k in ("token", "cache_folder", "revision"))
+    if not hub_capable:
+        pytest.skip(
+            "legacy Transformer.load(input_path); production path falls back to Transformer(...)"
+        )
+    unsupported = [
+        k
+        for k in ("token", "cache_folder", "revision", "trust_remote_code")
+        if not (accepts_var_kw or k in params)
+    ]
+    assert not unsupported, (
+        f"installed sentence_transformers Transformer.load no longer accepts {unsupported} "
+        f"and has no **kwargs; update _create_transformer_module (#6881) before it silently "
+        f"falls back to Transformer(...)."
+    )
+
+
+def _probe_texts():
+    return [
+        "roasted chickpeas in 20 kg bags",
+        "The capital of France is Paris.",
+        "A fast brown fox jumps over the lazy dog.",
+        "recette de tarte aux pommes traditionnelle",
+    ]
+
+
+def test_fast_sentence_transformer_matches_stock_st():
+    """End-to-end: FastSentenceTransformer embeddings and tokenization must match a
+    stock SentenceTransformer load of the same checkpoint. Opt-in (needs a model) and
+    GPU-only (FastSentenceTransformer requires CUDA), so it skips on CPU-only runners."""
+    model_id = os.environ.get("UNSLOTH_EMBEDDING_PARITY_MODEL")
+    if not model_id:
+        pytest.skip(
+            "set UNSLOTH_EMBEDDING_PARITY_MODEL to a chat-template embedding model "
+            "(HF id or local path) to run the #6881 parity test"
+        )
+
+    torch = pytest.importorskip("torch")
+    if not torch.cuda.is_available():
+        pytest.skip("FastSentenceTransformer requires CUDA; skipping on CPU-only runner")
+    np = pytest.importorskip("numpy")
+    pytest.importorskip("sentence_transformers")
+    from sentence_transformers import SentenceTransformer
+
+    device = "cuda"
+    # Prefer bf16 when the GPU supports it: fp16 overflows to NaN on bf16-native
+    # embedders such as EmbeddingGemma (Gemma3), which would mask real parity.
+    dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+    texts = _probe_texts()
+    max_seq_length = 256
+
+    # Control FIRST, before importing unsloth, so its global import patches never
+    # touch the stock reference (mirrors the issue's "restart runtime" repro).
+    ctrl = SentenceTransformer(model_id, device = device, model_kwargs = {"torch_dtype": dtype})
+    ctrl.max_seq_length = max_seq_length
+    ctrl_ids = ctrl.tokenize([texts[0]])["input_ids"][0].tolist()
+    ctrl_emb = np.asarray(
+        ctrl.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32
+    )
+
+    import unsloth  # noqa: F401
+    from unsloth import FastSentenceTransformer
+
+    fast = FastSentenceTransformer.from_pretrained(
+        model_id,
+        max_seq_length = max_seq_length,
+        dtype = dtype,
+        load_in_4bit = False,
+        load_in_16bit = True,
+    )
+    fast_ids = fast.tokenize([texts[0]])["input_ids"][0].tolist()
+    fast_emb = np.asarray(
+        fast.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32
+    )
+
+    # Identical tokenization = no chat-template wrapping slipped in (the #6881 defect).
+    assert fast_ids == ctrl_ids, (
+        f"tokenization diverged (chat-template wrapping regressed?):\n"
+        f"  stock: {ctrl_ids}\n  fast:  {fast_ids}"
+    )
+
+    cos = (ctrl_emb * fast_emb).sum(1) / (
+        np.linalg.norm(ctrl_emb, axis = 1) * np.linalg.norm(fast_emb, axis = 1)
+    )
+    assert float(cos.min()) > 0.99, (
+        f"embedding parity regressed: min cosine {float(cos.min()):.5f} <= 0.99 "
+        f"(per-text {[round(float(c), 5) for c in cos]})"
+    )

tests/python/test_tokenizers_and_torch_constraint.py

@@ -14,6 +14,7 @@ _TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent  # tests/
 _REPO_ROOT = _TESTS_DIR.parent  # unsloth/
 _INSTALL_SH = _REPO_ROOT / "install.sh"
 _INSTALL_PS1 = _REPO_ROOT / "install.ps1"
+_SETUP_PS1 = _REPO_ROOT / "studio" / "setup.ps1"
 _NO_TORCH_RT = _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt"
 
 
@@ -109,6 +110,56 @@ class TestStructuralInstallPs1Unchanged:
         assert '"torch>=2.4,<2.11.0"' in self._ps1
 
 
+class TestInstallPs1UvDefaultIndex:
+    """Installer-managed torch indexes must override inherited uv defaults."""
+
+    _ps1 = _read(_INSTALL_PS1)
+
+    def test_torch_installs_use_default_index(self):
+        assert "--default-index $TorchIndexUrl" in self._ps1
+        assert "--default-index $ROCmIndexUrl" in self._ps1
+
+    def test_torch_installs_do_not_use_deprecated_index_url(self):
+        assert "--index-url $TorchIndexUrl" not in self._ps1
+        assert "--index-url $ROCmIndexUrl" not in self._ps1
+
+    def test_torch_installs_neutralize_all_uv_index_env_vars(self):
+        # Extra-index vars outrank --default-index, so pinned installs must clear them.
+        for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
+            assert var in self._ps1
+        assert 'Remove-Item "Env:$n"' in self._ps1
+
+
+class TestSetupPs1FastInstallIndex:
+    """setup.ps1 Fast-Install must neutralize inherited uv indexes when pinning."""
+
+    _ps1 = _read(_SETUP_PS1)
+
+    def test_fast_install_clears_all_uv_index_env_vars(self):
+        for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
+            assert var in self._ps1
+        # Must truly remove the vars (child sees no value), not set them empty.
+        assert 'Remove-Item "Env:$n"' in self._ps1
+
+
+class TestInstallShUvDefaultIndex:
+    """Linux/Mac installer torch indexes must override inherited uv defaults."""
+
+    _sh = _read(_INSTALL_SH)
+
+    def test_torch_installs_use_default_index(self):
+        assert '--default-index "$TORCH_INDEX_URL"' in self._sh
+
+    def test_torch_installs_do_not_use_deprecated_index_url(self):
+        assert '--index-url "$TORCH_INDEX_URL"' not in self._sh
+
+    def test_torch_installs_neutralize_all_uv_index_env_vars(self):
+        # --default-index installs run with all uv index env vars unset via `env -u`.
+        assert (
+            "env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in self._sh
+        )
+
+
 # Group 2 -- Shell snippet tests (bash subprocess, mocked python)
 class TestTorchConstraintShell:
     """Test the TORCH_CONSTRAINT block via bash with mocked python minor versions."""

tests/sh/test_mac_intel_compat.sh

@@ -312,7 +312,7 @@ if [ "$SKIP_TORCH" = true ]; then
 else
     echo "==> Installing PyTorch ($TORCH_INDEX_URL)..."
     uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \
-        --index-url "$TORCH_INDEX_URL"
+        --default-index "$TORCH_INDEX_URL"
 fi
 TORCH_EOF
 

tests/studio/install/test_selection_logic.py

@@ -2684,68 +2684,6 @@ class TestBlackwellCuda124Exclusion:
         assert kept == [cpu]
 
 
-# N.1c3. direct_linux_release_plan -- no silent CPU on NVIDIA hosts
-
-
-class TestDirectLinuxNvidiaCpuGate:
-    """A linux-cpu-only release on an NVIDIA host must raise (caller walks back to a usable CUDA line), not silently CPU-install. CPU-only hosts keep the CPU bundle."""
-
-    def _bundle_cpu_only(self):
-        return make_release(
-            [
-                make_artifact(
-                    "llama-b8508-bin-ubuntu-x64.tar.gz",
-                    install_kind = "linux-cpu",
-                    runtime_line = None,
-                    coverage_class = None,
-                    supported_sms = [],
-                    min_sm = None,
-                    max_sm = None,
-                    bundle_profile = None,
-                ),
-            ]
-        )
-
-    def _patch(self, monkeypatch):
-        monkeypatch.setattr(
-            INSTALL_LLAMA_PREBUILT,
-            "parse_direct_linux_release_bundle",
-            lambda repo, release: self._bundle_cpu_only(),
-        )
-        monkeypatch.setattr(
-            INSTALL_LLAMA_PREBUILT,
-            "detect_torch_cuda_runtime_preference",
-            lambda host: CudaRuntimePreference(runtime_line = None, selection_log = []),
-        )
-        monkeypatch.setattr(
-            INSTALL_LLAMA_PREBUILT,
-            "detected_linux_runtime_lines",
-            lambda: (["cuda13"], {"cuda13": ["/usr/local/cuda/lib64"]}),
-        )
-
-    def test_nvidia_host_without_cuda_line_raises_for_walkback(self, monkeypatch):
-        self._patch(monkeypatch)
-        host = make_host(driver_cuda_version = (13, 1), compute_caps = ["100"])
-        with pytest.raises(PrebuiltFallback, match = "no compatible Linux prebuilt"):
-            INSTALL_LLAMA_PREBUILT.direct_linux_release_plan(
-                {"tag_name": "b8508"}, host, "unslothai/llama.cpp", "latest"
-            )
-
-    def test_cpu_host_still_gets_cpu_bundle(self, monkeypatch):
-        self._patch(monkeypatch)
-        host = make_host(
-            nvidia_smi = None,
-            driver_cuda_version = None,
-            compute_caps = [],
-            has_physical_nvidia = False,
-            has_usable_nvidia = False,
-        )
-        plan = INSTALL_LLAMA_PREBUILT.direct_linux_release_plan(
-            {"tag_name": "b8508"}, host, "unslothai/llama.cpp", "latest"
-        )
-        assert [a.install_kind for a in plan.attempts] == ["linux-cpu"]
-
-
 class TestLinuxPublishedAttemptsNvidiaCpuGate:
     """Live fork-manifest path: an NVIDIA host whose CUDA selection finds nothing gets an empty attempt list (source-builds with CUDA), not the manifest CPU bundle. CPU-only hosts still get the CPU bundle."""
 

tests/studio/playwright_chat_ui.py

@@ -1296,16 +1296,92 @@ with sync_playwright() as p:
     # still abort or interrupt this navigation, so the field wait below is the
     # final confirmation that we reached /login.
     _tolerated_nav = ("ERR_ABORTED", "interrupted by another navigation")
-    try:
-        page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
-    except Exception as exc:
-        if not any(t in str(exc) for t in _tolerated_nav):
-            raise
-        info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login")
-    pw_field = page.locator("#password")
-    pw_field.wait_for(state = "visible", timeout = 60_000)
-    pw_field.fill(NEW2)
-    page.locator('button[type="submit"]').click()
+    # A slow CI runner can make this re-login navigation time out even with the
+    # server healthy, so retry the whole goto/wait/fill/submit sequence (mirrors
+    # the change-password retry above). wait_for_health is a diagnostic pre-gate.
+    wait_for_health(BASE, timeout = 30.0, info = info)
+    relogin_err: Exception | None = None
+    for _relogin_attempt in range(3):
+        try:
+            try:
+                page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
+            except Exception as exc:
+                if not any(t in str(exc) for t in _tolerated_nav):
+                    raise
+                info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login")
+            pw_field = page.locator("#password")
+            pw_field.wait_for(state = "visible", timeout = 60_000)
+            pw_field.fill(NEW2)
+            # Wait on the login POST so a transient 4xx/5xx is caught and retried
+            # here, not swallowed until the out-of-loop composer wait.
+            status, _ = click_and_wait_for_response(
+                page,
+                url_substr = "/api/auth/login",
+                method = "POST",
+                do_click = lambda: page.locator('button[type="submit"]').click(),
+                timeout_ms = 30_000,
+                info = lambda m: print(f"[ui]   {m}", flush = True),
+            )
+            if status is not None and status >= 400:
+                raise AssertionError(
+                    f"login POST returned {status}; see console_errors={console_errors[:1]!r}"
+                )
+            relogin_err = None
+            break
+        except Exception as e:
+            relogin_err = e
+            try:
+                cur_url = page.url
+            except Exception:
+                cur_url = "<page closed>"
+            print(
+                f"[ui]   re-login attempt {_relogin_attempt + 1} failed: "
+                f"{type(e).__name__}: {str(e)[:200]}; page.url={cur_url}; "
+                f"page_errors={len(page_errors)} console_errors={len(console_errors)}",
+                flush = True,
+            )
+            if console_errors:
+                print(
+                    f"[ui]   first console.error: {console_errors[0][:200]!r}",
+                    flush = True,
+                )
+            if page_errors:
+                print(f"[ui]   first pageerror:    {page_errors[0][:200]!r}", flush = True)
+            try:
+                shoot(f"18-relogin-attempt-{_relogin_attempt + 1}-fail")
+            except Exception:
+                pass
+            if _relogin_attempt < 2:
+                # ERR_NO_BUFFER_SPACE needs the OS to recover socket
+                # buffers; back off 5s then 15s before retrying.
+                if "ERR_NO_BUFFER_SPACE" in str(e):
+                    backoff_s = 5 if _relogin_attempt == 0 else 15
+                    print(
+                        f"[ui]   ENOBUFS detected; sleeping {backoff_s}s "
+                        f"before retry to let OS recover socket buffers...",
+                        flush = True,
+                    )
+                    time.sleep(backoff_s)
+                # Replace the page if it died; otherwise next iteration's
+                # page.goto() handles the reload.
+                old_page = page
+                page = recover_or_replace_page(
+                    page,
+                    ctx,
+                    default_timeout_ms = 60_000,
+                    info = lambda m: print(f"[ui]   recovery: {m}", flush = True),
+                )
+                # A freshly created replacement page loses the pageerror/console
+                # listeners; re-attach so error tracking survives recovery.
+                if page is not old_page:
+                    page.on("pageerror", lambda e: page_errors.append(str(e)))
+                    page.on("console", _on_console)
+    if relogin_err is not None:
+        raise relogin_err
+    # Composer mount confirms the rotated session is authenticated. Kept OUTSIDE the
+    # retry: the loop breaks right after submit, so we never re-goto /login once login
+    # has set tokens -- that would hit the guest guard, redirect to /chat, and make a
+    # merely-slow composer look like a broken login.
     composer = page.locator('textarea[aria-label="Message input"]')
     composer.wait_for(state = "visible", timeout = 60_000)
     shoot("18-relogin-with-NEW2")

tests/test_fast_gemv_dispatch.py

@@ -0,0 +1,63 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""`get_lora_parameters` must not treat a `weight_scale` as a quant state for a weight that is
+already dequantized to bf16 (e.g. a compressed-tensors layer at forward time). Otherwise the
+bnb fast_gemv / fast_dequantize path reads a missing `absmax` and crashes.
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+# unsloth.kernels.utils imports bitsandbytes unconditionally, so skip the whole module up
+# front on runners without it (e.g. CPU-only) before importing unsloth, otherwise collection
+# errors instead of producing a skip. Any other import error still surfaces as a failure.
+pytest.importorskip("bitsandbytes")
+
+import unsloth  # noqa: F401  (sets UNSLOTH_IS_PRESENT before transformers)
+from unsloth.kernels.utils import get_lora_parameters_bias, _FP8_WEIGHT_DTYPES
+
+_FP8 = _FP8_WEIGHT_DTYPES[0] if _FP8_WEIGHT_DTYPES else None
+
+
+def _proj(weight, weight_scale = None):
+    proj = SimpleNamespace(weight = weight, bias = None, merged = False)
+    if weight_scale is not None:
+        proj.weight_scale = weight_scale
+    return proj
+
+
+def test_bf16_weight_scale_not_used_as_quant_state():
+    """A bf16 weight carrying a weight_scale (compressed-tensors) -> quant state must be None."""
+    proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16), torch.rand(2, 2))
+    W, W_quant = get_lora_parameters_bias(proj)[:2]
+    assert W_quant is None
+
+
+def test_fp8_weight_keeps_scale():
+    """An actual fp8 weight still resolves its weight_scale as the quant state."""
+    if _FP8 is None:
+        pytest.skip("no float8 dtype in this torch build")
+    scale = torch.rand(2, 2)
+    proj = _proj(torch.randn(4, 4).to(_FP8), scale)
+    W, W_quant = get_lora_parameters_bias(proj)[:2]
+    assert W_quant is scale
+
+
+def test_plain_bf16_has_no_quant_state():
+    proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16))
+    W, W_quant = get_lora_parameters_bias(proj)[:2]
+    assert W_quant is None

tests/test_fp8_restore_dropped_scale.py

@@ -0,0 +1,358 @@
+# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Restoring dropped block-fp8 `weight_scale_inv` tensors on load (#6200).
+
+Some block-scale fp8 checkpoints leave a Linear (e.g. `mlp.gate_proj`) unconverted, so its raw
+quantized values land in a plain bf16 weight and its `weight_scale_inv` is dropped, producing a
+garbage un-scaled weight. `_restore_dropped_fp8_scales` dequantizes such orphaned weights in place
+using the scale from the checkpoint. Runs offline on CPU with synthetic checkpoints.
+"""
+
+import json
+import os
+import tempfile
+from types import SimpleNamespace
+
+import torch
+from torch import nn
+from safetensors.torch import save_file
+
+# Import unsloth first to set UNSLOTH_IS_PRESENT env var.
+import unsloth
+from unsloth.models.loader_utils import _restore_dropped_fp8_scales, _FP8_DTYPES
+
+
+_SHARD = "model-00001-of-00001.safetensors"
+_FP8 = _FP8_DTYPES[0] if _FP8_DTYPES else None
+
+
+def _write_checkpoint(
+    path,
+    tensors,
+    filename = _SHARD,
+    include_index = True,
+):
+    save_file(tensors, os.path.join(path, filename))
+    if include_index:
+        weight_map = {name: filename for name in tensors}
+        with open(os.path.join(path, "model.safetensors.index.json"), "w") as f:
+            json.dump({"weight_map": weight_map}, f)
+
+
+def _fp8_config(block = (2, 2)):
+    return SimpleNamespace(
+        quantization_config = {
+            "quant_method": "fp8",
+            "weight_block_size": list(block),
+        }
+    )
+
+
+def _fp8_anchor():
+    """A module carrying a real fp8 weight, so the model looks like a genuine fp8 load."""
+    m = nn.Linear(2, 2, bias = False)
+    m.weight = nn.Parameter(torch.randn(2, 2).to(_FP8), requires_grad = False)
+    return m
+
+
+def _bf16_linear(out_f, in_f, raw):
+    m = nn.Linear(in_f, out_f, bias = False).to(torch.bfloat16)
+    with torch.no_grad():
+        m.weight.copy_(raw)
+    return m
+
+
+def _expand(scale, block, shape):
+    bs0, bs1 = block
+    expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(bs1, dim = 1)
+    return expanded[: shape[0], : shape[1]]
+
+
+def test_restore_dequantizes_orphaned_scale():
+    """A plain bf16 weight whose scale was dropped is dequantized in place."""
+    if _FP8 is None:
+        return
+    torch.manual_seed(0)
+    raw = torch.randn(4, 4, dtype = torch.bfloat16)
+    scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
+
+    model = nn.Module()
+    model.config = _fp8_config((2, 2))
+    model.anchor = _fp8_anchor()
+    model.layer = _bf16_linear(4, 4, raw)
+
+    with tempfile.TemporaryDirectory() as d:
+        _write_checkpoint(
+            d,
+            {
+                "layer.weight": raw.to(torch.float32),
+                "layer.weight_scale_inv": scale,
+            },
+        )
+        restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
+
+    assert restored == 1
+    expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (4, 4))).to(torch.bfloat16)
+    assert torch.equal(model.layer.weight.data, expected)
+
+
+def test_skips_already_fp8_weight():
+    """A correctly converted fp8 weight is skipped, never double-scaled."""
+    if _FP8 is None:
+        return
+    weight = torch.randn(4, 4).to(_FP8)
+    before = weight.clone()
+
+    model = nn.Module()
+    model.config = _fp8_config((2, 2))
+    model.layer = nn.Linear(4, 4, bias = False)
+    model.layer.weight = nn.Parameter(weight, requires_grad = False)
+
+    with tempfile.TemporaryDirectory() as d:
+        _write_checkpoint(d, {"layer.weight_scale_inv": torch.rand(2, 2)})
+        restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
+
+    assert restored == 0 and skipped == 1
+    assert torch.equal(model.layer.weight.data.float(), before.float())
+
+
+def test_skips_offloaded_meta_weight():
+    """A disk-offloaded layer (weight on the meta device) is skipped without error or restore."""
+    if _FP8 is None:
+        return
+    raw = torch.randn(4, 4, dtype = torch.bfloat16)
+    scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
+
+    model = nn.Module()
+    model.config = _fp8_config((2, 2))
+    model.anchor = _fp8_anchor()
+    model.layer = nn.Linear(4, 4, bias = False)
+    # Simulate an offloaded weight living on the meta device.
+    model.layer.weight = nn.Parameter(
+        torch.empty(4, 4, dtype = torch.bfloat16, device = "meta"), requires_grad = False
+    )
+
+    with tempfile.TemporaryDirectory() as d:
+        _write_checkpoint(
+            d,
+            {
+                "layer.weight": raw.to(torch.float32),
+                "layer.weight_scale_inv": scale,
+            },
+        )
+        restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
+
+    assert restored == 0
+    assert model.layer.weight.device.type == "meta"
+
+
+def test_noop_when_fully_dequantized():
+    """If the model has no fp8 weights at all (e.g. load_in_16bit dequantize), do not rescale."""
+    raw = torch.randn(4, 4, dtype = torch.bfloat16)
+    scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
+
+    model = nn.Module()
+    model.config = _fp8_config((2, 2))
+    model.layer = _bf16_linear(4, 4, raw)  # no fp8 anchor -> looks dequantized
+
+    with tempfile.TemporaryDirectory() as d:
+        _write_checkpoint(d, {"layer.weight_scale_inv": scale})
+        restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
+
+    assert (restored, skipped) == (0, 0)
+    assert torch.equal(model.layer.weight.data, raw)  # untouched
+
+
+def test_non_block_divisible_shape():
+    """Block scale is expanded then sliced to a non-divisible weight shape."""
+    if _FP8 is None:
+        return
+    raw = torch.randn(3, 4, dtype = torch.bfloat16)
+    scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
+
+    model = nn.Module()
+    model.config = _fp8_config((2, 2))
+    model.anchor = _fp8_anchor()
+    model.layer = _bf16_linear(3, 4, raw)  # weight shape [3, 4]
+
+    with tempfile.TemporaryDirectory() as d:
+        _write_checkpoint(d, {"layer.weight_scale_inv": scale})
+        restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
+
+    assert restored == 1
+    expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (3, 4))).to(torch.bfloat16)
+    assert torch.equal(model.layer.weight.data, expected)
+
+
+def test_transposed_scale_layout():
+    """A scale stored in the transposed block grid is transposed before use."""
+    if _FP8 is None:
+        return
+    raw = torch.randn(4, 2, dtype = torch.bfloat16)  # weight [4, 2] -> grid (2, 1)
+    scale_correct = torch.rand(2, 1, dtype = torch.float32) + 0.1
+    scale_stored = scale_correct.t().contiguous()  # stored transposed as (1, 2)
+
+    model = nn.Module()
+    model.config = _fp8_config((2, 2))
+    model.anchor = _fp8_anchor()
+    model.layer = _bf16_linear(4, 2, raw)
+
+    with tempfile.TemporaryDirectory() as d:
+        _write_checkpoint(d, {"layer.weight_scale_inv": scale_stored})
+        restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
+
+    assert restored == 1
+    expected = (raw.to(torch.float32) * _expand(scale_correct, (2, 2), (4, 2))).to(torch.bfloat16)
+    assert torch.equal(model.layer.weight.data, expected)
+
+
+def test_single_file_checkpoint_without_index():
+    """Unsharded model.safetensors (no index) is still scanned for dropped scales."""
+    if _FP8 is None:
+        return
+    raw = torch.randn(4, 4, dtype = torch.bfloat16)
+    scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
+
+    model = nn.Module()
+    model.config = _fp8_config((2, 2))
+    model.anchor = _fp8_anchor()
+    model.layer = _bf16_linear(4, 4, raw)
+
+    with tempfile.TemporaryDirectory() as d:
+        _write_checkpoint(
+            d, {"layer.weight_scale_inv": scale}, filename = "model.safetensors", include_index = False
+        )
+        restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
+
+    assert restored == 1
+    expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (4, 4))).to(torch.bfloat16)
+    assert torch.equal(model.layer.weight.data, expected)
+
+
+def test_scalar_block_size_config():
+    """A scalar weight_block_size (not a list) is handled without error."""
+    if _FP8 is None:
+        return
+    raw = torch.randn(4, 4, dtype = torch.bfloat16)
+    scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
+
+    model = nn.Module()
+    model.config = SimpleNamespace(
+        quantization_config = {"quant_method": "fp8", "weight_block_size": 2}
+    )
+    model.anchor = _fp8_anchor()
+    model.layer = _bf16_linear(4, 4, raw)
+
+    with tempfile.TemporaryDirectory() as d:
+        _write_checkpoint(d, {"layer.weight_scale_inv": scale})
+        restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
+
+    assert restored == 1
+
+
+def test_text_only_prefix_mapping():
+    """Checkpoint keys with a language_model prefix match the stripped text-only module names."""
+    if _FP8 is None:
+        return
+    raw = torch.randn(2, 2, dtype = torch.bfloat16)
+    scale = torch.rand(1, 1, dtype = torch.float32) + 0.1
+
+    model = nn.Module()
+    model.config = _fp8_config((2, 2))
+    model.anchor = _fp8_anchor()
+    model.model = nn.Module()
+    model.model.gate_proj = _bf16_linear(2, 2, raw)  # module lacks the language_model prefix
+
+    with tempfile.TemporaryDirectory() as d:
+        # checkpoint key carries the language_model wrapper the text-only load stripped
+        _write_checkpoint(d, {"model.language_model.gate_proj.weight_scale_inv": scale})
+        restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
+
+    assert restored == 1
+    expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (2, 2))).to(torch.bfloat16)
+    assert torch.equal(model.model.gate_proj.weight.data, expected)
+
+
+def test_skips_variant_load():
+    """A variant load (variant="fp8") is skipped to avoid applying default-checkpoint scales."""
+    if _FP8 is None:
+        return
+    raw = torch.randn(4, 4, dtype = torch.bfloat16)
+    scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
+    model = nn.Module()
+    model.config = _fp8_config((2, 2))
+    model.anchor = _fp8_anchor()
+    model.layer = _bf16_linear(4, 4, raw)
+    with tempfile.TemporaryDirectory() as d:
+        _write_checkpoint(d, {"layer.weight_scale_inv": scale})
+        result = _restore_dropped_fp8_scales(model, d, local_files_only = True, variant = "fp8")
+    assert result == (0, 0)
+    assert torch.equal(model.layer.weight.data, raw)  # untouched
+
+
+def test_vlm_language_model_model_alias():
+    """A checkpoint key language_model.model.* matches a model.language_model.* module."""
+    if _FP8 is None:
+        return
+    raw = torch.randn(2, 2, dtype = torch.bfloat16)
+    scale = torch.rand(1, 1, dtype = torch.float32) + 0.1
+    model = nn.Module()
+    model.config = _fp8_config((2, 2))
+    model.anchor = _fp8_anchor()
+    model.model = nn.Module()
+    model.model.language_model = nn.Module()
+    model.model.language_model.gate_proj = _bf16_linear(
+        2, 2, raw
+    )  # -> model.language_model.gate_proj
+    with tempfile.TemporaryDirectory() as d:
+        _write_checkpoint(d, {"language_model.model.gate_proj.weight_scale_inv": scale})
+        restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
+    assert restored == 1
+    expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (2, 2))).to(torch.bfloat16)
+    assert torch.equal(model.model.language_model.gate_proj.weight.data, expected)
+
+
+def test_noop_without_scale_keys():
+    if _FP8 is None:
+        return
+    model = nn.Module()
+    model.config = _fp8_config((2, 2))
+    model.anchor = _fp8_anchor()
+    model.layer = _bf16_linear(4, 4, torch.randn(4, 4, dtype = torch.bfloat16))
+    with tempfile.TemporaryDirectory() as d:
+        _write_checkpoint(d, {"layer.weight": torch.randn(4, 4)})
+        assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)
+
+
+def test_noop_without_index_or_single_file():
+    if _FP8 is None:
+        return
+    model = nn.Module()
+    model.config = _fp8_config((2, 2))
+    model.anchor = _fp8_anchor()
+    model.layer = _bf16_linear(4, 4, torch.randn(4, 4, dtype = torch.bfloat16))
+    with tempfile.TemporaryDirectory() as d:
+        assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)
+
+
+def test_noop_when_not_block_fp8():
+    """A non-fp8 (or non-block) quantization config is ignored."""
+    scale = torch.rand(2, 2)
+    model = nn.Module()
+    model.config = SimpleNamespace(quantization_config = {"quant_method": "compressed-tensors"})
+    model.layer = nn.Linear(4, 4, bias = False)
+    with tempfile.TemporaryDirectory() as d:
+        _write_checkpoint(d, {"layer.weight_scale_inv": scale})
+        assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)

tests/utils/test_rope_scaling_drift.py

@@ -257,6 +257,39 @@ def test_recompute_helper_scales_on_cpu():
     ), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled."
 
 
+def test_extended_rope_scaling_keeps_llama3_and_carries_theta():
+    # Long-context extension keeps native llama3, but falls back to linear for every other
+    # type (the patched attention constructor only rebuilds linear/llama3/longrope), and the
+    # linear dict carries rope_theta so transformers v5 does not fall back to base 10000.
+    from types import SimpleNamespace
+
+    from unsloth.models.llama import _extended_rope_scaling
+
+    # llama3 model: keep native scaling, do not synthesize linear.
+    scaling, native = _extended_rope_scaling(_make_config(LLAMA3_ROPE_SCALING), 2.0)
+    assert (
+        scaling is None and native == "llama3"
+    ), "must keep native llama3 scaling instead of overwriting it with linear."
+
+    # yarn is not rebuildable by the patcher -> keep the safe linear fallback, not native.
+    yarn = SimpleNamespace(rope_scaling = {"rope_type": "yarn", "factor": 2.0}, rope_theta = 500000.0)
+    scaling, _ = _extended_rope_scaling(yarn, 2.0)
+    assert scaling == {
+        "type": "linear",
+        "factor": 2.0,
+        "rope_theta": 500000.0,
+    }, f"yarn must fall back to linear (patcher cannot rebuild it), got {scaling}."
+
+    # plain RoPE with theta only under v5 rope_parameters: linear must carry rope_theta.
+    v5 = SimpleNamespace(rope_parameters = {"rope_type": "default", "rope_theta": 1000000.0})
+    scaling, _ = _extended_rope_scaling(v5, 2.0)
+    assert scaling == {
+        "type": "linear",
+        "factor": 2.0,
+        "rope_theta": 1000000.0,
+    }, f"linear override dropped rope_theta on v5 (got {scaling}); base would fall back to 10000."
+
+
 def test_extended_rotary_reads_config_factor():
     # LlamaExtendedRotaryEmbedding must honor the config factor, not hardcode 8
     # (Llama-3.2 uses 32); otherwise the subclass path re-drops scaling (#2405).

tests/version_compat/test_sentence_transformers_pinned_symbols.py

@@ -19,6 +19,8 @@ ST_TAGS = [
     "v5.2.3",
     "v5.3.0",
     "v5.4.1",
+    "v5.5.1",
+    "v5.6.0",
     "master",
 ]
 
@@ -120,6 +122,42 @@ def test_st_transformer_base_class_either_path(tag: str):
     )
 
 
+# Transformer.load classmethod: unsloth builds saved-ST modules through it (#6881).
+@pytest.mark.parametrize("tag", ST_TAGS)
+def test_st_transformer_load_accepts_unsloth_kwargs(tag: str):
+    """unsloth builds saved ST models via Transformer.load(...) so the saved
+    modality_config is honored (#6881). If .load stops accepting the hub kwargs it
+    passes (and has no **kwargs), update the fix before it silently regresses. Not
+    locating .load is a SKIP (may be inherited); the live test guards the install."""
+    candidates = [
+        "sentence_transformers/models/Transformer.py",
+        "sentence_transformers/models/transformer.py",
+        "sentence_transformers/base/modules/transformer.py",
+        "sentence_transformers/base/modules/module.py",
+    ]
+    for p in candidates:
+        src = fetch_text("UKPLab/sentence-transformers", tag, p)
+        if src is None or not has_def(src, "load", "func"):
+            continue
+        m = re.search(r"def\s+load\s*\((.*?)\)\s*(?:->[^:]*)?:", src, re.S)
+        if m is None:
+            continue
+        sig = m.group(1)
+        accepts_var_kw = "**" in sig
+        missing = [
+            kw
+            for kw in ("token", "cache_folder", "revision", "trust_remote_code")
+            if not (accepts_var_kw or re.search(rf"\b{re.escape(kw)}\b", sig))
+        ]
+        assert not missing, (
+            f"{tag}: Transformer.load in {p} no longer accepts {missing} and has no "
+            f"**kwargs; update unsloth.models.sentence_transformer._create_transformer_module "
+            f"(#6881) before it silently falls back to Transformer(...)."
+        )
+        return
+    pytest.skip(f"{tag}: Transformer.load not locatable in {candidates} (may be inherited)")
+
+
 # sentence_transformers.util: import_from_string + load_dir_path helpers unsloth calls.
 @pytest.mark.parametrize("tag", ST_TAGS)
 def test_st_util_helpers(tag: str):

unsloth/_gpu_init.py

@@ -26,6 +26,7 @@ already_imported = [mod for mod in critical_modules if mod in sys.modules]
 # Fix some issues before importing other packages
 from .import_fixes import (
     fix_message_factory_issue,
+    fix_torch_check_is_size,
     check_fbgemm_gpu_version,
     disable_broken_causal_conv1d,
     disable_broken_vllm,
@@ -72,6 +73,7 @@ fix_bitsandbytes_rocm_arch_detection()
 disable_broken_causal_conv1d()
 disable_broken_vllm()
 fix_message_factory_issue()
+fix_torch_check_is_size()
 check_fbgemm_gpu_version()
 torchvision_compatibility_check()
 fix_diffusers_warnings()
@@ -81,6 +83,7 @@ del fix_bitsandbytes_rocm_arch_detection
 del disable_broken_causal_conv1d
 del disable_broken_vllm
 del fix_message_factory_issue
+del fix_torch_check_is_size
 del check_fbgemm_gpu_version
 del torchvision_compatibility_check
 del fix_diffusers_warnings
@@ -173,6 +176,7 @@ from .import_fixes import (
     fix_vllm_guided_decoding_params,
     fix_vllm_pdl_blackwell,
     fix_triton_compiled_kernel_missing_attrs,
+    fix_dynamo_config_thread_visibility,
     patch_trunc_normal_precision_issue,
     ignore_logger_messages,
     patch_ipykernel_hf_xet,
@@ -203,6 +207,10 @@ fix_vllm_guided_decoding_params()
 fix_trl_vllm_ascend()
 fix_vllm_pdl_blackwell()
 fix_triton_compiled_kernel_missing_attrs()
+# Must run before unsloth_zoo's patch_torch_compile and the gpt-oss temporary
+# patches raise the dynamo recompile limits, so those settings reach the
+# autograd worker threads on torch >= 2.12.
+fix_dynamo_config_thread_visibility()
 patch_trunc_normal_precision_issue()
 ignore_logger_messages()
 patch_ipykernel_hf_xet()
@@ -233,6 +241,7 @@ del fix_vllm_guided_decoding_params
 del fix_trl_vllm_ascend
 del fix_vllm_pdl_blackwell
 del fix_triton_compiled_kernel_missing_attrs
+del fix_dynamo_config_thread_visibility
 del patch_trunc_normal_precision_issue
 del ignore_logger_messages
 del patch_ipykernel_hf_xet

unsloth/import_fixes.py

@@ -172,6 +172,10 @@ if not UNSLOTH_ENABLE_LOGGING:
     # Deprecation warnings from torchao
     warnings.filterwarnings("ignore", message = "`int4_weight_only` is deprecated")
     warnings.filterwarnings("ignore", message = "`int8_weight_only` is deprecated")
+    # torch._check_is_size FutureWarning (called by bitsandbytes 4-bit dequant)
+    warnings.filterwarnings(
+        "ignore", message = r"_check_is_size will be removed", category = FutureWarning
+    )
 
     # TorchAO deprecated import paths (https://github.com/pytorch/ao/issues/2752)
     warnings.filterwarnings(
@@ -253,6 +257,30 @@ if not UNSLOTH_ENABLE_LOGGING:
     )
 
 
+def fix_torch_check_is_size():
+    """Shim torch._check_is_size if a future torch removes it (bitsandbytes 4-bit
+    dequant calls it). The FutureWarning is silenced in suppress_cuda_printf."""
+    try:
+        import torch
+
+        if hasattr(torch, "_check_is_size"):
+            return
+
+        def _check_is_size(
+            i,
+            message = None,
+            *,
+            max = None,
+        ):
+            torch._check(i >= 0, message)
+            if max is not None:
+                torch._check(i <= max, message)
+
+        torch._check_is_size = _check_is_size
+    except Exception:
+        return
+
+
 # Fix up AttributeError: 'MessageFactory' object has no attribute 'GetPrototype'
 # MUST do this at the start primarily due to tensorflow causing issues
 def fix_message_factory_issue():
@@ -1064,6 +1092,135 @@ def fix_triton_compiled_kernel_missing_attrs():
     )
 
 
+def fix_dynamo_config_thread_visibility():
+    """torch 2.12 made torch._dynamo/_inductor config overrides thread-local
+    (ContextVars), so `config.recompile_limit = 1024` set on the main thread is
+    invisible to the autograd worker threads that run backward. Gradient
+    checkpointing recompiles fullgraph gpt-oss kernels there against the default
+    limit of 8, raising FailOnRecompileLimitHit at step 0. Mirror direct config
+    assignments into the process-global entry default (torch <= 2.11 semantics).
+    config.patch(...) and config.load_config(...) also assign via __setattr__ but
+    are thread-local by design, so skip mirroring while inside one (tracked per
+    thread). No-op below torch 2.12 and on any torch without this internal layout.
+    """
+    try:
+        import torch
+
+        if Version(torch.__version__) < Version("2.12.0"):
+            return
+        import torch._dynamo.config as _dynamo_config
+        from torch.utils._config_module import ConfigModule
+        from contextvars import ContextVar
+    except Exception:
+        return
+
+    try:
+        probe = getattr(_dynamo_config, "_config", {}).get("recompile_limit", None)
+        if probe is None or not isinstance(getattr(probe, "user_override", None), ContextVar):
+            # Overrides are not context-local on this torch; nothing to fix.
+            return
+        original_setattr = ConfigModule.__setattr__
+        if getattr(original_setattr, "__unsloth_patched__", False):
+            return
+    except Exception:
+        return
+
+    mirrored_modules = ("torch._dynamo.config", "torch._inductor.config")
+
+    # config.patch(...) and config.load_config(...) also assign via __setattr__, but
+    # their writes are thread-local by design; a per-thread depth counter marks them
+    # so they are not mirrored into the process-global default.
+    import threading
+
+    _scoped_depth = threading.local()
+
+    def _in_scoped_write():
+        return getattr(_scoped_depth, "n", 0) > 0
+
+    def _bump(delta):
+        _scoped_depth.n = getattr(_scoped_depth, "n", 0) + delta
+
+    original_patch = ConfigModule.patch
+    if not getattr(original_patch, "__unsloth_patched__", False):
+
+        @functools.wraps(original_patch)
+        def _patched_patch(self, *args, **kwargs):
+            ctx = original_patch(self, *args, **kwargs)
+            try:
+                cls = type(ctx)  # patch() builds a fresh ConfigPatch class each call
+                if not getattr(cls, "__unsloth_patch_wrapped__", False):
+                    _enter0, _exit0 = cls.__enter__, cls.__exit__
+
+                    def _enter(s, _e = _enter0):
+                        _bump(1)
+                        try:
+                            return _e(s)
+                        finally:
+                            _bump(-1)
+
+                    def _exit(
+                        s,
+                        *a,
+                        _x = _exit0,
+                    ):
+                        _bump(1)
+                        try:
+                            return _x(s, *a)
+                        finally:
+                            _bump(-1)
+
+                    cls.__enter__, cls.__exit__ = _enter, _exit
+                    cls.__unsloth_patch_wrapped__ = True
+            except Exception:
+                pass
+            return ctx
+
+        _patched_patch.__unsloth_patched__ = True
+        ConfigModule.patch = _patched_patch
+
+    # load_config restores a saved config by calling setattr per key (thread-local).
+    original_load_config = getattr(ConfigModule, "load_config", None)
+    if callable(original_load_config) and not getattr(
+        original_load_config, "__unsloth_patched__", False
+    ):
+
+        @functools.wraps(original_load_config)
+        def _patched_load_config(self, *args, **kwargs):
+            _bump(1)
+            try:
+                return original_load_config(self, *args, **kwargs)
+            finally:
+                _bump(-1)
+
+        _patched_load_config.__unsloth_patched__ = True
+        ConfigModule.load_config = _patched_load_config
+
+    @functools.wraps(original_setattr)
+    def _patched_setattr(self, name, value):
+        original_setattr(self, name, value)
+        if _in_scoped_write():
+            return  # transient patch / load_config write: keep it thread-local
+        # Aliases (cache_size_limit -> recompile_limit) re-enter with the real name.
+        if self.__dict__.get("__name__", None) in mirrored_modules:
+            try:
+                entry = self.__dict__["_config"].get(name, None)
+                if entry is not None and entry.alias is None:
+                    entry.default = value
+            except Exception:
+                pass
+
+    _patched_setattr.__unsloth_patched__ = True
+    ConfigModule.__setattr__ = _patched_setattr
+
+    # No replay of existing overrides: unsloth installs this before it sets any
+    # dynamo/inductor config, so the wrapper mirrors every later assignment. Replaying
+    # would also bake a still-active config.patch override into the global default.
+    logger.info(
+        "Unsloth: Patched torch config modules so dynamo/inductor settings "
+        "(e.g. recompile_limit) apply across threads on torch >= 2.12."
+    )
+
+
 def patch_trunc_normal_precision_issue():
     """
     Patch torch.nn.init.trunc_normal_ for low precision tensors to run init in fp32.
@@ -1323,8 +1480,7 @@ def fix_vllm_pdl_blackwell():
 
     if patched:
         logger.info(
-            f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - "
-            f"patched: {', '.join(patched)}"
+            f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - patched: {', '.join(patched)}"
         )
     else:
         # Just set the env var - vLLM might be an older version without supports_pdl

unsloth/kernels/utils.py

@@ -282,6 +282,21 @@ def QUANT_STATE(W):
     return getattr(W, "quant_state", None)
 
 
+# fp8 weight dtypes. A `weight_scale` / `weight_scale_inv` should only be treated as a
+# quant state when the weight itself is still fp8. compressed-tensors layers expose an
+# already-dequantized bf16 weight at forward time while keeping a `weight_scale` around;
+# reading that as a quant state routes a bf16 weight into the bitsandbytes fast_gemv /
+# fast_dequantize path, which then reads a missing `absmax` and crashes.
+_FP8_WEIGHT_DTYPES = tuple(
+    dtype
+    for dtype in (
+        getattr(torch, "float8_e4m3fn", None),
+        getattr(torch, "float8_e5m2", None),
+    )
+    if dtype is not None
+)
+
+
 def get_lora_parameters(proj):
     """Return (weight, weight quant_state, lora A, lora B, lora scale).
     With QAT enabled, also fake-quantizes the base layer and lora weights.
@@ -298,9 +313,11 @@ def get_lora_parameters(proj):
         if weight_fake_quantizer is not None:
             W = weight_fake_quantizer(W)
 
-    # Get quant state for 4bit or FP8
+    # Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the
+    # weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer)
+    # must not carry a scale as its quant state or fast_gemv will crash on it.
     W_quant = getattr(W, "quant_state", None)
-    if W_quant is None:
+    if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES:
         W_quant = getattr(base_layer, "weight_scale_inv", None)
         if W_quant is None:
             W_quant = getattr(base_layer, "weight_scale", None)
@@ -349,9 +366,11 @@ def get_lora_parameters_bias(proj):
     )  # (proj.base_layer if hasattr(proj, "base_layer") else proj)
     W = base_layer.weight
 
-    # Get quant state for 4bit or FP8
+    # Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the
+    # weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer)
+    # must not carry a scale as its quant state or fast_gemv will crash on it.
     W_quant = getattr(W, "quant_state", None)
-    if W_quant is None:
+    if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES:
         W_quant = getattr(base_layer, "weight_scale_inv", None)
         if W_quant is None:
             W_quant = getattr(base_layer, "weight_scale", None)

unsloth/models/llama.py

@@ -28,7 +28,11 @@ from ._utils import (
     is_bfloat16_supported,
     get_quant_type,
 )
-from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings
+from .loader_utils import (
+    _exclude_rope_inv_freq_from_ddp,
+    _get_fp8_mode_and_check_settings,
+    _restore_dropped_fp8_scales,
+)
 from ..utils.packing import (
     get_packed_info_from_kwargs,
     mask_packed_sequence_boundaries,
@@ -1651,6 +1655,26 @@ def _rope_scaling_as_dict(rope_scaling):
         return {}
 
 
+def _extended_rope_scaling(config, factor):
+    """RoPE scaling to extend a model past its native window. Keeps native llama3 as-is
+    (linear extension is far worse for long context); everything else gets linear. Returns
+    (scaling_or_None, type): None keeps llama3. The linear dict carries rope_theta so
+    transformers v5 (which stores it under rope_parameters) keeps the real base, not 10000.
+    Only llama3 is preserved because patch_llama_rope_scaling can only rebuild linear/llama3/
+    longrope and its longrope branch needs a top-level original_max_position_embeddings."""
+    existing = _rope_scaling_as_dict(
+        getattr(config, "rope_scaling", None) or getattr(config, "rope_parameters", None) or {}
+    )
+    existing_type = existing.get("rope_type") or existing.get("type")
+    if existing_type == "llama3":
+        return None, existing_type
+    return {
+        "type": "linear",
+        "factor": factor,
+        "rope_theta": _get_rope_theta(config),
+    }, existing_type
+
+
 def _llama3_inv_freq_from_config(
     config,
     rope_scaling,
@@ -2518,34 +2542,33 @@ class FastLlamaModel:
             max_seq_length = model_max_seq_length
 
         if (rope_scaling is None) and (max_seq_length > model_max_seq_length):
-            rope_scaling = max_seq_length / model_max_seq_length
+            factor = max_seq_length / model_max_seq_length
 
             if fast_inference:
                 raise NotImplementedError(
                     "Unsloth: Fast inference does not yet work with RoPE Scaling."
                 )
 
-            logger.warning_once(
-                f"Unsloth: {model_name} can only handle sequence lengths of at most "
-                f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of "
-                f"{round(rope_scaling, 3)}, it can be magically be extended to "
-                f"{max_seq_length}!"
-            )
-
-            # Warn RoPE scaling isn't allowed
-            if not has_rope_scaling:
-                raise RuntimeError(
-                    f"However, {model_name} doesn't support RoPE Scaling!\n"
-                    "Please file a feature request at https://github.com/unslothai/unsloth."
+            linear_scaling, native_type = _extended_rope_scaling(model_config, factor)
+            if linear_scaling is not None:
+                logger.warning_once(
+                    f"Unsloth: {model_name} can only handle sequence lengths of at most "
+                    f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of "
+                    f"{round(factor, 3)}, it can be magically be extended to "
+                    f"{max_seq_length}!"
+                )
+                if not has_rope_scaling:
+                    raise RuntimeError(
+                        f"However, {model_name} doesn't support RoPE Scaling!\n"
+                        "Please file a feature request at https://github.com/unslothai/unsloth."
+                    )
+                kwargs["rope_scaling"] = linear_scaling
+            else:
+                # Native llama3 scaling already handles long context; just widen the window.
+                logger.warning_once(
+                    f"Unsloth: extending {model_name} to {max_seq_length} using its native "
+                    f"{native_type} RoPE scaling."
                 )
-
-            rope_scaling = {
-                "type": "linear",
-                "factor": rope_scaling,
-            }
-
-            # Add to kwargs
-            kwargs["rope_scaling"] = rope_scaling
 
         from .loader_utils import (
             check_and_disable_bitsandbytes_loading,
@@ -2659,6 +2682,18 @@ class FastLlamaModel:
                     offload_embedding = False,
                     fast_inference = fast_inference,
                 )
+                # Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200).
+                _restore_dropped_fp8_scales(
+                    model,
+                    model_name,
+                    local_files_only = kwargs.get("local_files_only", False),
+                    token = token,
+                    # Weights load from the default branch (revision not forwarded), so read scales from there too.
+                    revision = None,
+                    subfolder = kwargs.get("subfolder"),
+                    cache_dir = kwargs.get("cache_dir"),
+                    variant = kwargs.get("variant"),
+                )
             elif not fast_inference:
                 if user_config is not None:
                     # Transformers 5.x @strict model init rejects extra kwargs next
@@ -2697,6 +2732,18 @@ class FastLlamaModel:
                     offload_embedding = False,
                     fast_inference = False,
                 )
+                # Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200).
+                _restore_dropped_fp8_scales(
+                    model,
+                    model_name,
+                    local_files_only = kwargs.get("local_files_only", False),
+                    token = token,
+                    # Weights load from the default branch (revision not forwarded), so read scales from there too.
+                    revision = None,
+                    subfolder = kwargs.get("subfolder"),
+                    cache_dir = kwargs.get("cache_dir"),
+                    variant = kwargs.get("variant"),
+                )
                 model.fast_generate = make_fast_generate_wrapper(model.generate)
                 model.fast_generate_batches = None
             else:

unsloth/models/loader_utils.py

@@ -367,6 +367,287 @@ def _tag_model_with_fp8_torchao_config(model: torch.nn.Module, fp8_mode: str):
         pass
 
 
+_FP8_DTYPES = tuple(
+    dtype
+    for dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e5m2", None))
+    if dtype is not None
+)
+
+
+def _fp8_block_size_from_config(model):
+    """Return the [block_out, block_in] block size of an fp8 checkpoint, or None if not block-fp8."""
+    config = getattr(model, "config", None)
+    quant = getattr(config, "quantization_config", None)
+    if quant is None:
+        return None
+    if hasattr(quant, "to_dict"):
+        quant = quant.to_dict()
+    if not isinstance(quant, dict):
+        return None
+    if quant.get("quant_method") != "fp8":
+        return None
+    block = quant.get("weight_block_size")
+    if not block:
+        return None
+    if isinstance(block, (int, float)):
+        block = [block, block]
+    elif isinstance(block, (list, tuple)):
+        if len(block) == 1:
+            block = [block[0], block[0]]
+        elif len(block) < 2:
+            return None
+    else:
+        return None
+    return [int(block[0]), int(block[1])]
+
+
+def _load_fp8_weight_map(
+    model_name,
+    local_files_only,
+    token,
+    revision = None,
+    subfolder = None,
+    cache_dir = None,
+):
+    """Return the checkpoint's tensor->file map, using the same snapshot the load used.
+
+    Prefers the sharded `model.safetensors.index.json`; falls back to a single `model.safetensors`
+    (every tensor maps to that one file) so unsharded checkpoints are covered too.
+    """
+
+    def _local_path(filename):
+        return (
+            os.path.join(model_name, subfolder, filename)
+            if subfolder
+            else os.path.join(model_name, filename)
+        )
+
+    def _remote_path(filename):
+        from huggingface_hub import hf_hub_download
+        return hf_hub_download(
+            model_name,
+            filename,
+            revision = revision,
+            subfolder = subfolder,
+            cache_dir = cache_dir,
+            local_files_only = local_files_only,
+            token = token,
+        )
+
+    index_file = "model.safetensors.index.json"
+    single_file = "model.safetensors"
+    is_local = os.path.isdir(model_name)
+
+    # Sharded checkpoint.
+    if is_local and os.path.exists(_local_path(index_file)):
+        index_path = _local_path(index_file)
+    elif not is_local:
+        try:
+            index_path = _remote_path(index_file)
+        except Exception:
+            index_path = None
+    else:
+        index_path = None
+    if index_path is not None:
+        import json
+        with open(index_path, "r") as f:
+            return json.load(f).get("weight_map", None)
+
+    # Unsharded single file: map every tensor to it.
+    try:
+        if is_local and os.path.exists(_local_path(single_file)):
+            single_path = _local_path(single_file)
+        elif not is_local:
+            single_path = _remote_path(single_file)
+        else:
+            return None
+        from safetensors import safe_open
+        with safe_open(single_path, framework = "pt") as f:
+            return {key: single_file for key in f.keys()}
+    except Exception:
+        return None
+
+
+def _resolve_fp8_shard(
+    model_name,
+    shard,
+    local_files_only,
+    token,
+    revision = None,
+    subfolder = None,
+    cache_dir = None,
+):
+    """Resolve a checkpoint shard filename to a local path (repo id or local dir)."""
+    if os.path.isdir(model_name):
+        return (
+            os.path.join(model_name, subfolder, shard)
+            if subfolder
+            else os.path.join(model_name, shard)
+        )
+    from huggingface_hub import hf_hub_download
+
+    return hf_hub_download(
+        model_name,
+        shard,
+        revision = revision,
+        subfolder = subfolder,
+        cache_dir = cache_dir,
+        local_files_only = local_files_only,
+        token = token,
+    )
+
+
+def _match_fp8_module(module_by_name, base):
+    """Resolve a checkpoint module name to a live module, allowing for VLM key remappings.
+
+    VLM loads can name the text tower differently from the checkpoint keys: `text_only=True`
+    strips the `language_model.` wrapper (so `model.language_model.layers.*` -> `model.layers.*`),
+    and full VLM loads may expose `model.language_model.*` while the checkpoint stores
+    `language_model.model.*`. Try the raw key first, then a few safe remappings.
+    """
+    if base in module_by_name:
+        return module_by_name[base]
+    candidates = []
+    if "language_model." in base:
+        candidates.append(base.replace("language_model.", "", 1))  # text-only: drop wrapper
+    if "language_model.model." in base:
+        candidates.append(base.replace("language_model.model.", "model.language_model.", 1))
+    if base.startswith("language_model."):
+        candidates.append("model." + base)  # add model. prefix
+    for candidate in candidates:
+        if candidate in module_by_name:
+            return module_by_name[candidate]
+    return None
+
+
+def _restore_dropped_fp8_scales(
+    model,
+    model_name,
+    *,
+    local_files_only = False,
+    token = None,
+    revision = None,
+    subfolder = None,
+    cache_dir = None,
+    variant = None,
+):
+    """Re-apply block-fp8 `weight_scale_inv` tensors that transformers dropped on load.
+
+    On some block-scale fp8 checkpoints (e.g. Qwen3.6-27B-FP8, issue #6200) transformers fails to
+    convert a Linear (such as `mlp.gate_proj`) to an fp8 module, loading the raw quantized values
+    into a plain bf16 weight and discarding its `weight_scale_inv` as an unexpected key. The weight
+    is then used un-scaled, producing a garbage model. For every checkpoint scale whose live weight
+    is not fp8, dequantize the orphaned weight in place. Modules that were converted correctly keep
+    an fp8 weight and are skipped, so a healthy checkpoint is a no-op. Returns (restored, skipped).
+    """
+    try:
+        block = _fp8_block_size_from_config(model)
+        if block is None or not _FP8_DTYPES:
+            return (0, 0)
+        # A variant load reads variant-named files; skip to avoid applying default scales to them.
+        if variant:
+            return (0, 0)
+        # No fp8 params means the checkpoint was dequantized on purpose (e.g. load_in_16bit);
+        # re-applying a scale would corrupt those already-correct 16bit weights, so do nothing.
+        if not any(p.dtype in _FP8_DTYPES for p in model.parameters()):
+            return (0, 0)
+        weight_map = _load_fp8_weight_map(
+            model_name, local_files_only, token, revision, subfolder, cache_dir
+        )
+        if not weight_map:
+            return (0, 0)
+
+        scale_keys = {k: v for k, v in weight_map.items() if k.endswith(".weight_scale_inv")}
+        if not scale_keys:
+            return (0, 0)
+
+        module_by_name = dict(model.named_modules())
+        bs0, bs1 = block
+        restored = 0
+        skipped = 0
+        failed = 0
+        offloaded = 0
+        shard_cache = {}
+        for scale_key, shard in scale_keys.items():
+            base = scale_key[: -len(".weight_scale_inv")]
+            module = _match_fp8_module(module_by_name, base)
+            if module is None:
+                continue
+            weight = getattr(module, "weight", None)
+            if not isinstance(weight, torch.Tensor) or weight.ndim != 2:
+                continue
+            if weight.device.type == "meta":
+                # Disk-offloaded layer: weight lives on meta until forward, so it cannot be
+                # scaled in place here. Count and warn rather than silently leave it unscaled.
+                offloaded += 1
+                continue
+            if weight.dtype in _FP8_DTYPES:
+                # Correctly converted fp8 module: the fp8 path already handles the scale.
+                skipped += 1
+                continue
+
+            # Errors after this point are per-tensor: warn and continue, never abort or hide them.
+            try:
+                if shard not in shard_cache:
+                    from safetensors import safe_open
+                    shard_path = _resolve_fp8_shard(
+                        model_name,
+                        shard,
+                        local_files_only,
+                        token,
+                        revision,
+                        subfolder,
+                        cache_dir,
+                    )
+                    shard_cache[shard] = safe_open(shard_path, framework = "pt")
+                scale = shard_cache[shard].get_tensor(scale_key).to(torch.float32)
+
+                out_features, in_features = weight.shape
+                out_blocks = (out_features + bs0 - 1) // bs0
+                in_blocks = (in_features + bs1 - 1) // bs1
+                if tuple(scale.shape) == (out_blocks, in_blocks):
+                    pass
+                elif tuple(scale.shape) == (in_blocks, out_blocks) and out_blocks != in_blocks:
+                    # Transposed block layout: same handling as the fp8 forward path.
+                    scale = scale.t().contiguous()
+                else:
+                    # Shape does not match the block grid: skip rather than apply a wrong scale.
+                    continue
+                scale = scale.to(weight.device)
+                with torch.no_grad():
+                    if out_features % bs0 == 0 and in_features % bs1 == 0:
+                        # Memory-frugal path: multiply block views in place against the broadcast
+                        # fp32 scale, avoiding a full expanded scale and fp32 copy that could OOM.
+                        # The in-place multiply promotes to fp32, matching the fallback exactly.
+                        module.weight.data.view(out_blocks, bs0, in_blocks, bs1).mul_(
+                            scale[:, None, :, None]
+                        )
+                    else:
+                        scale_expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(
+                            bs1, dim = 1
+                        )[:out_features, :in_features]
+                        module.weight.data = (weight.to(torch.float32) * scale_expanded).to(
+                            weight.dtype
+                        )
+                restored += 1
+            except Exception:
+                failed += 1
+                continue
+
+        if restored > 0:
+            print(f"Unsloth: Restored {restored} dropped FP8 weight_scale_inv tensor(s) on load")
+        if failed > 0:
+            print(f"Unsloth: {failed} dropped FP8 weight_scale_inv tensor(s) could not be restored")
+        if offloaded > 0:
+            print(
+                f"Unsloth: {offloaded} dropped FP8 weight_scale_inv tensor(s) skipped because the "
+                "layer is disk-offloaded; load without disk offload so the scales can be restored"
+            )
+        return (restored, skipped)
+    except Exception:
+        return (0, 0)
+
+
 def check_and_disable_bitsandbytes_loading(
     model_config,
     load_in_4bit = True,

unsloth/models/sentence_transformer.py

@@ -990,7 +990,17 @@ class FastSentenceTransformer(FastModel):
             return None
 
     @staticmethod
-    def _create_transformer_module(model_name, model, tokenizer, max_seq_length, trust_remote_code):
+    def _create_transformer_module(
+        model_name,
+        model,
+        tokenizer,
+        max_seq_length,
+        trust_remote_code,
+        token = None,
+        cache_dir = None,
+        revision = None,
+        module_subfolder = "",
+    ):
         """Helper to create and configure a Transformer module."""
         from sentence_transformers.models import Transformer
 
@@ -1077,7 +1087,45 @@ class FastSentenceTransformer(FastModel):
                 elif "tokenizer_args" in transformer_init_params:
                     transformer_kwargs["tokenizer_args"] = trust_remote_code_kwargs.copy()
 
-                transformer_module = Transformer(model_name, **transformer_kwargs)
+                # Build via Transformer.load so the saved modality_config is honored: plain
+                # Transformer(...) makes ST 5.x infer a "message" modality for chat-template
+                # models (e.g. Qwen3-Embedding), chat-wrapping inputs and degrading embeddings
+                # (#6881). Only use .load when it resolves a Hub id (accepts the kwargs or
+                # **kwargs); legacy ST 3.x/4.x load(input_path) is local-only with no modality
+                # bug, so fall back to the constructor.
+                transformer_module = None
+                transformer_load = getattr(Transformer, "load", None)
+                has_modules_json = (
+                    FastSentenceTransformer._module_path(
+                        model_name, token, cache_dir = cache_dir, revision = revision
+                    )
+                    is not None
+                )
+                if callable(transformer_load) and has_modules_json:
+                    load_params = inspect.signature(transformer_load).parameters
+                    accepts_var_kw = any(
+                        p.kind is inspect.Parameter.VAR_KEYWORD for p in load_params.values()
+                    )
+                    hub_capable = accepts_var_kw or any(
+                        key in load_params for key in ("token", "cache_folder", "revision")
+                    )
+                    if hub_capable:
+                        load_kwargs = {
+                            "token": token,
+                            "cache_folder": cache_dir,
+                            "revision": revision,
+                            "trust_remote_code": trust_remote_code,
+                            **transformer_kwargs,
+                        }
+                        # Resolve config/tokenizer from the module's saved subfolder
+                        # (modules.json "path"), like stock ST; "" (root) is a no-op.
+                        if module_subfolder:
+                            load_kwargs["subfolder"] = module_subfolder
+                        if not accepts_var_kw:
+                            load_kwargs = {k: v for k, v in load_kwargs.items() if k in load_params}
+                        transformer_module = Transformer.load(model_name, **load_kwargs)
+                if transformer_module is None:
+                    transformer_module = Transformer(model_name, **transformer_kwargs)
             finally:
                 # Restore original Auto* loading immediately
                 AutoModel.from_pretrained = original_model_from_pretrained
@@ -1191,6 +1239,10 @@ class FastSentenceTransformer(FastModel):
                         tokenizer,
                         max_seq_length,
                         trust_remote_code,
+                        token,
+                        cache_dir,
+                        revision,
+                        module_subfolder = module_config.get("path") or "",
                     )
                     modules[name] = transformer_module
                 else:
@@ -1226,7 +1278,14 @@ class FastSentenceTransformer(FastModel):
         )
 
         transformer_module = FastSentenceTransformer._create_transformer_module(
-            model_name, model, tokenizer, max_seq_length, trust_remote_code
+            model_name,
+            model,
+            tokenizer,
+            max_seq_length,
+            trust_remote_code,
+            token,
+            cache_dir,
+            revision,
         )
         modules["0"] = transformer_module
 

unsloth/models/vision.py

@@ -41,7 +41,11 @@ from ._utils import (
     set_task_config_attr,
 )
 from ._utils import *
-from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings
+from .loader_utils import (
+    _exclude_rope_inv_freq_from_ddp,
+    _get_fp8_mode_and_check_settings,
+    _restore_dropped_fp8_scales,
+)
 from ..save import patch_saving_functions
 from ..models.loader_utils import is_distributed
 from unsloth_zoo.gradient_checkpointing import (
@@ -1192,6 +1196,17 @@ class FastBaseModel:
                     offload_embedding = offload_embedding,
                     fast_inference = fast_inference,
                 )
+                # Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200).
+                _restore_dropped_fp8_scales(
+                    model,
+                    model_name,
+                    local_files_only = local_files_only,
+                    token = token,
+                    revision = kwargs.get("revision"),
+                    subfolder = kwargs.get("subfolder"),
+                    cache_dir = kwargs.get("cache_dir"),
+                    variant = kwargs.get("variant"),
+                )
                 if hasattr(model, "generate"):
                     model.fast_generate = make_fast_generate_wrapper(model.generate)
                     model.fast_generate_batches = error_out_no_vllm

unsloth_cli/commands/start.py

@@ -133,6 +133,21 @@ _YOLO_OPTION = typer.Option(
         "flag/config. Any of the three spellings works for any agent."
     ),
 )
+_PERSIST_OPTION = typer.Option(
+    False,
+    "--persist/--no-persist",
+    help = (
+        "Keep this agent's Unsloth-managed session dir so you can resume it later. "
+        "codex/openclaw/hermes/pi have their whole home relocated into an Unsloth dir "
+        "that is a throwaway temp dir (wiped on exit) by default; with --persist it "
+        "lives under the Unsloth agents dir and survives, so their own resume can reopen "
+        "it. claude and opencode keep sessions in your own stores (~/.claude, "
+        "~/.local/share/opencode), so they already resume regardless. To reopen a "
+        "session, pass the agent's own resume command through, e.g. "
+        "`unsloth start codex --persist resume` or `claude --resume <id>`; those flow to "
+        "the agent unchanged."
+    ),
+)
 
 # Per-agent CLI flag for "run tools without prompting". opencode and openclaw have no
 # such flag (config only) and are handled in their config writers, so they are absent.
@@ -989,6 +1004,12 @@ def _refresh_windows_path() -> None:
         os.environ["PATH"] = os.pathsep.join(entries)
 
 
+def _install_source(install_hint: str) -> Optional[str]:
+    """The first http(s) URL an install hint fetches, or None (e.g. an npm install)."""
+    match = re.search(r"https?://[^\s'\")]+", install_hint)
+    return match.group(0) if match else None
+
+
 def _install_agent(name: str, install_hint: str) -> Optional[str]:
     # Missing agent under --launch: offer to run its documented install command, then
     # re-resolve it on PATH. Consent-based (we never auto-run a remote install script
@@ -997,7 +1018,18 @@ def _install_agent(name: str, install_hint: str) -> Optional[str]:
     if not sys.stdin.isatty():
         return None
     typer.echo(f"`{name}` is not installed.")
-    if not typer.confirm(f"Install it now with `{install_hint}`?", default = False):
+    # Make the supply-chain risk explicit before the prompt: these are the vendors'
+    # own installers (curl | bash, irm | iex, npm), run with the user's privileges,
+    # and nothing checks a signature or hash on the fetched content. Naming the source
+    # turns a blind "yes" into informed consent.
+    source = _install_source(install_hint)
+    warning = (
+        f"This will download and RUN a script from {source} with your privileges"
+        if source
+        else f"This will RUN `{install_hint}` with your privileges"
+    )
+    typer.secho(f"{warning}; there is no signature or hash check.", fg = "yellow", err = True)
+    if not typer.confirm(f"Install `{name}` now with `{install_hint}`?", default = False):
         return None
     # Run each hint through the shell it is written for: PowerShell (irm | iex, or npm)
     # on Windows, /bin/sh (curl | bash, or npm) everywhere else.
@@ -1116,15 +1148,20 @@ def _agents_config_root() -> Path:
 
 
 @contextlib.contextmanager
-def _session_config(agent: str, launch: bool):
+def _session_config(
+    agent: str,
+    launch: bool,
+    persist: bool = False,
+):
     """Yield a private directory for an agent's session config (never the user's own).
 
-    launch: an ephemeral temp dir removed after the agent process exits, so nothing
-    persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run later
-    on this machine), reused across runs. Either way the user's real ~/.<agent>
-    config is left untouched.
+    launch (default): an ephemeral temp dir removed after the agent process exits, so
+    nothing persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run
+    later on this machine), reused across runs. persist (from --persist): use that same
+    stable dir even for a launch, so the agent's session survives the exit and can be
+    resumed next time. Either way the user's real ~/.<agent> config is left untouched.
     """
-    if launch:
+    if launch and not persist:
         path = Path(tempfile.mkdtemp(prefix = f"unsloth-{agent}-"))
         try:
             yield path
@@ -1436,6 +1473,7 @@ def claude(
     tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
     serve: bool = _SERVE_OPTION,
     yolo: bool = _YOLO_OPTION,
+    persist: bool = _PERSIST_OPTION,
 ):
     """Point Claude Code at the running Studio server and start it."""
     base, key, entry = _connect(
@@ -1480,6 +1518,9 @@ def claude(
     # --yolo (or its aliases) maps to Claude's own --dangerously-skip-permissions.
     # IS_SANDBOX is left unset on purpose: Claude refuses bypass mode as root unless a
     # sandbox is detected, and we don't want to falsely claim one on the user's host.
+    # claude keeps its history in ~/.claude/projects, which --settings/env never
+    # relocate, so a session already survives exit; resume it with `claude --continue`
+    # or `--resume <id>` passed through.
     command = [
         "claude",
         "--model",
@@ -1516,6 +1557,7 @@ def codex(
     tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
     serve: bool = _SERVE_OPTION,
     yolo: bool = _YOLO_OPTION,
+    persist: bool = _PERSIST_OPTION,
 ):
     """Point OpenAI Codex at the running Studio server and start it."""
     base, key, entry = _connect(
@@ -1541,7 +1583,7 @@ def codex(
         *_yolo_command_flags("codex", yolo),
         *ctx.args,
     ]
-    with _session_config("codex", launch) as home:
+    with _session_config("codex", launch, persist = persist) as home:
         write_codex_config(base, entry, home)
         env = {_CODEX_ENV_KEY: key, "CODEX_HOME": str(home)}
         _run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex")
@@ -1559,6 +1601,7 @@ def openclaw(
     tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
     serve: bool = _SERVE_OPTION,
     yolo: bool = _YOLO_OPTION,
+    persist: bool = _PERSIST_OPTION,
 ):
     """Point OpenClaw at the running Studio server and start it."""
     base, key, entry = _connect(
@@ -1584,7 +1627,7 @@ def openclaw(
         if os.name == "nt"
         else "curl -fsSL https://openclaw.ai/install.sh | bash"
     )
-    with _session_config("openclaw", launch) as cfg:
+    with _session_config("openclaw", launch, persist = persist) as cfg:
         config_path = cfg / "openclaw.json"
         # key lives in the config, not the env; --yolo writes the exec policy here too.
         write_openclaw_config(base, key, entry, config_path, yolo = yolo)
@@ -1605,6 +1648,7 @@ def opencode(
     tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
     serve: bool = _SERVE_OPTION,
     yolo: bool = _YOLO_OPTION,
+    persist: bool = _PERSIST_OPTION,
 ):
     """Point OpenCode at the running Studio server and start it."""
     base, key, entry = _connect(
@@ -1628,7 +1672,9 @@ def opencode(
         command = ["opencode", "--model", opencode_model]
     else:
         command = ["opencode"]
-    with _session_config("opencode", launch) as cfg:
+    # opencode keeps sessions in ~/.local/share/opencode (never relocated), so resume
+    # already survives exit; reopen the last one by passing `opencode --continue` through.
+    with _session_config("opencode", launch, persist = persist) as cfg:
         config_path = cfg / "opencode.json"
         # OPENCODE_CONFIG is an overlay (loaded between the user's global and project
         # configs), so this adds the Unsloth provider/model for the session without
@@ -1680,6 +1726,7 @@ def hermes(
     tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
     serve: bool = _SERVE_OPTION,
     yolo: bool = _YOLO_OPTION,
+    persist: bool = _PERSIST_OPTION,
 ):
     """Point Hermes (Nous Research) at the running Studio server and start it."""
     base, key, entry = _connect(
@@ -1691,7 +1738,7 @@ def hermes(
     )
     command = ["hermes", *_yolo_command_flags("hermes", yolo), *ctx.args]
     install_hint = _hermes_install_hint()
-    with _session_config("hermes", launch) as home:
+    with _session_config("hermes", launch, persist = persist) as home:
         # HERMES_HOME relocates hermes' whole home dir (config.yaml, sessions, state)
         # like CODEX_HOME, so the user's ~/.hermes is left untouched for the session.
         write_hermes_config(base, entry, home / "config.yaml")
@@ -1711,6 +1758,7 @@ def pi(
     tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
     serve: bool = _SERVE_OPTION,
     yolo: bool = _YOLO_OPTION,
+    persist: bool = _PERSIST_OPTION,
 ):
     """Point Pi (coding agent) at the running Studio server and start it."""
     base, key, entry = _connect(
@@ -1735,7 +1783,7 @@ def pi(
     # --ignore-scripts matches Pi's documented install recipe (its README notes Pi needs
     # no install scripts), so accepting the prompt skips dependency lifecycle scripts.
     install_hint = "npm install -g --ignore-scripts @earendil-works/pi-coding-agent"
-    with _session_config("pi", launch) as home:
+    with _session_config("pi", launch, persist = persist) as home:
         # Pi resolves its config dir from PI_CODING_AGENT_DIR first (getAgentDir() prefers
         # it over $HOME/.pi/agent), so pin it at the session dir: an inherited
         # PI_CODING_AGENT_DIR in the user's shell would otherwise send Pi to their real

unsloth_cli/tests/test_start.py

@@ -128,6 +128,32 @@ def test_install_agent_uses_powershell_on_windows(monkeypatch):
     assert ran == [["powershell", "-NoProfile", "-Command", install_hint]]
 
 
+def test_install_agent_warns_and_names_remote_source(monkeypatch, capsys):
+    # Before the confirm, a remote installer must name the URL it fetches so the
+    # user consents to a specific source rather than blindly accepting.
+    monkeypatch.setattr(start.os, "name", "nt")
+    monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True))
+    monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False)  # decline: nothing runs
+    hint = "& ([scriptblock]::Create((irm https://hermes-agent.nousresearch.com/install.ps1))) -SkipSetup"
+    assert start._install_agent("hermes", hint) is None
+    err = capsys.readouterr().err
+    assert "https://hermes-agent.nousresearch.com/install.ps1" in err
+    assert "download and RUN" in err
+    assert "signature or hash" in err
+
+
+def test_install_agent_warns_for_package_installer(monkeypatch, capsys):
+    # An npm-style installer has no URL to fetch, but still runs with the user's
+    # privileges, so the warning names the command instead.
+    monkeypatch.setattr(start.os, "name", "posix")
+    monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True))
+    monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False)
+    assert start._install_agent("codex", "npm install -g @openai/codex") is None
+    err = capsys.readouterr().err
+    assert "npm install -g @openai/codex" in err
+    assert "with your privileges" in err
+
+
 def test_hermes_install_hint_is_windows_native_on_windows(monkeypatch):
     monkeypatch.setattr(start.os, "name", "nt")
 
@@ -2522,3 +2548,145 @@ def test_session_config_no_launch_preserves_existing_state(fake_studio, tmp_path
     with start._session_config("codex", launch = False) as home2:
         assert home2 == home
         assert (home2 / "sessions" / "live.sqlite").read_text() == "state"
+
+
+# ── --persist: persist the agent session so it can be resumed ────────────────
+def test_session_config_persist_uses_stable_dir_and_survives(monkeypatch, tmp_path):
+    # --persist routes a launch to the stable Unsloth agents dir (the one --no-launch
+    # already uses) instead of a throwaway temp dir, and never wipes it on exit.
+    monkeypatch.setattr(start, "_agents_config_root", lambda: tmp_path / "agents")
+    with start._session_config("codex", launch = True, persist = True) as home:
+        assert home == tmp_path / "agents" / "codex"
+        (home / "marker").write_text("kept")
+    assert home.exists()
+    assert (home / "marker").read_text() == "kept"
+
+
+def test_session_config_default_launch_is_ephemeral():
+    # Default launch (no --persist) still uses a throwaway temp dir wiped on exit.
+    with start._session_config("codex", launch = True) as home:
+        assert home.exists()
+        assert "unsloth-codex-" in home.name
+    assert not home.exists()
+
+
+# The temp-dir agents: --persist points each one's home/state env at the stable dir;
+# without it, at an ephemeral temp path. opencode is handled separately (only its
+# config overlay is relocated; its session data was never in the temp dir).
+_RESUME_ENV_VAR = {
+    "codex": "CODEX_HOME",
+    "openclaw": "OPENCLAW_STATE_DIR",
+    "hermes": "HERMES_HOME",
+    "pi": "HOME",
+}
+
+
+def _capture_launch(monkeypatch, argv):
+    captured = {}
+
+    def run(
+        command,
+        env = None,
+        **kwargs,
+    ):
+        captured["command"] = command
+        captured["env"] = env
+        return SimpleNamespace(returncode = 0)
+
+    monkeypatch.setattr(start.subprocess, "run", run)
+    result = CliRunner().invoke(start.start_app, argv)
+    assert result.exit_code == 0, result.output
+    return captured
+
+
+@pytest.mark.parametrize("agent", sorted(_RESUME_ENV_VAR))
+def test_resume_persists_agent_home_to_stable_dir(agent, fake_studio, tmp_path, monkeypatch):
+    monkeypatch.setattr(start.shutil, "which", lambda _: f"/usr/local/bin/{agent}")
+    captured = _capture_launch(monkeypatch, [agent, "--persist"])
+    stable = tmp_path / "agents" / agent
+    assert captured["env"][_RESUME_ENV_VAR[agent]] == str(stable)
+    # The stable dir survives the agent exit, so the session can be resumed.
+    assert stable.exists()
+
+
+@pytest.mark.parametrize("agent", sorted(_RESUME_ENV_VAR))
+def test_default_launch_home_is_ephemeral(agent, fake_studio, tmp_path, monkeypatch):
+    monkeypatch.setattr(start.shutil, "which", lambda _: f"/usr/local/bin/{agent}")
+    captured = _capture_launch(monkeypatch, [agent])
+    home = captured["env"][_RESUME_ENV_VAR[agent]]
+    assert f"unsloth-{agent}-" in home
+    assert str(tmp_path / "agents") not in home
+
+
+def test_resume_opencode_config_in_stable_dir(fake_studio, tmp_path, monkeypatch):
+    # opencode's session data lives in ~/.local/share/opencode (never relocated), so
+    # resume already survives exit; --persist also stabilizes its config overlay dir.
+    monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode")
+    captured = _capture_launch(monkeypatch, ["opencode", "--persist"])
+    stable = tmp_path / "agents" / "opencode"
+    assert captured["env"]["OPENCODE_CONFIG"] == str(stable / "opencode.json")
+    assert stable.exists()
+
+
+def test_persist_bare_codex_launch_has_no_resume_token(fake_studio, monkeypatch):
+    # A bare `--persist` only persists the session dir; it must NOT auto-append a native
+    # resume token, or the very first launch (no session yet) would send codex down its
+    # no-session error path. The user resumes explicitly: `unsloth start codex --persist resume`.
+    monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex")
+    captured = _capture_launch(monkeypatch, ["codex", "--persist"])
+    assert "resume" not in captured["command"]
+    # command[0] is the resolved executable path; assert the argv after it.
+    assert captured["command"][1:] == ["--oss", "--profile", start._CODEX_PROFILE]
+
+
+def test_persist_bare_opencode_launch_has_no_resume_token(fake_studio, monkeypatch):
+    monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode")
+    captured = _capture_launch(monkeypatch, ["opencode", "--persist"])
+    assert "--continue" not in captured["command"]
+    assert captured["command"][1:] == ["--model", f"{start._OPENCODE_PROVIDER}/{MODEL['id']}"]
+
+
+def test_persist_bare_claude_launch_has_no_resume_token(fake_studio, monkeypatch):
+    monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
+    monkeypatch.setattr(start, "_claude_flags", lambda: [])
+    captured = _capture_launch(monkeypatch, ["claude", "--persist"])
+    assert "--continue" not in captured["command"]
+    assert captured["command"][1:] == ["--model", MODEL["id"]]
+
+
+def test_resume_with_passthrough_does_not_auto_append(fake_studio, monkeypatch):
+    # When the caller drives their own subcommand, --persist only persists the dir; it
+    # must not inject a resume token that would collide with the user's command.
+    monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex")
+    captured = _capture_launch(monkeypatch, ["codex", "--persist", "exec", "hello"])
+    assert "resume" not in captured["command"]
+    assert captured["command"][-2:] == ["exec", "hello"]
+
+
+def test_default_launch_has_no_resume_token(fake_studio, monkeypatch):
+    monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex")
+    captured = _capture_launch(monkeypatch, ["codex"])
+    assert "resume" not in captured["command"]
+
+
+def test_resume_persist_only_agents_have_no_resume_token(fake_studio, monkeypatch):
+    # openclaw/hermes persist their session dir but have no non-interactive resume
+    # selector, so --persist must not append a token; their own picker resumes.
+    for agent in ("openclaw", "hermes"):
+        monkeypatch.setattr(start.shutil, "which", lambda _, a = agent: f"/usr/local/bin/{a}")
+        captured = _capture_launch(monkeypatch, [agent, "--persist"])
+        assert "resume" not in captured["command"]
+        assert "--continue" not in captured["command"]
+
+
+def test_native_resume_flag_passes_through_unchanged(fake_studio, monkeypatch):
+    # The persistence flag is --persist, NOT --resume, so an agent's own
+    # `--resume <id>` (e.g. `unsloth start claude --resume <guid>`) still flows
+    # through to the agent verbatim and is not swallowed as a Studio option.
+    monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
+    monkeypatch.setattr(start, "_claude_flags", lambda: [])
+    captured = _capture_launch(monkeypatch, ["claude", "--resume", "some-session-guid"])
+    assert captured["command"][-2:] == ["--resume", "some-session-guid"]
+    # Studio never auto-appends its own resume token when the user drives resume.
+    assert captured["command"].count("--resume") == 1
+    assert "--continue" not in captured["command"]