From 8b8980a607c6e77de7c87b1cba4b88b04bee4dcb Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 05:03:18 +0000 Subject: [PATCH] Fix/adjust diffusion: round 12 local-path GGUF + per-variant delete + MPS + base namespace for PR #5754 Round 12 reviewer findings. Backend correctness (P1) * core/inference/diffusion.py load_model: GGUF branch now handles an absolute local directory passed as repo_id by joining Path(repo_id) / gguf_filename directly instead of handing the path to hf_hub_download (which raises HFValidationError because the path is not 'namespace/repo'). Closes round 12 review #1 -- the load request advertised 'local path' support but actually only worked for Hub repo ids. Delete guard precision (P1) * routes/models.py /delete-finetuned + /delete-cached: diffusion guard now consults gguf_filename from status() and ALLOWS per-variant deletes that target a different quant than the one the loaded pipeline is reading. Loading 'Q4_K_S' no longer blocks deleting 'Q8_0' from the same repo / export directory (round 12 reviews #3 and #4). Accelerator (P2) * core/inference/diffusion.py _drain_cuda_cache: also calls torch.mps.empty_cache() when the MPS backend is the active accelerator. Apple Silicon swaps now actually return held VRAM instead of leaving it pinned in the Metal allocator (round 12 review #10). Smart base repo (P2) * core/inference/diffusion.py _smart_base_repo: only inspects the LAST segment of the repo id / path for the 'base' / '9b' tokens. A namespace like baseorg/FLUX.2-klein-4B-GGUF or a parent directory like /home/me/.cache/base/... no longer falsely selects the Base variant (round 12 review #9). --- studio/backend/core/inference/diffusion.py | 57 ++++++++++++++++++---- studio/backend/routes/models.py | 51 +++++++++++++++++-- 2 files changed, 94 insertions(+), 14 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index e2da94b362..367ae4de56 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -155,12 +155,17 @@ def _smart_base_repo(fam: DiffusionFamily, repo_id: str) -> str: containing "9b" gets the 9B base, "base-4b" / "base-9b" map to the Base variants, everything else falls back to the family default (Apache 2.0 4B Base). + + Only the LAST segment of the repo id / path is inspected so a + namespace or parent directory like ``baseorg/...`` or + ``/home/me/.cache/base/...`` does not falsely select the Base + variant (round 12 review #9). """ if fam.name != "flux.2-klein": return fam.base_repo - lower = (repo_id or "").lower() - is_9b = "9b" in lower - is_base = "base" in lower + last_segment = (repo_id or "").rstrip("/").rsplit("/", 1)[-1].lower() + is_9b = "9b" in last_segment + is_base = "base" in last_segment if is_9b and is_base: return "black-forest-labs/FLUX.2-klein-base-9B" if is_9b: @@ -529,11 +534,28 @@ class DiffusionBackend: f"Family {fam.name} does not have a GGUF transformer " "path wired in this build; load the full repo instead." ) - local_gguf_path = hf_hub_download( - repo_id = repo_id, - filename = gguf_filename, - token = hf_token, - ) + # DiffusionLoadRequest.repo_id is documented to + # accept either a Hub repo id OR a local + # absolute path (Studio export, downloaded HF + # snapshot, etc.). Only the Hub case wants + # hf_hub_download -- a local repo path passed + # to it raises HFValidationError because + # "/abs/path" is not "namespace/repo". + repo_id_path = Path(repo_id).expanduser() + if repo_id_path.is_absolute() and repo_id_path.is_dir(): + candidate = repo_id_path / gguf_filename + if not candidate.is_file(): + raise RuntimeError( + f"Local repo path '{repo_id}' does not contain " + f"'{gguf_filename}'." + ) + local_gguf_path = str(candidate) + else: + local_gguf_path = hf_hub_download( + repo_id = repo_id, + filename = gguf_filename, + token = hf_token, + ) # All cheap failure points (bad gguf_filename, missing # pipeline / transformer class, gated download token, @@ -1007,12 +1029,16 @@ def _release(obj: Any) -> None: def _drain_cuda_cache() -> None: - """Hand freed weights back to the CUDA allocator. + """Hand freed weights back to the active accelerator's allocator. Call this AFTER every reference to the freed object has been dropped (caller's local + attribute) and a ``gc.collect()`` has fired __del__. Calling earlier would empty an already-pinned - cache and not actually release the memory.""" + cache and not actually release the memory. + + Handles CUDA *and* MPS (Apple Silicon) so a diffusion swap on + macOS actually returns VRAM to the Metal allocator. + """ try: import torch @@ -1020,6 +1046,17 @@ def _drain_cuda_cache() -> None: torch.cuda.empty_cache() except Exception: pass + try: + import torch + + mps_backend = getattr(getattr(torch, "backends", None), "mps", None) + if mps_backend is not None and mps_backend.is_available(): + mps_module = getattr(torch, "mps", None) + empty_cache = getattr(mps_module, "empty_cache", None) if mps_module else None + if empty_cache is not None: + empty_cache() + except Exception: + pass # ─── Module-level singleton ─────────────────────────────────────────── diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 7a1b238065..82bfa8865b 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -2002,6 +2002,19 @@ async def delete_finetuned_model( if v: candidates.append(v) target_str = str(target_path) + # Per-variant deletes only touch ``_delete_gguf_variant_ + # files(target_path, gguf_variant)`` which removes a + # specific quant file. If the loaded pipeline uses a + # DIFFERENT variant from the same directory, the delete + # is safe. Round 12 review #3. + loaded_gguf = ( + diff_status.get("gguf_filename") or "" + ).lower() + wants_variant = ( + export_type == "gguf" + and gguf_variant + and loaded_gguf + ) for candidate in candidates: try: candidate_path = Path(candidate).expanduser() @@ -2022,6 +2035,15 @@ async def delete_finetuned_model( or _is_path_under(candidate_resolved, target_path) or _is_path_under(target_path, candidate_resolved) ): + # Allow per-variant deletes that target a + # different quant than the loaded one. + if wants_variant: + variant_low = gguf_variant.lower() + loaded_label = ( + _extract_quant_label(loaded_gguf) or "" + ).lower() + if loaded_label and loaded_label != variant_low: + continue raise HTTPException( status_code = 400, detail = "Unload the diffusion image model before deleting", @@ -2769,10 +2791,31 @@ async def delete_cached_model( } owned.discard("") if needle in owned: - raise HTTPException( - status_code = 400, - detail = "Unload the diffusion image model before deleting", - ) + # Per-variant delete only touches the requested + # quant via ``_delete_gguf_variant_files``. If the + # loaded pipeline uses a DIFFERENT variant from the + # same repo, the delete is safe. Round 12 review #4. + loaded_gguf = ( + diff_status.get("gguf_filename") or "" + ).lower() + if variant and loaded_gguf: + variant_low = variant.lower() + loaded_label = ( + _extract_quant_label(loaded_gguf) or "" + ).lower() + if loaded_label and loaded_label != variant_low: + # Different quant from the same repo -> allow. + pass + else: + raise HTTPException( + status_code = 400, + detail = "Unload the diffusion image model before deleting", + ) + else: + raise HTTPException( + status_code = 400, + detail = "Unload the diffusion image model before deleting", + ) except HTTPException: raise except Exception as e: