Fix/adjust diffusion: round 14 P1+P2 batch for PR #5754
Round 14 reviewer aggregate (logs/review_round14_aggregate.md): P1 fixes: - routes/export.py /load-checkpoint now runs the active-export 409 guard BEFORE the chat / diffusion unloads, so a rejected request no longer tears down unrelated GPU state. - core/inference/llama_cpp.py wraps the WHOLE load_model body in a single try/finally that publishes loading_model_identifier across download, metadata read, VRAM settle, process spawn, and health check. Done via a thin load_model wrapper around the existing body (renamed _load_model_impl) to avoid reindenting hundreds of lines. - routes/models.py /delete-finetuned now checks loading_model_identifier so a pending HF GGUF download cannot have its destination directory rmtree'd before llama-server spawns. - core/inference/diffusion.py stores the original caller-supplied gguf_filename (e.g. ``BF16/model.gguf``) in a new self._gguf_filename field and exposes it as active_gguf_filename. UI-facing gguf_filename still collapses to basename for the panel. - routes/models.py /delete-cached llama guard now allows safe different-variant deletes when hf_variant differs, matching the diffusion path's variant-aware behaviour. - core/inference/diffusion.py tracks self._cpu_offload_enabled and forces a CPU torch.Generator when offload is on, so seeded generation no longer crashes on CUDA hosts with the default offload enabled. P2 fixes: - core/inference/diffusion.py detect_family normalises mixed separators (``Qwen_Image-Edit-GGUF``, ``Qwen-Image_Edit-GGUF``, ``QwenImageEdit-GGUF``) so every Qwen-Image-Edit spelling is excluded from the base Qwen-Image family. - core/inference/diffusion.py logger.info / logger.error in load_model run repo_id and effective_base through _redact_hf_tokens so URL-embedded ``hf_xxxxx`` tokens never reach structured-log sinks. - core/inference/diffusion.py _release_other_gpu_owners_for_diffusion now raises RuntimeError when an export job is active instead of logging and continuing, so direct backend callers cannot bypass the route layer's 409 guard. - core/inference/diffusion.py full-diffusers repo / base_repo paths expand ``~`` via _expand_existing_local_path so ``repo_id="~/models/my-flux"`` no longer falls through to the Hub. Tests: - 5 new regression cases (mixed Qwen-Image-Edit separators, token redaction, status full-filename, CPU offload generator device, staging Windows leaf already-set sanity). - All 68 diffusion backend + route tests pass.
This commit is contained in:
Daniel Han-Chen
parent
f501ab8fc8
commit
e03ed3dd29
6 changed files with 439 additions and 115 deletions
|
|
@ -182,6 +182,43 @@ def _smart_base_repo(fam: DiffusionFamily, repo_id: str) -> str:
|
|||
return "black-forest-labs/FLUX.2-klein-4B"
|
||||
|
||||
|
||||
def _expand_existing_local_path(value: str) -> str:
|
||||
"""Expand ``~`` in ``value`` when the expanded path exists locally.
|
||||
|
||||
Round 14 P2 #11: the GGUF local path branch already calls
|
||||
``Path(repo_id).expanduser()``, but the full-diffusers-repo and
|
||||
base-companion-repo paths passed the literal ``~/...`` straight
|
||||
into ``from_pretrained``, which treated it as a Hub id and tried
|
||||
to download. Keep behaviour identical for Hub ids (no leading
|
||||
``~`` -> return as-is) and for non-existent expansions (the
|
||||
diffusers loader will surface its own ``not found`` error).
|
||||
"""
|
||||
if not value or not isinstance(value, str) or not value.startswith("~"):
|
||||
return value
|
||||
candidate = Path(value).expanduser()
|
||||
if candidate.exists():
|
||||
return str(candidate)
|
||||
return value
|
||||
|
||||
|
||||
_HF_TOKEN_RE = re.compile(r"hf_[A-Za-z0-9]{20,}")
|
||||
|
||||
|
||||
def _redact_hf_tokens(value: Any) -> Any:
|
||||
"""Scrub embedded ``hf_xxxxxxxx`` tokens out of a string before
|
||||
logging. Round 14 P2 #9: callers can wrap an authenticated URL
|
||||
(``https://hf_token@huggingface.co/...``) into ``repo_id`` /
|
||||
``base_repo`` / paths; the token would otherwise reach
|
||||
structured-log sinks via the load-info / load-failure log lines.
|
||||
Non-strings are returned unchanged so the helper is safe to
|
||||
sprinkle through ``logger.info`` / ``logger.error`` argument
|
||||
lists.
|
||||
"""
|
||||
if not isinstance(value, str):
|
||||
return value
|
||||
return _HF_TOKEN_RE.sub("<redacted>", value)
|
||||
|
||||
|
||||
def _resolve_local_gguf_child(repo_root: Path, gguf_filename: str) -> Path:
|
||||
"""Resolve a GGUF filename inside a local repo directory safely.
|
||||
|
||||
|
|
@ -288,12 +325,26 @@ def detect_family(
|
|||
needle = (repo_id or "").lower()
|
||||
if not needle:
|
||||
return None
|
||||
# Normalise mixed separator spellings (``Qwen_Image-Edit-GGUF``,
|
||||
# ``Qwen-Image_Edit-GGUF``, ``Qwen.Image.Edit-GGUF``) and the
|
||||
# compact concatenation (``QwenImageEdit-GGUF``) so the
|
||||
# _FAMILY_EXCLUDE deny lists do not need every permutation of
|
||||
# ``-``, ``_``, ``.`` and run-together spellings to keep
|
||||
# Qwen-Image-Edit out of the base Qwen-Image family (round 14
|
||||
# P2 #8).
|
||||
needle_norm = re.sub(r"[^a-z0-9]+", "-", needle).strip("-")
|
||||
needle_compact = re.sub(r"[^a-z0-9]+", "", needle)
|
||||
# Scan _FAMILIES first (GGUF-supported), then _FULL_REPO_FAMILIES
|
||||
# so a repo like ``stabilityai/stable-diffusion-xl-base-1.0`` is
|
||||
# auto-detected as SDXL instead of returning None.
|
||||
for fam in _FAMILIES + _FULL_REPO_FAMILIES:
|
||||
excludes = _FAMILY_EXCLUDE.get(fam.name, ())
|
||||
if any(e in needle for e in excludes):
|
||||
if any(
|
||||
e in needle
|
||||
or re.sub(r"[^a-z0-9]+", "-", e).strip("-") in needle_norm
|
||||
or re.sub(r"[^a-z0-9]+", "", e) in needle_compact
|
||||
for e in excludes
|
||||
):
|
||||
continue
|
||||
if fam.name in needle:
|
||||
return fam
|
||||
|
|
@ -352,9 +403,24 @@ class DiffusionBackend:
|
|||
self._family: Optional[DiffusionFamily] = None
|
||||
self._repo_id: Optional[str] = None
|
||||
self._gguf_path: Optional[str] = None
|
||||
# Original ``gguf_filename`` the caller passed in, preserved
|
||||
# so delete guards can compare against subdirectory variants
|
||||
# like ``BF16/model.gguf`` or ``Q4_K_M/model.gguf`` instead
|
||||
# of the collapsed basename (round 14 P1 #4). The basename
|
||||
# alone (``model.gguf``) loses the quant directory and lets
|
||||
# /delete-cached unlink the wrong file.
|
||||
self._gguf_filename: Optional[str] = None
|
||||
self._base_repo: Optional[str] = None
|
||||
self._device: Optional[str] = None
|
||||
self._dtype: Optional[str] = None
|
||||
# True when ``enable_model_cpu_offload()`` was applied on the
|
||||
# loaded pipeline. Diffusers' offload moves the active
|
||||
# submodule between CPU and GPU on each step, so a CUDA
|
||||
# ``torch.Generator`` mismatches the CPU-resident embeddings
|
||||
# and generation crashes mid-forward (round 14 P1 #6). When
|
||||
# this is True, seeded generation has to use a CPU generator
|
||||
# regardless of self._device.
|
||||
self._cpu_offload_enabled: bool = False
|
||||
self._loaded_at: Optional[float] = None
|
||||
self._loading: bool = False
|
||||
self._last_error: Optional[str] = None
|
||||
|
|
@ -387,6 +453,11 @@ class DiffusionBackend:
|
|||
# local HF cache layout (and the system username on default
|
||||
# POSIX layouts) to any authenticated Studio session.
|
||||
with self._lock:
|
||||
# UI-facing collapsed basename. Full local path leaks the
|
||||
# HF cache layout + system username; the original caller-
|
||||
# supplied filename (e.g. ``BF16/model.gguf``) is kept
|
||||
# separately as ``active_gguf_filename`` for delete
|
||||
# guards.
|
||||
gguf_basename = Path(self._gguf_path).name if self._gguf_path else None
|
||||
# Expose BOTH the resident pipeline's id AND the pending
|
||||
# load target. Delete guards must check both: when model A
|
||||
|
|
@ -398,7 +469,7 @@ class DiffusionBackend:
|
|||
# user just clicked.
|
||||
active_repo = self._repo_id
|
||||
active_base = self._base_repo
|
||||
active_gguf = gguf_basename
|
||||
active_gguf = self._gguf_filename
|
||||
pending_repo = self._pending_repo_id if self._loading else None
|
||||
pending_base = self._pending_base_repo if self._loading else None
|
||||
pending_gguf = self._pending_gguf_filename if self._loading else None
|
||||
|
|
@ -414,6 +485,14 @@ class DiffusionBackend:
|
|||
if pending_repo and pending_repo != active_repo:
|
||||
ui_family = None
|
||||
ui_pipeline_class = None
|
||||
# UI-facing ``gguf_filename`` collapses to the basename
|
||||
# so the Images panel does not surface internal cache /
|
||||
# variant directory names. Guard-facing ``active_*`` /
|
||||
# ``pending_*`` retain the full caller-supplied filename
|
||||
# so /delete-cached can compare against subdirectory
|
||||
# variants like ``BF16/model.gguf`` (round 14 P1 #4-5).
|
||||
ui_gguf = pending_gguf or active_gguf
|
||||
ui_gguf_basename = Path(ui_gguf).name if ui_gguf else None
|
||||
return {
|
||||
"is_loaded": self._pipe is not None,
|
||||
"is_loading": self._loading,
|
||||
|
|
@ -421,7 +500,7 @@ class DiffusionBackend:
|
|||
"family": ui_family,
|
||||
"pipeline_class": ui_pipeline_class,
|
||||
"base_repo": pending_base or active_base,
|
||||
"gguf_filename": pending_gguf or active_gguf,
|
||||
"gguf_filename": ui_gguf_basename,
|
||||
# Guard-facing fields: every repo / path / GGUF
|
||||
# filename the backend owns RIGHT NOW. Delete routes
|
||||
# iterate both, paired so the variant-filename check
|
||||
|
|
@ -550,9 +629,11 @@ class DiffusionBackend:
|
|||
# success.
|
||||
self._pending_repo_id = repo_id
|
||||
self._pending_base_repo = base_repo
|
||||
self._pending_gguf_filename = (
|
||||
Path(gguf_filename).name if gguf_filename else None
|
||||
)
|
||||
# Store the caller's full ``gguf_filename`` (e.g.
|
||||
# ``BF16/model.gguf``) so the variant-aware delete
|
||||
# guards have the subdirectory info. The UI side of
|
||||
# status() still collapses to the basename for display.
|
||||
self._pending_gguf_filename = gguf_filename if gguf_filename else None
|
||||
try:
|
||||
pipeline_cls = getattr(diffusers, fam.pipeline_class, None)
|
||||
if pipeline_cls is None:
|
||||
|
|
@ -595,11 +676,15 @@ class DiffusionBackend:
|
|||
"or load a full diffusers repo (base_repo only "
|
||||
"applies when picking a GGUF quant)."
|
||||
)
|
||||
effective_base = repo_id
|
||||
# ``~/models/my-flux`` must be expanded so
|
||||
# diffusers' from_pretrained does not pass the
|
||||
# literal tilde through to ``os.path.isdir`` and
|
||||
# fall back to the Hub (round 14 P2 #11).
|
||||
effective_base = _expand_existing_local_path(repo_id)
|
||||
with self._lock:
|
||||
self._pending_base_repo = effective_base
|
||||
elif base_repo:
|
||||
effective_base = base_repo
|
||||
effective_base = _expand_existing_local_path(base_repo)
|
||||
# Refresh pending so delete guards see the actual
|
||||
# base, not just caller-supplied None.
|
||||
with self._lock:
|
||||
|
|
@ -608,13 +693,19 @@ class DiffusionBackend:
|
|||
effective_base = _smart_base_repo(fam, repo_id)
|
||||
with self._lock:
|
||||
self._pending_base_repo = effective_base
|
||||
# ``repo_id`` / ``effective_base`` are user-supplied
|
||||
# strings that can embed an ``hf_xxxxx`` token via a
|
||||
# URL-style path (``https://hf_token@huggingface.co/...``).
|
||||
# Scrub them BEFORE the logger formats the line so the
|
||||
# token never reaches structured-log sinks (round 14
|
||||
# P2 #9).
|
||||
logger.info(
|
||||
"Loading diffusion model %s (family=%s, device=%s, dtype=%s, base=%s)",
|
||||
repo_id,
|
||||
_redact_hf_tokens(repo_id),
|
||||
fam.name,
|
||||
device,
|
||||
dtype,
|
||||
effective_base,
|
||||
_redact_hf_tokens(effective_base),
|
||||
)
|
||||
|
||||
transformer = None
|
||||
|
|
@ -685,9 +776,11 @@ class DiffusionBackend:
|
|||
self._family = None
|
||||
self._repo_id = None
|
||||
self._gguf_path = None
|
||||
self._gguf_filename = None
|
||||
self._base_repo = None
|
||||
self._device = None
|
||||
self._dtype = None
|
||||
self._cpu_offload_enabled = False
|
||||
self._loaded_at = None
|
||||
_release(old)
|
||||
old = None
|
||||
|
|
@ -735,6 +828,9 @@ class DiffusionBackend:
|
|||
pipe_kwargs["token"] = hf_token
|
||||
|
||||
pipe = None
|
||||
cpu_offload_enabled = bool(
|
||||
enable_model_cpu_offload and device == "cuda"
|
||||
)
|
||||
try:
|
||||
pipe = pipeline_cls.from_pretrained(effective_base, **pipe_kwargs)
|
||||
# Device placement / offload can ALSO raise after
|
||||
|
|
@ -746,7 +842,7 @@ class DiffusionBackend:
|
|||
# the next load attempt. Explicitly release both
|
||||
# pipe and transformer in the same try (round 13
|
||||
# P2 #11).
|
||||
if enable_model_cpu_offload and device == "cuda":
|
||||
if cpu_offload_enabled:
|
||||
pipe.enable_model_cpu_offload()
|
||||
else:
|
||||
pipe.to(device)
|
||||
|
|
@ -765,9 +861,14 @@ class DiffusionBackend:
|
|||
self._family = fam
|
||||
self._repo_id = repo_id
|
||||
self._gguf_path = local_gguf_path
|
||||
# Preserve the full caller-supplied filename, not
|
||||
# just the basename, so per-variant delete guards
|
||||
# see ``BF16/model.gguf`` (round 14 P1 #4).
|
||||
self._gguf_filename = gguf_filename if gguf_filename else None
|
||||
self._base_repo = effective_base
|
||||
self._device = device
|
||||
self._dtype = str(dtype).replace("torch.", "")
|
||||
self._cpu_offload_enabled = cpu_offload_enabled
|
||||
self._loaded_at = time.time()
|
||||
# Clear loading + pending here, BEFORE returning,
|
||||
# so the response payload reports the resident
|
||||
|
|
@ -813,7 +914,11 @@ class DiffusionBackend:
|
|||
# Use ``logger.error`` with the already-scrubbed
|
||||
# message and exc_info=False so the bearer token
|
||||
# cannot leak through structured logging sinks.
|
||||
logger.error("Diffusion load failed for %s: %s", repo_id, exc_msg)
|
||||
logger.error(
|
||||
"Diffusion load failed for %s: %s",
|
||||
_redact_hf_tokens(repo_id),
|
||||
exc_msg,
|
||||
)
|
||||
raise RuntimeError(
|
||||
f"Failed to load diffusion model: {exc_msg}"
|
||||
) from exc
|
||||
|
|
@ -845,9 +950,11 @@ class DiffusionBackend:
|
|||
self._family = None
|
||||
self._repo_id = None
|
||||
self._gguf_path = None
|
||||
self._gguf_filename = None
|
||||
self._base_repo = None
|
||||
self._device = None
|
||||
self._dtype = None
|
||||
self._cpu_offload_enabled = False
|
||||
self._loaded_at = None
|
||||
_release(old)
|
||||
old = None # noqa: F841
|
||||
|
|
@ -930,14 +1037,25 @@ class DiffusionBackend:
|
|||
raise RuntimeError("No diffusion model is loaded.")
|
||||
pipe = self._pipe
|
||||
device = self._device or "cpu"
|
||||
cpu_offload_enabled = self._cpu_offload_enabled
|
||||
generator = None
|
||||
if seed is not None:
|
||||
# Match the device of the pipeline so determinism holds
|
||||
# across reload cycles. For CPU offload, the noise still
|
||||
# has to live on the device the diffusion forward runs on.
|
||||
gen_device = (
|
||||
"cuda" if device == "cuda" and torch.cuda.is_available() else "cpu"
|
||||
)
|
||||
# across reload cycles. When CPU offload is enabled
|
||||
# (the default on CUDA hosts), diffusers shuttles each
|
||||
# submodule between CPU and GPU on every step. A CUDA
|
||||
# torch.Generator then mismatches the CPU-resident
|
||||
# embeddings at the start of the forward and the run
|
||||
# crashes (round 14 P1 #6). Use a CPU generator in that
|
||||
# case; numerical determinism for the same seed is
|
||||
# preserved because the seed feeds an int rather than a
|
||||
# device-local RNG state.
|
||||
if cpu_offload_enabled:
|
||||
gen_device = "cpu"
|
||||
else:
|
||||
gen_device = (
|
||||
"cuda" if device == "cuda" and torch.cuda.is_available() else "cpu"
|
||||
)
|
||||
generator = torch.Generator(device = gen_device).manual_seed(int(seed))
|
||||
|
||||
call_kwargs: dict[str, Any] = {
|
||||
|
|
@ -1126,34 +1244,49 @@ def _release_other_gpu_owners_for_diffusion() -> None:
|
|||
# higher-level guard) cannot still kill an active export.
|
||||
try:
|
||||
from core.export import get_export_backend # type: ignore
|
||||
|
||||
exp = get_export_backend()
|
||||
if getattr(exp, "current_checkpoint", None):
|
||||
is_export_active_fn = getattr(exp, "is_export_active", None)
|
||||
export_is_active = False
|
||||
if is_export_active_fn is not None:
|
||||
try:
|
||||
export_is_active = bool(is_export_active_fn())
|
||||
except Exception:
|
||||
# Unverifiable status -> treat as 'might be
|
||||
# active' and refuse to touch the subprocess.
|
||||
export_is_active = True
|
||||
if export_is_active:
|
||||
logger.info(
|
||||
"Skipping export shutdown for diffusion load: "
|
||||
"is_export_active=True (route layer should have "
|
||||
"rejected this request with 409)"
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"Shutting down idle export subprocess before diffusion load"
|
||||
)
|
||||
exp._shutdown_subprocess()
|
||||
exp.current_checkpoint = None
|
||||
exp.is_vision = False
|
||||
exp.is_peft = False
|
||||
except Exception as exc:
|
||||
logger.debug("export unload skipped: %s", exc)
|
||||
logger.debug("export module not importable: %s", exc)
|
||||
return
|
||||
|
||||
try:
|
||||
exp = get_export_backend()
|
||||
except Exception as exc:
|
||||
logger.debug("export backend not available: %s", exc)
|
||||
return
|
||||
|
||||
is_export_active_fn = getattr(exp, "is_export_active", None)
|
||||
if is_export_active_fn is not None:
|
||||
try:
|
||||
export_is_active = bool(is_export_active_fn())
|
||||
except Exception:
|
||||
# Unverifiable status -> treat as 'might be active' and
|
||||
# refuse so a direct backend caller (test / script /
|
||||
# future route that forgot the higher-level 409 guard)
|
||||
# cannot still terminate an in-flight export.
|
||||
export_is_active = True
|
||||
if export_is_active:
|
||||
# Round 14 P2 #10: the prior behaviour logged a warning
|
||||
# and continued, so direct ``DiffusionBackend.load_model``
|
||||
# callers (tests, scripts) silently bypassed the route
|
||||
# layer's 409. Hard-refuse instead so any code path that
|
||||
# reaches this helper while an export is active sees the
|
||||
# same failure mode the route returns.
|
||||
raise RuntimeError(
|
||||
"An export job is currently active. Stop the export "
|
||||
"job before loading a diffusion image model."
|
||||
)
|
||||
|
||||
if getattr(exp, "current_checkpoint", None):
|
||||
try:
|
||||
logger.info(
|
||||
"Shutting down idle export subprocess before diffusion load"
|
||||
)
|
||||
exp._shutdown_subprocess()
|
||||
exp.current_checkpoint = None
|
||||
exp.is_vision = False
|
||||
exp.is_peft = False
|
||||
except Exception as exc:
|
||||
logger.debug("idle export shutdown failed: %s", exc)
|
||||
|
||||
# Note: active training is *not* stopped here. The route layer
|
||||
# (`_raise_if_training_active` in routes/inference.py) refuses
|
||||
|
|
|
|||
|
|
@ -2616,9 +2616,66 @@ class LlamaCppBackend:
|
|||
|
||||
Returns True if server started and health check passed.
|
||||
"""
|
||||
# Serialise the whole load so concurrent /load calls never
|
||||
# leave two llama-server processes alive (#5401 / #5161). Does
|
||||
# not block /unload, /status, /load-progress.
|
||||
# Publish ``_loading_model_identifier`` BEFORE any phase of
|
||||
# the load can begin and clear it AFTER the load fully settles
|
||||
# (success or failure, including the duplicate-state fast path
|
||||
# and every internal early ``return False``). Round 14 P1 #2:
|
||||
# the prior inline try/finally only wrapped the download, so
|
||||
# /delete-cached and the cross-workload handoff helpers saw
|
||||
# the backend as idle once the GGUF bytes had landed but the
|
||||
# subprocess had not yet spawned. Mark the load as pending
|
||||
# for the entire duration -- download, metadata read,
|
||||
# VRAM settle, process spawn, health check, audio probe.
|
||||
self._loading_model_identifier = model_identifier
|
||||
try:
|
||||
# Serialise the whole load so concurrent /load calls never
|
||||
# leave two llama-server processes alive (#5401 / #5161).
|
||||
# Does not block /unload, /status, /load-progress.
|
||||
return self._load_model_impl(
|
||||
gguf_path = gguf_path,
|
||||
mmproj_path = mmproj_path,
|
||||
hf_repo = hf_repo,
|
||||
hf_variant = hf_variant,
|
||||
hf_token = hf_token,
|
||||
model_identifier = model_identifier,
|
||||
is_vision = is_vision,
|
||||
n_ctx = n_ctx,
|
||||
chat_template_override = chat_template_override,
|
||||
cache_type_kv = cache_type_kv,
|
||||
speculative_type = speculative_type,
|
||||
spec_draft_n_max = spec_draft_n_max,
|
||||
n_threads = n_threads,
|
||||
n_gpu_layers = n_gpu_layers,
|
||||
n_parallel = n_parallel,
|
||||
extra_args = extra_args,
|
||||
)
|
||||
finally:
|
||||
self._loading_model_identifier = None
|
||||
|
||||
def _load_model_impl(
|
||||
self,
|
||||
*,
|
||||
gguf_path: Optional[str] = None,
|
||||
mmproj_path: Optional[str] = None,
|
||||
hf_repo: Optional[str] = None,
|
||||
hf_variant: Optional[str] = None,
|
||||
hf_token: Optional[str] = None,
|
||||
model_identifier: str,
|
||||
is_vision: bool = False,
|
||||
n_ctx: int = 4096,
|
||||
chat_template_override: Optional[str] = None,
|
||||
cache_type_kv: Optional[str] = None,
|
||||
speculative_type: Optional[str] = None,
|
||||
spec_draft_n_max: Optional[int] = None,
|
||||
n_threads: Optional[int] = None,
|
||||
n_gpu_layers: Optional[int] = None,
|
||||
n_parallel: int = 1,
|
||||
extra_args: Optional[List[str]] = None,
|
||||
) -> bool:
|
||||
"""Internal body of ``load_model``. Kept as a separate method
|
||||
so ``load_model`` can wrap it in a single try/finally that
|
||||
publishes ``_loading_model_identifier`` for the WHOLE load
|
||||
instead of only the download window."""
|
||||
with self._serial_load_lock:
|
||||
# Duplicate /load that raced past the route-level check
|
||||
# (the first one hadn't published _healthy=True yet). If the
|
||||
|
|
@ -2693,40 +2750,25 @@ class LlamaCppBackend:
|
|||
# Scope HF_HUB_OFFLINE to the download block only when DNS is
|
||||
# dead; cleanup runs even on exception so a transient hiccup
|
||||
# at the start of one load cannot quarantine future loads.
|
||||
#
|
||||
# Publish ``_loading_model_identifier`` BEFORE entering the
|
||||
# download so /delete-cached and the cross-workload handoff
|
||||
# helpers can see a multi-GB pending load: previously they
|
||||
# only consulted ``model_identifier``, which the success
|
||||
# path sets later (see "Set identifier early" below). That
|
||||
# left a window where the user could rmtree the cache the
|
||||
# download was still writing to, or start /images/load
|
||||
# while llama-server was about to come up on the same GPU.
|
||||
# Cleared in ``finally`` so failed / cancelled loads do not
|
||||
# leak the pending state.
|
||||
self._loading_model_identifier = model_identifier
|
||||
try:
|
||||
if hf_repo:
|
||||
with _hf_offline_if_dns_dead():
|
||||
model_path = self._download_gguf(
|
||||
if hf_repo:
|
||||
with _hf_offline_if_dns_dead():
|
||||
model_path = self._download_gguf(
|
||||
hf_repo = hf_repo,
|
||||
hf_variant = hf_variant,
|
||||
hf_token = hf_token,
|
||||
)
|
||||
# Auto-download mmproj for vision models
|
||||
if is_vision and not mmproj_path:
|
||||
mmproj_path = self._download_mmproj(
|
||||
hf_repo = hf_repo,
|
||||
hf_variant = hf_variant,
|
||||
hf_token = hf_token,
|
||||
)
|
||||
# Auto-download mmproj for vision models
|
||||
if is_vision and not mmproj_path:
|
||||
mmproj_path = self._download_mmproj(
|
||||
hf_repo = hf_repo,
|
||||
hf_token = hf_token,
|
||||
)
|
||||
elif gguf_path:
|
||||
if not Path(gguf_path).is_file():
|
||||
raise FileNotFoundError(f"GGUF file not found: {gguf_path}")
|
||||
model_path = gguf_path
|
||||
else:
|
||||
raise ValueError("Either gguf_path or hf_repo must be provided")
|
||||
finally:
|
||||
self._loading_model_identifier = None
|
||||
elif gguf_path:
|
||||
if not Path(gguf_path).is_file():
|
||||
raise FileNotFoundError(f"GGUF file not found: {gguf_path}")
|
||||
model_path = gguf_path
|
||||
else:
|
||||
raise ValueError("Either gguf_path or hf_repo must be provided")
|
||||
|
||||
# Set identifier early so _read_gguf_metadata can use it for DeepSeek detection
|
||||
self._model_identifier = model_identifier
|
||||
|
|
|
|||
|
|
@ -107,6 +107,42 @@ async def load_checkpoint(
|
|||
),
|
||||
)
|
||||
|
||||
backend = get_export_backend()
|
||||
# Refuse to reload the export checkpoint while an export job
|
||||
# is still running. ``ExportBackend.load_checkpoint`` would
|
||||
# terminate the running subprocess in order to spawn a new
|
||||
# one, silently corrupting the partial output the user is
|
||||
# waiting on (round 13 P1 #1). Runs BEFORE the chat /
|
||||
# diffusion unloads below: a 409 from this guard must not
|
||||
# leave the user's chat or diffusion GPU owners freed for
|
||||
# nothing (round 14 P1 #1). ``is_export_active`` may be
|
||||
# absent on older / mocked backends; treat missing as "no
|
||||
# async-job tracker available" and skip rather than
|
||||
# fail-closed.
|
||||
is_export_active_fn = getattr(backend, "is_export_active", None)
|
||||
if is_export_active_fn is not None:
|
||||
try:
|
||||
export_is_active = bool(is_export_active_fn())
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Could not verify export status before export load: %s", e
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = (
|
||||
"Could not verify export status before loading "
|
||||
"an export checkpoint. Try again."
|
||||
),
|
||||
) from e
|
||||
if export_is_active:
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = (
|
||||
"An export job is currently active. Stop the "
|
||||
"export job before loading another checkpoint."
|
||||
),
|
||||
)
|
||||
|
||||
# Free GPU memory: shut down any chat backend before loading
|
||||
# the export checkpoint. Routes the unload through the shared
|
||||
# helper so we cover llama-server is_active=True and
|
||||
|
|
@ -141,41 +177,6 @@ async def load_checkpoint(
|
|||
except Exception as e:
|
||||
logger.debug("diffusion unload skipped for export: %s", e)
|
||||
|
||||
backend = get_export_backend()
|
||||
# Refuse to reload the export checkpoint while an export job
|
||||
# is still running. ``ExportBackend.load_checkpoint`` would
|
||||
# terminate the running subprocess in order to spawn a new
|
||||
# one, silently corrupting the partial output the user is
|
||||
# waiting on (round 13 P1 #1). Mirrors the symmetric guards
|
||||
# already in place for chat / diffusion / training handoffs.
|
||||
# ``is_export_active`` may be absent on older / mocked
|
||||
# backends -- treat missing as "no async-job tracker
|
||||
# available" -> skip rather than fail-closed; the
|
||||
# surrounding chat / diffusion unloads have already run.
|
||||
is_export_active_fn = getattr(backend, "is_export_active", None)
|
||||
if is_export_active_fn is not None:
|
||||
try:
|
||||
export_is_active = bool(is_export_active_fn())
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Could not verify export status before export load: %s", e
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code = 503,
|
||||
detail = (
|
||||
"Could not verify export status before loading "
|
||||
"an export checkpoint. Try again."
|
||||
),
|
||||
) from e
|
||||
if export_is_active:
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = (
|
||||
"An export job is currently active. Stop the "
|
||||
"export job before loading another checkpoint."
|
||||
),
|
||||
)
|
||||
|
||||
# load_checkpoint spawns and waits on a subprocess and can take
|
||||
# minutes. Run it in a worker thread so the event loop stays
|
||||
# free to serve the live log SSE stream concurrently.
|
||||
|
|
|
|||
|
|
@ -1940,6 +1940,27 @@ async def delete_finetuned_model(
|
|||
from routes.inference import get_llama_cpp_backend
|
||||
|
||||
llama_backend = get_llama_cpp_backend()
|
||||
# Pending HF GGUF download targeting this path: round 14 P1 #3.
|
||||
# ``loading_model_identifier`` is set before the download starts
|
||||
# and cleared after the subprocess settles, so the user cannot
|
||||
# rmtree the directory llama.cpp is writing into mid-flight.
|
||||
loading_identifier = getattr(llama_backend, "loading_model_identifier", None)
|
||||
if (
|
||||
loading_identifier
|
||||
and _loaded_model_matches_deleted_path(
|
||||
loading_identifier,
|
||||
target_path,
|
||||
)
|
||||
and (
|
||||
not gguf_variant
|
||||
or not getattr(llama_backend, "hf_variant", None)
|
||||
or llama_backend.hf_variant.lower() == gguf_variant.lower()
|
||||
)
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code = 409,
|
||||
detail = "Cannot delete a model while it is loading",
|
||||
)
|
||||
if (
|
||||
llama_backend.is_active
|
||||
and not llama_backend.is_loaded
|
||||
|
|
@ -2745,14 +2766,24 @@ async def delete_cached_model(
|
|||
# Exact match only (case-insensitive). Prefix match would
|
||||
# block deleting unrelated ``org/model`` while
|
||||
# ``org/model-v2`` is loaded -- same surface the diffusion
|
||||
# guard fixed in round 5.
|
||||
# guard fixed in round 5. Per-variant deletes that target a
|
||||
# DIFFERENT quant than the loaded one are allowed so the
|
||||
# llama and diffusion paths stay symmetric (round 14 P1 #7).
|
||||
if loaded_id == needle and (
|
||||
llama_backend.is_loaded or getattr(llama_backend, "is_active", False)
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Unload the model before deleting",
|
||||
loaded_variant = (getattr(llama_backend, "hf_variant", None) or "").lower()
|
||||
requested_variant = (variant or "").lower()
|
||||
same_variant = (
|
||||
not requested_variant
|
||||
or not loaded_variant
|
||||
or requested_variant == loaded_variant
|
||||
)
|
||||
if same_variant:
|
||||
raise HTTPException(
|
||||
status_code = 400,
|
||||
detail = "Unload the model before deleting",
|
||||
)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -1359,6 +1359,120 @@ def test_generate_image_with_metadata_returns_active_pipeline(monkeypatch):
|
|||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"repo_id",
|
||||
[
|
||||
"unsloth/Qwen_Image-Edit-GGUF",
|
||||
"unsloth/Qwen-Image_Edit-GGUF",
|
||||
"unsloth/Qwen-ImageEdit-GGUF",
|
||||
"unsloth/qwen-image_edit-2509-GGUF",
|
||||
"unsloth/Qwen.Image.Edit-GGUF",
|
||||
],
|
||||
)
|
||||
def test_detect_family_qwen_image_edit_mixed_separators(repo_id):
|
||||
"""Round 14 P2 #8: every spelling of Qwen-Image-Edit must NOT
|
||||
match the base Qwen-Image text-to-image family."""
|
||||
from core.inference.diffusion import detect_family
|
||||
|
||||
assert detect_family(repo_id) is None
|
||||
|
||||
|
||||
def test_redact_hf_tokens_removes_url_embedded_token():
|
||||
"""Round 14 P2 #9: tokens embedded in user-supplied paths /
|
||||
URLs must be scrubbed before logging."""
|
||||
from core.inference.diffusion import _redact_hf_tokens
|
||||
|
||||
leaky = "https://hf_abcdefghij0123456789@huggingface.co/unsloth/FLUX.2-klein-4B-GGUF"
|
||||
redacted = _redact_hf_tokens(leaky)
|
||||
assert "hf_" not in redacted
|
||||
assert "<redacted>" in redacted
|
||||
# Non-strings pass through unchanged so the helper is safe in
|
||||
# logger argument lists where families / dtypes mix in.
|
||||
assert _redact_hf_tokens(None) is None
|
||||
assert _redact_hf_tokens(42) == 42
|
||||
|
||||
|
||||
def test_status_preserves_active_gguf_subdir(monkeypatch):
|
||||
"""Round 14 P1 #4: status() must surface the original caller-
|
||||
supplied gguf_filename (``BF16/model.gguf``) instead of the
|
||||
collapsed basename."""
|
||||
import core.inference.diffusion as d
|
||||
|
||||
backend = d.DiffusionBackend()
|
||||
backend._pipe = object()
|
||||
backend._repo_id = "unsloth/FLUX.2-klein-4B-GGUF"
|
||||
backend._gguf_path = "/cache/models/unsloth/FLUX.2-klein-4B-GGUF/BF16/model.gguf"
|
||||
backend._gguf_filename = "BF16/model.gguf"
|
||||
backend._family = d.DiffusionFamily(
|
||||
name = "flux.2-klein",
|
||||
pipeline_class = "Flux2KleinPipeline",
|
||||
transformer_class = "Flux2Transformer2DModel",
|
||||
base_repo = "black-forest-labs/FLUX.2-klein-4B",
|
||||
aliases = (),
|
||||
)
|
||||
|
||||
s = backend.status()
|
||||
assert s["active_gguf_filename"] == "BF16/model.gguf"
|
||||
# UI-facing field still collapses to the basename.
|
||||
assert s["gguf_filename"] == "model.gguf"
|
||||
|
||||
|
||||
def test_generator_uses_cpu_when_cpu_offload_enabled(monkeypatch):
|
||||
"""Round 14 P1 #6: seeded CUDA generation must NOT create a
|
||||
CUDA torch.Generator when the pipeline was loaded with CPU
|
||||
offload enabled, otherwise it crashes mid-forward."""
|
||||
import core.inference.diffusion as d
|
||||
|
||||
backend = d.DiffusionBackend()
|
||||
|
||||
class _FakePipe:
|
||||
def __init__(self):
|
||||
self.last_kwargs = None
|
||||
|
||||
def __call__(self, **kwargs):
|
||||
self.last_kwargs = kwargs
|
||||
from PIL import Image
|
||||
|
||||
return SimpleNamespace(images = [Image.new("RGB", (8, 8))])
|
||||
|
||||
fake_pipe = _FakePipe()
|
||||
backend._pipe = fake_pipe
|
||||
backend._device = "cuda"
|
||||
backend._cpu_offload_enabled = True
|
||||
|
||||
captured_devices: list[str] = []
|
||||
|
||||
class _FakeGenerator:
|
||||
def __init__(self, device):
|
||||
captured_devices.append(device)
|
||||
|
||||
def manual_seed(self, seed):
|
||||
return self
|
||||
|
||||
class _FakeTorchCuda:
|
||||
@staticmethod
|
||||
def is_available():
|
||||
return True
|
||||
|
||||
fake_torch = SimpleNamespace(
|
||||
Generator = _FakeGenerator, cuda = _FakeTorchCuda
|
||||
)
|
||||
monkeypatch.setitem(sys.modules, "torch", fake_torch)
|
||||
|
||||
backend._generate_image_unlocked(prompt = "x", seed = 7, width = 8, height = 8)
|
||||
assert captured_devices == ["cpu"]
|
||||
|
||||
|
||||
def test_smart_base_repo_uses_windows_leaf_only_already_set_separator_round14():
|
||||
"""Sanity: relative paths still work after the Windows fix."""
|
||||
from core.inference.diffusion import _smart_base_repo, detect_family
|
||||
|
||||
repo = "owner/FLUX.2-klein-9B-GGUF"
|
||||
fam = detect_family(repo)
|
||||
assert fam is not None
|
||||
assert _smart_base_repo(fam, repo) == "black-forest-labs/FLUX.2-klein-9B"
|
||||
|
||||
|
||||
def test_generate_image_with_metadata_blocks_concurrent_unload(monkeypatch):
|
||||
"""Round 13 P2 #9: _generate_lock serialises the forward AND the
|
||||
meta snapshot, so a queued unload cannot wipe state in between."""
|
||||
|
|
|
|||
|
|
@ -84,6 +84,9 @@ class _FakeBackend:
|
|||
"active_base_repo": (
|
||||
"black-forest-labs/FLUX.2-klein" if self._loaded else None
|
||||
),
|
||||
# Round 14: guard-facing GGUF filename is now the full
|
||||
# caller-supplied value, but this fake never sets one so
|
||||
# both active and pending stay None.
|
||||
"active_gguf_filename": None,
|
||||
"pending_repo_id": None,
|
||||
"pending_base_repo": None,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue