Studio: fix dataset upload data loss, caption over-count, and custom-root sd.cpp uninstall

- Diffusion dataset upload now streams each file into a sibling temp file and
  atomically os.replace()s it into place only after the whole file is written and
  within the size cap. A mid-batch 413 (or any abort) removes the temp, never an
  example already stored under the same name, so re-uploading a too-large batch can
  no longer truncate or delete a previously uploaded image.
- _diffusion_dataset_summary counts an image as captioned only when it resolves to a
  non-empty caption via the same sidecar-over-metadata precedence the trainer uses. An
  empty (tombstone) sidecar shadows a metadata row and makes the trainer skip the
  image, so counting it over-reported caption_count and mislabeled an effectively
  uncaptioned dataset as captioned.
- uninstall.sh/.ps1 now remove a custom/env-mode Studio's native diffusion build that
  installs beside the root as a stable-diffusion.cpp sibling (find_sd_cpp_binary
  resolves it from the Studio home's parent), guarded by the same unsafe-path check,
  and stop processes locking the default-mode stable-diffusion.cpp before removing it.

Adds regression tests for the upload data-loss and caption-count paths and a hermetic
shell test for the custom-root stable-diffusion.cpp removal.

scripts/uninstall.ps1

@@ -366,7 +366,7 @@ function Uninstall-UnslothStudio {
     _StopStudioProcesses -KnownRoots $knownRoots
     # Also stop anything holding a handle on the exact paths we delete (llama-server,
     # the CLI shim, an mp-fork python with a venv DLL) so the dir delete isn't refused.
-    _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache, $defaultNode))
+    _StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultSdCpp, $defaultCache, $defaultNode))
 
     # ── Remove custom-root install trees ──
     _Step "Removing data and install directories..."
@@ -380,6 +380,16 @@ function Uninstall-UnslothStudio {
             continue
         }
         _RemovePath $r
+        # Native diffusion (stable-diffusion.cpp) for a custom/env-mode Studio installs beside
+        # the root at <parent>\stable-diffusion.cpp -- find_sd_cpp_binary resolves it from
+        # UNSLOTH_STUDIO_HOME.parent (sd_cpp_engine.py) -- so removing only the root leaves the
+        # build behind. Derive and remove the sibling, guarding the parent path the same way.
+        $customSdCpp = Join-Path (Split-Path -LiteralPath $r -Parent) "stable-diffusion.cpp"
+        if (_IsUnsafeRoot $customSdCpp) {
+            _Substep "refusing to remove unsafe path: $customSdCpp" "Yellow"
+        } else {
+            _RemovePath $customSdCpp
+        }
     }
     # Default install dir (always at %USERPROFILE%\.unsloth\studio when present).
     if ($defaultStudioHome) { _RemovePath $defaultStudioHome }

scripts/uninstall.sh

@@ -210,6 +210,16 @@ _custom_studio_roots | while IFS= read -r _custom_root; do
         continue
     fi
     _remove_path "$_custom_root"
+    # Native diffusion (stable-diffusion.cpp) for a custom/env-mode Studio installs beside
+    # the root at <parent>/stable-diffusion.cpp -- find_sd_cpp_binary resolves it from
+    # UNSLOTH_STUDIO_HOME.parent (sd_cpp_engine.py) -- so removing only the root leaves the
+    # build behind. Derive and remove the sibling, guarding the parent path the same way.
+    _custom_sd_cpp="$(dirname "$_custom_root")/stable-diffusion.cpp"
+    if _is_unsafe_root "$_custom_sd_cpp"; then
+        echo "  refusing to remove unsafe path: $_custom_sd_cpp" >&2
+    else
+        _remove_path "$_custom_sd_cpp"
+    fi
 done
 _remove_path "$HOME/.unsloth/studio"
 # Default-mode shared llama.cpp build + cache are siblings of studio (not removed

studio/backend/routes/training.py

@@ -1388,20 +1388,45 @@ _DIFFUSION_DATASET_IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp", ".bmp"}
 _DIFFUSION_DATASET_TEXT_EXTS = {".txt", ".caption", ".jsonl"}
 
 
+def _resolve_dataset_caption(
+    folder: Path, image_path: Path, meta_captions: dict[str, str]
+) -> Optional[str]:
+    """Resolve an image's caption using the same sidecar > metadata precedence the trainer
+    applies in ``discover_image_caption_pairs``. A per-image .txt/.caption sidecar wins and
+    is stripped, so an empty (tombstone) sidecar shadows metadata and yields "" -- the
+    trainer then skips that image (``if caption:``), so it must not count as captioned."""
+    caption: Optional[str] = None
+    for ext in (".txt", ".caption"):
+        sidecar = image_path.with_suffix(ext)
+        if sidecar.is_file():
+            try:
+                caption = sidecar.read_text(encoding = "utf-8").strip()
+            except OSError:
+                caption = None
+            break
+    if caption is None:
+        try:
+            rel = image_path.relative_to(folder).as_posix()
+        except ValueError:
+            rel = None
+        caption = meta_captions.get(image_path.name) or (
+            meta_captions.get(rel) if rel is not None else None
+        )
+    return caption
+
+
 def _diffusion_dataset_summary(folder: Path) -> DiffusionDatasetSummary:
-    # Count an image as captioned when a metadata/captions.jsonl row or a per-image
-    # sidecar (.txt / .caption) resolves a caption for it -- the same sources the
-    # trainer reads. Counting metadata-only captions here keeps a metadata-captioned
-    # dataset from reporting caption_count=0 and being treated as uncaptioned.
+    # Count an image as captioned only when it resolves to a NON-EMPTY caption via the same
+    # sidecar > metadata precedence the trainer uses -- an empty tombstone sidecar shadows a
+    # metadata row and makes the trainer skip the image, so counting it here would over-report
+    # caption_count and mislabel an effectively-uncaptioned dataset as captioned.
     meta_captions = _load_metadata_captions(folder)
     images = captions = 0
     for f in folder.iterdir():
         if not f.is_file() or f.suffix.lower() not in _DIFFUSION_DATASET_IMAGE_EXTS:
             continue
         images += 1
-        if f.name in meta_captions or any(
-            f.with_suffix(ext).is_file() for ext in (".txt", ".caption")
-        ):
+        if _resolve_dataset_caption(folder, f, meta_captions):
             captions += 1
     return DiffusionDatasetSummary(
         name = folder.name, path = str(folder), image_count = images, caption_count = captions
@@ -1475,6 +1500,9 @@ async def upload_diffusion_dataset(
     named folder under the Studio datasets root, creating it if needed. Repeat uploads
     into the same name accumulate, so large datasets can arrive in batches. The returned
     name can be passed directly as ``data_dir`` to /diffusion/start."""
+    import os
+    import tempfile
+
     from utils.paths import datasets_root
     from utils.upload_limits import get_upload_limit_bytes, get_upload_limit_label
 
@@ -1496,9 +1524,15 @@ async def upload_diffusion_dataset(
                 detail = f"Unsupported file '{f.filename}'. Allowed: {exts}",
             )
         dest = folder / filename
+        # Stream into a sibling temp file and only atomically promote it once the whole file
+        # is written and within the limit. A mid-stream 413 (or any abort) then removes the
+        # TEMP file, never dest, so re-uploading a batch that trips the limit can no longer
+        # truncate/delete an example that was already stored under the same name.
+        fd, tmp_name = tempfile.mkstemp(dir = folder, prefix = ".upload-", suffix = ext)
+        tmp = Path(tmp_name)
         complete = False
         try:
-            with open(dest, "wb") as out:
+            with os.fdopen(fd, "wb") as out:
                 while chunk := await f.read(1024 * 1024):
                     total_bytes += len(chunk)
                     if total_bytes > limit_bytes:
@@ -1511,11 +1545,12 @@ async def upload_diffusion_dataset(
                             ),
                         )
                     out.write(chunk)
+            os.replace(tmp, dest)
             complete = True
         finally:
             if not complete:
                 try:
-                    dest.unlink(missing_ok = True)
+                    tmp.unlink(missing_ok = True)
                 except OSError:
                     pass
         uploaded += 1

studio/backend/tests/test_diffusion_training.py

@@ -698,6 +698,62 @@ def test_diffusion_dataset_upload_accumulates(client, dataset_roots):
     assert r.json()["image_count"] == 3
 
 
+def test_diffusion_dataset_upload_over_cap_keeps_existing_example(
+    client, dataset_roots, monkeypatch
+):
+    # A re-upload that trips the size cap mid-write must not destroy an example already
+    # stored under the same name: the write goes to a sibling temp file and only atomically
+    # replaces the original on success, so the 413 leaves the prior good bytes intact.
+    import utils.upload_limits as ul
+
+    ds_root, _ = dataset_roots
+    folder = ds_root / "my style"
+    folder.mkdir()
+    (folder / "cat.png").write_bytes(b"ORIGINAL-CAT-BYTES")
+
+    monkeypatch.setattr(ul, "get_upload_limit_bytes", lambda: 8)
+    monkeypatch.setattr(ul, "get_upload_limit_label", lambda: "8B")
+
+    r = client.post(
+        "/api/train/diffusion/dataset",
+        data = {"name": "my style"},
+        files = [("files", ("cat.png", b"x" * 64, "image/png"))],
+    )
+    assert r.status_code == 413, r.text
+    # The pre-existing example survives untouched, and no temp file is left behind.
+    assert (folder / "cat.png").read_bytes() == b"ORIGINAL-CAT-BYTES"
+    assert sorted(p.name for p in folder.iterdir()) == ["cat.png"]
+
+
+def test_diffusion_info_empty_sidecar_shadows_metadata_caption(client, dataset_roots):
+    # An empty (tombstone) .txt sidecar shadows a metadata row -- the trainer strips it and
+    # skips the image -- so the summary must not count it as captioned (which would report a
+    # dataset as captioned that the trainer would reject as having no captioned images).
+    ds_root, _ = dataset_roots
+    folder = ds_root / "tombstoned"
+    folder.mkdir()
+    (folder / "a.png").write_bytes(b"x")
+    (folder / "b.png").write_bytes(b"x")
+    (folder / "c.png").write_bytes(b"x")
+    (folder / "metadata.jsonl").write_text(
+        json.dumps({"file_name": "a.png", "text": "cap a"})
+        + "\n"
+        + json.dumps({"file_name": "c.png", "text": "cap c"})
+        + "\n",
+        encoding = "utf-8",
+    )
+    # a.png: metadata caption but an empty sidecar tombstone -> uncaptioned.
+    (folder / "a.txt").write_text("   ", encoding = "utf-8")
+    # b.png: real sidecar caption. c.png: metadata only. Both captioned.
+    (folder / "b.txt").write_text("cap b", encoding = "utf-8")
+
+    r = client.get("/api/train/diffusion/info")
+    assert r.status_code == 200, r.text
+    summary = next(d for d in r.json()["datasets"] if d["name"] == "tombstoned")
+    assert summary["image_count"] == 3
+    assert summary["caption_count"] == 2
+
+
 def test_diffusion_dataset_upload_rejects_traversal_names(client, dataset_roots):
     for bad in ("../evil", "a/b", ".hidden", " "):
         r = client.post(

tests/sh/test_uninstall_sd_cpp_custom_root.sh

@@ -0,0 +1,85 @@
+#!/bin/bash
+# SPDX-License-Identifier: AGPL-3.0-only
+# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
+# Unit test for custom/env-mode stable-diffusion.cpp removal in scripts/uninstall.sh.
+#
+# A custom Studio (UNSLOTH_STUDIO_HOME=<root>) installs its native diffusion build beside
+# the root at <parent>/stable-diffusion.cpp -- find_sd_cpp_binary resolves it from
+# UNSLOTH_STUDIO_HOME.parent (sd_cpp_engine.py). Uninstall must remove that sibling too, or
+# a stale build lingers and a fresh install's finder can pick it up. Tested hermetically:
+# the real custom-root removal loop + its helpers are extracted from uninstall.sh and run
+# against per-test fixtures. Follows the extract-via-sed pattern of test_uninstall_shared_icon.sh.
+set -e
+
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+UNINSTALL_SH="$SCRIPT_DIR/../../scripts/uninstall.sh"
+PASS=0
+FAIL=0
+
+_TMP_ROOT=$(mktemp -d)
+trap 'rm -rf "$_TMP_ROOT"' EXIT
+# Deterministic deny-list checks: keep $HOME clear of the fixture trees.
+HOME="$_TMP_ROOT/home"
+mkdir -p "$HOME"
+
+assert_nodir() { _l="$1"; [ -d "$2" ] && { echo "  FAIL: $_l (still present: $2)"; FAIL=$((FAIL+1)); } || { echo "  PASS: $_l"; PASS=$((PASS+1)); }; }
+assert_dir()   { _l="$1"; [ -d "$2" ] && { echo "  PASS: $_l"; PASS=$((PASS+1)); } || { echo "  FAIL: $_l (missing dir $2)"; FAIL=$((FAIL+1)); }; }
+
+# Extract the helpers the loop depends on, plus the real custom-root removal loop.
+HELPERS_FILE=$(mktemp -p "$_TMP_ROOT")
+{
+    sed -n '/^_remove_path() {/,/^}/p'      "$UNINSTALL_SH"
+    sed -n '/^_is_studio_root() {/,/^}/p'   "$UNINSTALL_SH"
+    sed -n '/^_is_unsafe_root() {/,/^}/p'   "$UNINSTALL_SH"
+} > "$HELPERS_FILE"
+LOOP_FILE=$(mktemp -p "$_TMP_ROOT")
+sed -n '/^_custom_studio_roots | while IFS= read -r _custom_root; do/,/^done/p' "$UNINSTALL_SH" > "$LOOP_FILE"
+
+# shellcheck disable=SC1090
+. "$HELPERS_FILE"
+
+# make_studio <root> : a valid custom Studio root (share/studio.conf owner marker) plus its
+# sibling <parent>/stable-diffusion.cpp build, each with a file so removal is observable.
+make_studio() {
+    mkdir -p "$1/share"
+    : > "$1/share/studio.conf"
+    _sib="$(dirname "$1")/stable-diffusion.cpp"
+    mkdir -p "$_sib"
+    : > "$_sib/sd-cli"
+}
+run_loop() {
+    # shellcheck disable=SC1090
+    . "$LOOP_FILE"
+}
+
+# 1. Single custom root -> root AND its sibling stable-diffusion.cpp both removed.
+p1="$_TMP_ROOT/inst1"
+make_studio "$p1/studioA"
+: > "$p1/keep.txt"  # unrelated sibling file must be untouched
+_custom_studio_roots() { printf '%s\n' "$p1/studioA"; }
+run_loop
+assert_nodir "single custom root removed"                 "$p1/studioA"
+assert_nodir "custom-root sibling stable-diffusion.cpp removed" "$p1/stable-diffusion.cpp"
+[ -f "$p1/keep.txt" ] && { echo "  PASS: unrelated sibling file kept"; PASS=$((PASS+1)); } || { echo "  FAIL: unrelated sibling file removed"; FAIL=$((FAIL+1)); }
+
+# 2. Two custom roots sharing a parent share one sd.cpp -> all removed, no error on the
+#    second (already-gone) removal.
+p2="$_TMP_ROOT/inst2"
+make_studio "$p2/studioB"
+make_studio "$p2/studioC"  # same parent -> same sibling sd.cpp
+_custom_studio_roots() { printf '%s\n%s\n' "$p2/studioB" "$p2/studioC"; }
+run_loop
+assert_nodir "shared-parent root B removed"               "$p2/studioB"
+assert_nodir "shared-parent root C removed"               "$p2/studioC"
+assert_nodir "shared sibling stable-diffusion.cpp removed" "$p2/stable-diffusion.cpp"
+
+# 3. Default-mode sd.cpp (a bare ~/.unsloth/stable-diffusion.cpp with no custom root) is NOT
+#    touched by the custom-root loop -- it is removed by the separate default-mode line.
+mkdir -p "$HOME/.unsloth/stable-diffusion.cpp"
+_custom_studio_roots() { printf '%s\n' "$p1/studioA"; }  # a now-removed root -> guard skips
+run_loop
+assert_dir "default-mode sd.cpp untouched by custom loop" "$HOME/.unsloth/stable-diffusion.cpp"
+
+echo ""
+echo "Results: $PASS passed, $FAIL failed"
+[ "$FAIL" = 0 ]