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).
This commit is contained in:
Daniel Han-Chen
parent
921c60232e
commit
8b8980a607
2 changed files with 94 additions and 14 deletions
|
|
@ -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 ───────────────────────────────────────────
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue