Fix image-generation load, training, and inference bugs

scripts/nvfp4_t211_probe.py

@@ -24,7 +24,7 @@
 
 BASE = "Tongyi-MAI/Z-Image-Turbo"
 PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
-OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/nvfp4_t211_images")
+OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "nvfp4_t211_images"
 
 
 # ----------------------------------------------------------------------------- diag

scripts/sparse_accum_probe.py

@@ -27,7 +27,7 @@
 
 BASE = "Tongyi-MAI/Z-Image-Turbo"
 PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
-OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/sparse_images")
+OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "sparse_images"
 
 
 def _psnr(a, b):

scripts/uninstall.ps1

@@ -315,7 +315,7 @@ function Uninstall-UnslothStudio {
         # 2. A loaded module under a target root (orphaned mp-fork python holding a
         #    venv DLL). Scoped to names that load our DLLs to keep the scan fast.
         try {
-            $cands = Get-Process -Name python, pythonw, unsloth, llama-server, llama-cli -ErrorAction SilentlyContinue
+            $cands = Get-Process -Name python, pythonw, unsloth, llama-server, llama-cli, sd-cli, sd-server -ErrorAction SilentlyContinue
             foreach ($proc in $cands) {
                 $hit = $false
                 try {
@@ -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,8 @@ function Uninstall-UnslothStudio {
             continue
         }
         _RemovePath $r
+        # The native diffusion sibling (<custom root>.parent\stable-diffusion.cpp) is left
+        # in place: sd.cpp writes no owner marker, so auto-deleting it could destroy a clone.
     }
     # Default install dir (always at %USERPROFILE%\.unsloth\studio when present).
     if ($defaultStudioHome) { _RemovePath $defaultStudioHome }

scripts/uninstall.sh

@@ -217,8 +217,9 @@ _remove_path "$HOME/.unsloth/studio"
 # when absent. A user-set UNSLOTH_LLAMA_CPP_PATH is intentionally kept.
 _remove_path "$HOME/.unsloth/llama.cpp"
 # Default-mode native diffusion (stable-diffusion.cpp / sd-cli) build, a sibling of
-# studio like llama.cpp (install_sd_cpp_prebuilt.default_install_dir()). No-op in
-# env/custom mode and when absent. A user-set UNSLOTH_SD_CPP_PATH is kept.
+# studio like llama.cpp. Only the default location is removed; the env/custom-mode
+# sibling (<custom root>.parent/stable-diffusion.cpp) is left in place since sd.cpp
+# writes no owner marker and auto-deleting it could destroy a user-managed clone.
 _remove_path "$HOME/.unsloth/stable-diffusion.cpp"
 _remove_path "$HOME/.unsloth/.cache"
 # Isolated Node.js runtime (install_node_prebuilt.py), a sibling of studio in

studio/backend/core/inference/diffusion.py

@@ -406,13 +406,11 @@ def _dense_quant_prefetch_needed(self, fam: DiffusionFamily, kwargs: dict) -> bo
         """True when ``load_pipeline`` may take the dense transformer-quant path, so
         the prefetch should also pull the base repo's ``transformer/`` shards.
 
-        Those shards are excluded from the prefetch by default (the GGUF supplies
-        the transformer), but ``_load_dense_quant_pipeline`` fetches them with
-        ``from_pretrained(subfolder = "transformer")`` under the load lock during
-        "finalizing", after the previous pipeline was already evicted, where
-        unload/cancellation cannot preempt the download. Mirrors the dense-path
-        gates in ``load_pipeline``: quant requested and supported for this device,
-        and no pre-quantized checkpoint that would shortcut the dense build."""
+        Those shards are excluded from the prefetch by default (the GGUF supplies the
+        transformer), but ``_load_dense_quant_pipeline`` fetches them later under the
+        load lock, where unload/cancellation cannot preempt the download. Checks only
+        the dense-path gates knowable pre-download; skips the ``offload_policy`` gate
+        since that needs the GGUF's on-disk size, not known until after this runs."""
         mode = normalize_transformer_quant(kwargs.get("transformer_quant"))
         if mode is None:
             return False
@@ -822,10 +820,14 @@ def base_filter(rfilename: str) -> bool:
         return total, base_files
 
     @staticmethod
-    def _cache_bytes(repo_id: str) -> int:
+    def _hub_cache_repo_dir(repo_id: str) -> Path:
+        """Local HF hub cache dir for ``repo_id``."""
         from huggingface_hub import constants
+        return Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}"
 
-        blobs = Path(constants.HF_HUB_CACHE) / f"models--{repo_id.replace('/', '--')}" / "blobs"
+    @staticmethod
+    def _cache_bytes(repo_id: str) -> int:
+        blobs = DiffusionBackend._hub_cache_repo_dir(repo_id) / "blobs"
         total = 0
         try:
             for entry in blobs.iterdir():
@@ -865,15 +867,25 @@ def _local_dir_weight_bytes(path: Path, *, exclude_transformer: bool) -> int:
     def _companion_cache_bytes(base: str) -> int:
         """Resident companion (VAE + text-encoder) size for the memory plan.
 
-        For a hub base repo this is the cached blob total (``_cache_bytes``). For a
-        LOCAL diffusers base directory the blob cache is empty, so sum the on-disk
-        component weights instead, excluding ``transformer/`` (the GGUF supplies the
-        transformer). Without this a local base folds its multi-GB VAE / text-encoder
-        weights to zero and auto planning can pick a resident placement that OOMs."""
+        Excludes ``transformer/`` (supplied by the GGUF/single file, not resident here) --
+        otherwise the dense-quant prefetch's cached transformer shards would inflate this
+        and wrongly force offload. Walks the snapshot dir, not the flat ``blobs/`` cache,
+        since only the snapshot preserves the subfolder split needed to exclude it."""
         local = Path(base).expanduser()
         if local.is_dir():
             return DiffusionBackend._local_dir_weight_bytes(local, exclude_transformer = True)
-        return DiffusionBackend._cache_bytes(base)
+        snapshots = DiffusionBackend._hub_cache_repo_dir(base) / "snapshots"
+        if not snapshots.is_dir():
+            return 0
+        # Multiple revisions may be cached; the active one is the fullest, so take the max.
+        return max(
+            (
+                DiffusionBackend._local_dir_weight_bytes(rev, exclude_transformer = True)
+                for rev in snapshots.iterdir()
+                if rev.is_dir()
+            ),
+            default = 0,
+        )
 
     # ── Synchronous load / generate / unload ───────────────────────────────
 
@@ -1035,6 +1047,19 @@ def load_pipeline(
                             clear_gpu_cache()
                         except Exception:  # noqa: BLE001
                             pass
+                elif (
+                    kind == "gguf"
+                    and normalize_transformer_quant(transformer_quant) is not None
+                    and dense_transformer_supported(target)
+                    and plan.offload_policy != OFFLOAD_NONE
+                ):
+                    # memory_mode forcing offload silently drops the requested quant; warn so it's diagnosable.
+                    logger.warning(
+                        "diffusion.transformer_quant: %s requested but memory_mode forces "
+                        "offload (%s); loading GGUF without dense quant",
+                        normalize_transformer_quant(transformer_quant),
+                        plan.offload_policy,
+                    )
 
                 if pipe is None:
                     if kind == "pipeline":
@@ -1427,7 +1452,8 @@ def _plan_memory(
 
         The size estimate is per-kind: diffusers keeps GGUF weights packed (per-matmul
         transient dequant), so a GGUF loads near its on-disk size; a safetensors
-        single-file loads near its on-disk size (it carries its dtype); and a full
+        single-file loads near its on-disk size (it carries its dtype), except an fp8
+        transformer file that gets upcast to bf16 on load (~2x resident); and a full
         pipeline is one cached download (transformer + companions), already compressed."""
         device_memory = snapshot_device_memory(target)
         if kind == "pipeline":
@@ -1445,9 +1471,15 @@ def _plan_memory(
             companion_mib = None
         else:
             if kind == "single_file":
-                # Safetensors single-file: no dequant expansion (it carries its dtype).
+                # An fp8 transformer checkpoint loads via from_single_file with a bf16
+                # compute dtype and no quantization_config, so diffusers upcasts it to
+                # bf16 (~2x resident); detect it from the basename. Excludes the
+                # single-file-is-pipeline (SDXL) case, which is already a bf16 pipeline.
+                fp8_upcast = not getattr(fam, "single_file_is_pipeline", False) and (
+                    "fp8" in Path(single_file_path).name.lower() if single_file_path else False
+                )
                 transformer_resident = estimate_safetensors_dense_mib(
-                    file_size_mib(single_file_path)
+                    file_size_mib(single_file_path), fp8_upcast = fp8_upcast
                 )
             else:
                 transformer_resident = estimate_gguf_resident_mib(file_size_mib(single_file_path))
@@ -1537,6 +1569,13 @@ def _controlnet_pipe(self, state: _LoadState, resolved_cn: Any, cancel: threadin
         if cn_model is None:
             if cancel.is_set():
                 raise RuntimeError(DIFFUSION_CANCELLED_MSG)
+            # Keep at most one ControlNet resident: evict the previous module + its
+            # from_pipe wrapper before loading the new one, or swapping ControlNets
+            # within a base-model load accumulates until OOM.
+            if self._cn_models or self._cn_pipes:
+                self._cn_models.clear()
+                self._cn_pipes.clear()
+                clear_gpu_cache()
             # resolve_controlnet accepts a bare owner/name repo without the non-GGUF base
             # trust gate, and from_pretrained below downloads and deserializes it. A
             # malicious pickle .bin would execute on load, so run the same Hub malware
@@ -1724,7 +1763,11 @@ def _reset_step_cache(pipe: Any) -> None:
         resolution/batch changed, or a stale-cache reuse otherwise. Best-effort: a
         transformer without the hook (uncached load) is a silent no-op."""
         transformer = getattr(pipe, "transformer", None)
-        reset = getattr(transformer, "reset_stateful_hooks", None)
+        # ``_reset_stateful_cache`` is the transformer-level entry point; the public
+        # ``reset_stateful_hooks`` lives only on the HookRegistry, kept as a fallback.
+        reset = getattr(transformer, "_reset_stateful_cache", None) or getattr(
+            transformer, "reset_stateful_hooks", None
+        )
         if callable(reset):
             try:
                 reset()
@@ -1832,6 +1875,12 @@ def generate(
                         raise ValueError(
                             f"{state.family.name} is an image-editing model: provide an input image."
                         )
+                    if mask_image is not None:
+                        # The edit family has no inpaint pipeline; a mask would be silently dropped.
+                        raise ValueError(
+                            f"{state.family.name} is an image-editing model and does not "
+                            "support masks (mask_image)."
+                        )
                     workflow = "edit"
                     init_pil = _decode_b64_image(init_image, mode = "RGB")
                 elif mask_image is not None and init_image is not None:

studio/backend/core/inference/diffusion_controlnet.py

@@ -204,8 +204,8 @@ def resolve_controlnet(spec_id: str, *, family: Optional[str] = None) -> Resolve
 
 # Union ControlNet mode indices. A single "union" model covers several control modes and
 # selects the active one via an integer ``control_mode`` argument; these are the standard
-# indices used by the FLUX.1 / Qwen-Image union ControlNets. "passthrough" (an already-made
-# map) carries no intrinsic mode, so it maps to nothing (the caller omits control_mode).
+# indices used by the FLUX.1 / Qwen-Image union ControlNets. "passthrough" carries no
+# intrinsic mode, so union_control_mode() defaults it to 0.
 _UNION_CONTROL_MODES: dict[str, int] = {
     "canny": 0,
     "tile": 1,
@@ -220,13 +220,13 @@ def resolve_controlnet(spec_id: str, *, family: Optional[str] = None) -> Resolve
 def union_control_mode(spec_id: str, control_type: str) -> Optional[int]:
     """The integer ``control_mode`` for a union ControlNet, or None.
 
-    Returns a mode only for a curated *union* catalog entry AND a control type that maps to a
-    known index; otherwise None so the caller omits the kwarg (a non-union ControlNet has a
-    single fixed mode, and 'passthrough' does not name one). Pure lookup, no network."""
+    A union model requires a concrete mode (diffusers raises on None), so a curated union
+    entry always gets an index, defaulting to 0 for types like 'passthrough' that carry
+    none. A non-union entry returns None so the caller omits the kwarg."""
     entry = _catalog_by_id().get(spec_id)
     if entry is None or not entry.is_union:
         return None
-    return _UNION_CONTROL_MODES.get((control_type or "").strip().lower())
+    return _UNION_CONTROL_MODES.get((control_type or "").strip().lower(), 0)
 
 
 def preprocess_control(image: Any, control_type: str) -> Any:

studio/backend/core/inference/diffusion_device.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import os
 from dataclasses import dataclass
 from typing import Any, Optional
 
@@ -234,6 +235,11 @@ def _mps_or_cpu_target(torch: Any) -> DiffusionDeviceTarget:
         mps_available = False
 
     if mps_available:
+        # torch reads PYTORCH_MPS_HIGH_WATERMARK_RATIO once, at the first MPS allocation
+        # (the bfloat16 probe below), so it must be relaxed before that. Otherwise the
+        # allocator caps the process at ~1.7x recommendedMaxWorkingSetSize and can OOM a
+        # model that would otherwise fit in unified RAM. setdefault respects an override.
+        os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.0")
         # Prefer bfloat16; otherwise fall back to float32, NEVER silent float16.
         # Modern diffusion transformers (Z-Image, FLUX.2, ...) produce activations
         # far outside float16's finite range (~6.5e4) -- Z-Image's MLP
@@ -254,7 +260,9 @@ def _mps_or_cpu_target(torch: Any) -> DiffusionDeviceTarget:
 
 
 def _cpu_target(torch: Any, dtype: Any = None) -> DiffusionDeviceTarget:
-    if dtype is None:
+    # torch is None on the no-torch CPU fallback; leave dtype=None then (matching the
+    # no-torch DiffusionDeviceTarget elsewhere) rather than crashing on torch.float32.
+    if dtype is None and torch is not None:
         dtype = torch.float32
     return DiffusionDeviceTarget(
         device = "cpu",

studio/backend/core/inference/diffusion_lora.py

@@ -266,6 +266,13 @@ def _pick_repo_weight_file(repo_id: str, hf_token: Optional[str]) -> str:
     raise FileNotFoundError(f"no .safetensors/.gguf LoRA file found in '{repo_id}'")
 
 
+def _scrub_hub_url(msg: str) -> str:
+    """Strip embedded http(s) URLs from a Hub error message before it hits a 400 body."""
+    cleaned = re.sub(r"https?://\S+", "", msg)
+    # Collapse the whitespace / stray separators the URL removal leaves behind.
+    return re.sub(r"\s{2,}", " ", cleaned).strip()
+
+
 def resolve_specs(
     specs: list[tuple[str, float]],
     *,
@@ -274,20 +281,31 @@ def resolve_specs(
 ) -> list[ResolvedLora]:
     """Resolve request (id, weight) pairs, dropping zero-weight entries.
 
-    A stale / unknown id raises FileNotFoundError inside resolve_one; convert it to
-    ValueError so the route (which maps only ValueError to a 400) reports bad client
-    input instead of a generic 500. A Hub download can also raise
-    ``RuntimeError("Cancelled")`` when the user unloads / starts a superseding load
-    mid-download; convert that to the diffusion cancellation sentinel so the route
-    maps it to a 409 instead of a generic server error toast."""
+    Maps the named not-found/gated Hub errors (bad repo/revision/file/gating) to a 400
+    and scrubs the URL from the message; deliberately does NOT catch the base
+    HfHubHTTPError so a Hub 5xx stays a 500. A mid-download cancel also maps to a 409
+    instead of a generic 500."""
+    from huggingface_hub.errors import (
+        EntryNotFoundError,
+        GatedRepoError,
+        RepositoryNotFoundError,
+        RevisionNotFoundError,
+    )
+
     out: list[ResolvedLora] = []
     try:
         for spec_id, weight in specs:
             if weight == 0:
                 continue
             out.append(resolve_one(spec_id, weight, hf_token = hf_token, cancel_event = cancel_event))
-    except FileNotFoundError as exc:
-        raise ValueError(str(exc)) from exc
+    except (
+        FileNotFoundError,
+        RepositoryNotFoundError,
+        RevisionNotFoundError,
+        EntryNotFoundError,
+        GatedRepoError,
+    ) as exc:
+        raise ValueError(_scrub_hub_url(str(exc))) from exc
     except RuntimeError as exc:
         if str(exc) == "Cancelled":
             raise RuntimeError(DIFFUSION_CANCELLED_MSG) from exc

studio/backend/core/inference/diffusion_memory.py

@@ -236,14 +236,19 @@ def estimate_gguf_resident_mib(storage_mib: Optional[int]) -> Optional[int]:
     return int(storage_mib * 1.05)  # small margin for allocator + bf16 norms/biases
 
 
-def estimate_safetensors_dense_mib(storage_mib: Optional[int]) -> Optional[int]:
+def estimate_safetensors_dense_mib(
+    storage_mib: Optional[int], *, fp8_upcast: bool = False
+) -> Optional[int]:
     """Resident size of a safetensors checkpoint, in MiB.
 
-    Unlike a GGUF (which is dequantised to bf16/fp16 on load, so a 4-bit file
-    expands ~4x), a safetensors checkpoint loads near its on-disk size: a dense
-    bf16 file is already bf16, and a bnb-4bit / fp8 file stays compressed in VRAM.
-    So the on-disk size is the estimate, returned unchanged (None passes through).
-    """
+    Unlike a GGUF (dequantised on load, so a 4-bit file expands ~4x), a safetensors
+    checkpoint usually loads near its on-disk size (None passes through unchanged).
+    Exception: ``fp8_upcast`` -- an fp8 single-file transformer loads with no
+    quantization_config, so diffusers upcasts it to bf16 (~2x on-disk resident)."""
+    if storage_mib is None:
+        return None
+    if fp8_upcast:
+        return storage_mib * 2
     return storage_mib
 
 

studio/backend/core/inference/sd_cpp_backend.py

@@ -752,6 +752,10 @@ def generate(
                 "img2img / inpaint / reference / upscale are not yet supported on the native "
                 "sd.cpp engine; run on a GPU (diffusers) for image-conditioned workflows."
             )
+        # strength 0 (or None) disables ControlNet (documented on the request model, matches
+        # the diffusers path), so it must be a no-op here too, not a hard 400.
+        if controlnet is not None and controlnet[3] in (None, 0, 0.0):
+            controlnet = None
         if controlnet is not None:
             raise ValueError(
                 "ControlNet is not yet supported on the native sd.cpp engine; run on a GPU "