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.
This commit is contained in:
Daniel Han 2026-07-07 02:16:56 +00:00
commit fb94a79337
5 changed files with 206 additions and 10 deletions

View file

@ -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 }

View file

@ -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