From 2ef9b0e09fc0f3b73f8c7acd721f9809348b7737 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 07:31:51 +0000 Subject: [PATCH] Fix/adjust diffusion: round 16 P1+P2 batch for PR #5754 Round 16 reviewer aggregate (logs/review_round16_aggregate.md): P1 fixes: - routes/models.py /delete-cached llama guard pairs loading_id with loading_hf_variant so deleting a different cached quant (Q8_0) while another variant (Q4_K_M) is loading is no longer blocked. - core/inference/diffusion.py load_model now calls _release_other_gpu_owners_for_diffusion BEFORE _release_chat_backend_for_diffusion. The other-owners helper RAISES on active training/export, so a route -> worker race or direct backend caller no longer drops the user's chat model before the diffusion load is refused. - routes/models.py /delete-cached diffusion guard fails CLOSED (503) on HF cache scan failure instead of silently falling through to repo-id-only matching, which could miss a loaded local snapshot path. - routes/inference.py _release_llama_for and _release_safetensors_chat_for now raise 503 on actual unload failure (exception or False return), so new GPU workloads do not start while the old chat process still owns VRAM. - core/inference/diffusion.py status() now takes include_internal=False by default and only exposes the guard-facing active_*/pending_* paths when callers opt in. The public /api/inference/images/status route gets the redacted payload; routes/models.py delete guards pass include_internal=True so they still see the raw paths. - core/inference/diffusion.py generate_image_with_metadata routes the response model through _display_repo_id so /images/generate cannot echo back an absolute local path. P2 fixes: - routes/inference.py /images/load now maps backend "Could not verify training/export status" to 503 instead of 409, matching the route-level pre-check. - core/inference/diffusion.py _release_other_gpu_owners_for_diffusion raises "Could not verify export status" when the is_export_active() probe itself raises, instead of silently treating it as active export. - core/inference/diffusion.py detect_family compares compact family spellings (Flux2Klein) against per-token compact strings so unsloth/Flux2Klein-GGUF matches the flux.2-klein family without matching the embedded substring inside flux.20. - main.py installs a RequestValidationError handler that scrubs hf_xxxxx tokens out of the 422 response body so a rejected ``repo_id`` containing a URL-embedded HF token does not echo it back to the browser. Tests: - 3 new regression cases (Flux2Klein compact alias, public status redaction, generate_image_with_metadata redaction). - All 75 diffusion backend + route tests pass. --- studio/backend/core/inference/diffusion.py | 123 +++++++++------ studio/backend/main.py | 37 +++++ studio/backend/routes/inference.py | 142 ++++++++++++------ studio/backend/routes/models.py | 46 +++++- .../backend/tests/test_diffusion_backend.py | 104 +++++++++++-- 5 files changed, 343 insertions(+), 109 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 195ded40a6..ca3c3d917b 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -357,34 +357,37 @@ def detect_family( # P2 #8). needle_norm = re.sub(r"[^a-z0-9]+", "-", needle).strip("-") needle_compact = re.sub(r"[^a-z0-9]+", "", needle) + # Per-token compact strings let ``unsloth/Flux2Klein-GGUF`` match + # the ``flux2klein`` alias: the whole-needle compact is + # ``unslothflux2kleingguf`` and the regex boundary check rejects + # the embedded match, but the token ``Flux2Klein`` (between the + # ``/`` and the ``-``) compacts to exactly ``flux2klein`` (round + # 16 P2 #9). + needle_compact_tokens = { + re.sub(r"[^a-z0-9]+", "", token) + for token in re.split(r"[^a-z0-9]+", needle) + if token + } def _matches_family_token(term: str) -> bool: """Token-boundary match on the normalised needle. Prevents ``owner/flux.20-model`` from matching ``flux.2`` because ``flux.20`` does not have a separator after ``flux-2`` - (round 15 P2 #8). Falls back to compact equality so aliases - like ``qwenimage`` still match ``unsloth/QwenImage-GGUF``.""" + (round 15 P2 #8). Compact spellings (``flux2klein``) match + only when they appear as a complete repo-name token, not + as a substring of a longer token (round 16 P2 #9).""" term_norm = re.sub(r"[^a-z0-9]+", "-", term.lower()).strip("-") if not term_norm: return False if re.search(rf"(^|-){re.escape(term_norm)}($|-)", needle_norm): return True term_compact = re.sub(r"[^a-z0-9]+", "", term.lower()) - if term_compact and term_compact in needle_compact: - # Compact contiguous match: ``qwenimage`` in - # ``qwenimage-gguf`` -> qwenimage-compact in needle_compact. - # Use word boundary on the compact form too: the compact - # ``flux2`` must not match inside ``flux20``. - return ( - bool( - re.search( - rf"(^|[^0-9a-z]){re.escape(term_compact)}([^0-9a-z]|$)", - needle_compact, - ) - ) - or term_compact == needle_compact - ) - return False + if not term_compact: + return False + return ( + term_compact in needle_compact_tokens + or term_compact == needle_compact + ) # Scan _FAMILIES first (GGUF-supported), then _FULL_REPO_FAMILIES # so a repo like ``stabilityai/stable-diffusion-xl-base-1.0`` is @@ -496,7 +499,7 @@ class DiffusionBackend: def repo_id(self) -> Optional[str]: return self._repo_id - def status(self) -> dict[str, Any]: + def status(self, *, include_internal: bool = False) -> dict[str, Any]: # Take _lock so the snapshot cannot observe a torn state where # _pipe was already swapped but _family/_repo_id haven't been # updated yet (or vice versa). Frontend polling at 1 Hz would @@ -504,6 +507,14 @@ class DiffusionBackend: # Only echo the GGUF basename; full absolute path leaks the # local HF cache layout (and the system username on default # POSIX layouts) to any authenticated Studio session. + # + # Round 16 P1 #5: the guard-facing ``active_*`` / ``pending_*`` + # fields hold the EXACT raw path (so /delete-cached can match + # an HF snapshot mmap) but are NOT safe to surface to the + # browser. Callers that need the raw path (route-internal + # delete guards) pass ``include_internal=True``; the public + # ``/api/inference/images/status`` route always uses the + # public payload. with self._lock: # UI-facing collapsed basename. Full local path leaks the # HF cache layout + system username; the original caller- @@ -552,7 +563,7 @@ class DiffusionBackend: # guard-facing ``active_*`` / ``pending_*`` fields below # preserve the exact value so delete guards still match # against the snapshot path. - return { + payload: dict[str, Any] = { "is_loaded": self._pipe is not None, "is_loading": self._loading, "repo_id": _display_repo_id(pending_repo or active_repo), @@ -560,23 +571,30 @@ class DiffusionBackend: "pipeline_class": ui_pipeline_class, "base_repo": _display_repo_id(pending_base or active_base), "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 - # is compared against the SAME repo that owns it - # (round 13 P1 #3-5). - "active_repo_id": active_repo, - "active_base_repo": active_base, - "active_gguf_filename": active_gguf, - "pending_repo_id": pending_repo, - "pending_base_repo": pending_base, - "pending_gguf_filename": pending_gguf, "device": self._device, "dtype": self._dtype, "loaded_at": self._loaded_at, "last_error": self._last_error, "supported_families": supported_families(), } + if include_internal: + # Guard-facing fields: every repo / path / GGUF + # filename the backend owns RIGHT NOW. Delete routes + # iterate both, paired so the variant-filename check + # is compared against the SAME repo that owns it + # (round 13 P1 #3-5). Round 16 P1 #5: never returned + # by the public /images/status route. + payload.update( + { + "active_repo_id": active_repo, + "active_base_repo": active_base, + "active_gguf_filename": active_gguf, + "pending_repo_id": pending_repo, + "pending_base_repo": pending_base, + "pending_gguf_filename": pending_gguf, + } + ) + return payload def _pick_device_and_dtype(self) -> tuple[str, "Any"]: """Pick (device, dtype) for the current host. @@ -808,20 +826,29 @@ class DiffusionBackend: # transient Hub error on the GGUF download) have now # been validated. Anything past this line allocates # GPU memory, so: - # 1. Release competing GPU owners (chat + export). - # 2. Release any *previous* diffusion pipeline so the + # 1. Verify training is idle and the export job (if + # any) is also idle. ``_release_other_gpu_owners + # _for_diffusion`` RAISES on conflict, so it must + # run BEFORE we unload chat (round 16 P1 #2): a + # route precheck -> worker race could otherwise + # drop the user's chat model only to bail out + # because training started in between, and a + # direct ``DiffusionBackend.load_model`` caller + # that did not run the route prechecks would also + # leave chat unloaded for nothing. + # 2. Release the chat backend (llama-server + the + # safetensors orchestrator) now that we know the + # load can actually proceed. + # 3. Release any *previous* diffusion pipeline so the # new transformer / new from_pretrained does not # race the old pipe for VRAM. Switching between # FLUX.2 klein 4B and 9B on a 16-24 GB GPU OOMs # otherwise: from_single_file allocates the new # transformer while the old pipeline still owns # its weights. - # 3. THEN call from_single_file / from_pretrained. - # Training is *not* unloaded here: the route layer - # refuses /images/load with HTTP 409 when training is - # active so the user keeps their long run. - _release_chat_backend_for_diffusion() + # 4. THEN call from_single_file / from_pretrained. _release_other_gpu_owners_for_diffusion() + _release_chat_backend_for_diffusion() old = self._pipe if old is not None: @@ -1172,8 +1199,12 @@ class DiffusionBackend: with self._generate_lock: image = self._generate_image_unlocked(**kwargs) with self._lock: + # Round 16 P1 #6: route ``model`` through + # _display_repo_id so a generation response for a + # locally-loaded model cannot echo back an absolute + # filesystem path to the browser. meta = { - "model": self._repo_id, + "model": _display_repo_id(self._repo_id), "family": self._family.name if self._family else None, } return image, meta @@ -1343,12 +1374,16 @@ def _release_other_gpu_owners_for_diffusion() -> 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 + except Exception as exc: + # Round 16 P2 #8: distinguish unverifiable status from + # active export. The previous "treat as active" mapping + # surfaced as a misleading 409 conflict; raise a + # "Could not verify" RuntimeError so the route layer + # maps it to 503 (retryable) instead. + raise RuntimeError( + "Could not verify export status before loading a " + "diffusion image model." + ) from exc if export_is_active: # Round 14 P2 #10: the prior behaviour logged a warning # and continued, so direct ``DiffusionBackend.load_model`` diff --git a/studio/backend/main.py b/studio/backend/main.py index 004ae404cd..ad727e0998 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -293,6 +293,43 @@ app = FastAPI( lifespan = lifespan, ) + +# ── Validation error scrubber ──────────────────────────────────── +# Round 16 P2 #10: FastAPI's default RequestValidationError handler +# echoes the rejected ``input`` value back in the 422 body. A +# request like +# {"repo_id": "https://hf_token@huggingface.co/owner/repo"} +# is rejected by ``DiffusionLoadRequest._no_embedded_hf_tokens``, +# but the rejected URL would still appear in the response payload, +# leaking the token to the browser console / network log. Wrap the +# handler so any ``hf_xxxxx`` substring is replaced with +# ```` before serialisation. Scoped to the response body +# only; the underlying validator behaviour is unchanged. +from fastapi.exceptions import RequestValidationError as _RequestValidationError # noqa: E402 +from fastapi.responses import JSONResponse as _JSONResponse # noqa: E402 +import re as _re_validation # noqa: E402 + + +_HF_TOKEN_VALIDATION_RE = _re_validation.compile(r"hf_[A-Za-z0-9]{20,}") + + +def _scrub_validation_obj(value): + if isinstance(value, str): + return _HF_TOKEN_VALIDATION_RE.sub("", value) + if isinstance(value, list): + return [_scrub_validation_obj(v) for v in value] + if isinstance(value, dict): + return {k: _scrub_validation_obj(v) for k, v in value.items()} + return value + + +@app.exception_handler(_RequestValidationError) +async def _validation_error_scrubbing_handler(request, exc): + return _JSONResponse( + status_code = 422, + content = {"detail": _scrub_validation_obj(exc.errors())}, + ) + # Initialize structured logging from loggers.config import LogConfig from loggers.handlers import LoggingMiddleware diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index 41fc87c424..8e021d7813 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -366,60 +366,103 @@ async def _release_llama_for(workload: str) -> None: /export/load could start while a long ``_download_gguf`` was in flight; llama-server would then come up afterwards and double-own the GPU. + + Round 16 P1 #4: a missing or unavailable llama backend is a + silent no-op (fresh install / no GGUF use), but an unload that + actually FAILS raises 503 so the caller does not start a new GPU + workload while llama-server is still resident. """ try: llama = get_llama_cpp_backend() - is_loaded = bool(getattr(llama, "is_loaded", False)) - is_active = bool(getattr(llama, "is_active", False)) - is_loading = bool(getattr(llama, "loading_model_identifier", None)) - if is_loaded or is_active or is_loading: - logger.info( - "Unloading GGUF chat (loaded=%s active=%s loading=%s) before %s load", - is_loaded, - is_active, - is_loading, - workload, - ) - await asyncio.to_thread(llama.unload_model) - except Exception as e: - logger.debug("llama-server unload skipped for %s: %s", workload, e) + except Exception as exc: + logger.debug("llama-server unavailable for %s: %s", workload, exc) + return + + is_loaded = bool(getattr(llama, "is_loaded", False)) + is_active = bool(getattr(llama, "is_active", False)) + is_loading = bool(getattr(llama, "loading_model_identifier", None)) + if not (is_loaded or is_active or is_loading): + return + + logger.info( + "Unloading GGUF chat (loaded=%s active=%s loading=%s) before %s load", + is_loaded, + is_active, + is_loading, + workload, + ) + try: + await asyncio.to_thread(llama.unload_model) + except Exception as exc: + logger.warning( + "Failed to unload GGUF chat before %s load: %s", workload, exc + ) + raise HTTPException( + status_code = 503, + detail = ( + f"Could not unload the existing GGUF chat model before " + f"starting {workload}." + ), + ) from exc async def _release_safetensors_chat_for(workload: str) -> None: """Unload the safetensors / Unsloth chat backend (drains both ``active_model_name`` and ``loading_models``) if it owns the GPU. + + Round 16 P1 #4: ``unload_model`` returning ``False`` (subprocess + wedged, IPC timeout) used to be silently ignored, leaving the + old chat model resident while a new GPU workload started on top. + Treat ``False`` as failure and raise 503 so the caller retries + instead of double-owning VRAM. """ try: from core.inference import get_inference_backend as _gib # type: ignore inf = _gib() - active_model_name = getattr(inf, "active_model_name", None) - loading_models = set(getattr(inf, "loading_models", set()) or set()) - if active_model_name: - logger.info( - "Unloading safetensors chat '%s' before %s load", - active_model_name, - workload, + except Exception as exc: + logger.debug("safetensors unavailable for %s: %s", workload, exc) + return + + async def _unload_required(model_name: str) -> None: + try: + ok = await asyncio.to_thread(inf.unload_model, model_name) + except Exception as exc: + raise HTTPException( + status_code = 503, + detail = ( + f"Could not unload safetensors chat model " + f"'{model_name}' before starting {workload}." + ), + ) from exc + if ok is False: + raise HTTPException( + status_code = 503, + detail = ( + f"Safetensors backend refused to unload " + f"'{model_name}' before starting {workload}. " + "Try again." + ), ) - await asyncio.to_thread(inf.unload_model, active_model_name) - for loading in loading_models: - if loading == active_model_name: - continue - try: - logger.info( - "Unloading in-flight safetensors chat '%s' before %s load", - loading, - workload, - ) - await asyncio.to_thread(inf.unload_model, loading) - except Exception as inner: - logger.debug( - "loading safetensors unload skipped for %s: %s", - loading, - inner, - ) - except Exception as e: - logger.debug("safetensors unload skipped for %s: %s", workload, e) + + active_model_name = getattr(inf, "active_model_name", None) + loading_models = set(getattr(inf, "loading_models", set()) or set()) + if active_model_name: + logger.info( + "Unloading safetensors chat '%s' before %s load", + active_model_name, + workload, + ) + await _unload_required(active_model_name) + for loading in loading_models: + if loading == active_model_name: + continue + logger.info( + "Unloading in-flight safetensors chat '%s' before %s load", + loading, + workload, + ) + await _unload_required(loading) async def _release_chat_for(workload: str) -> None: @@ -1942,18 +1985,21 @@ async def diffusion_load( ) return JSONResponse(content = status) except RuntimeError as exc: - # Round 15 P2 #7: if a training run / export job starts - # between the route-level pre-check and the backend worker, - # ``_release_other_gpu_owners_for_diffusion`` raises a - # RuntimeError that should surface as a 409 conflict (the - # same status the route layer returns), not 400. Match the - # known conflict strings the backend raises. + # Round 15 P2 #7 / round 16 P2 #7: backend-level conflict + # checks raise RuntimeError that surfaces here. Distinguish: + # - "Could not verify ..." -> 503 (retryable, status check + # itself failed), matching the route-level pre-check. + # - explicit "currently active" -> 409 conflict. + # - anything else -> 400 (bad request). detail = str(exc) + if ( + "Could not verify training status" in detail + or "Could not verify export status" in detail + ): + raise HTTPException(status_code = 503, detail = detail) from exc if ( "export job is currently active" in detail or "Training is currently active" in detail - or "Could not verify training status" in detail - or "Could not verify export status" in detail ): raise HTTPException(status_code = 409, detail = detail) from exc raise HTTPException(status_code = 400, detail = detail) from exc diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 10e6e26822..1ebf525f56 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -2062,7 +2062,9 @@ async def delete_finetuned_model( from core.inference.diffusion import get_diffusion_backend diff_backend = get_diffusion_backend() - diff_status = diff_backend.status() + # include_internal=True so we can iterate active_*/pending_* + # raw paths against ``target_path`` (round 16 P1 #5). + diff_status = diff_backend.status(include_internal = True) if diff_status.get("is_loaded") or diff_status.get("is_loading"): target_str = str(target_path) # Pair each owned repo / path with the GGUF variant it @@ -2757,17 +2759,32 @@ async def delete_cached_model( loading_id = ( getattr(llama_backend, "loading_model_identifier", None) or "" ).lower() + loading_variant = ( + getattr(llama_backend, "loading_hf_variant", None) or "" + ).lower() # Also consult the pending-load identifier: a multi-GB HF # download stays in ``loading_model_identifier`` until the # download completes, before ``model_identifier`` is set # (round 13 P1 #6). Without this check the cache directory # the download was writing into could be rmtree'd mid-flight. + # Round 16 P1 #1: pair against ``loading_hf_variant`` so a + # delete of a DIFFERENT cached quant from the same repo + # (loading Q4_K_M, deleting cached Q8_0) is allowed; only + # block when the requested variant matches what is being + # downloaded. Mirrors the /delete-finetuned pairing. needle = repo_id.lower() + requested_variant = (variant or "").lower() if loading_id == needle: - raise HTTPException( - status_code = 409, - detail = "Cannot delete a model while it is loading", + same_loading_variant = ( + not requested_variant + or not loading_variant + or requested_variant == loading_variant ) + if same_loading_variant: + raise HTTPException( + status_code = 409, + detail = "Cannot delete a model while it is loading", + ) # Exact match only (case-insensitive). Prefix match would # block deleting unrelated ``org/model`` while # ``org/model-v2`` is loaded -- same surface the diffusion @@ -2778,7 +2795,6 @@ async def delete_cached_model( llama_backend.is_loaded or getattr(llama_backend, "is_active", False) ): loaded_variant = (getattr(llama_backend, "hf_variant", None) or "").lower() - requested_variant = (variant or "").lower() same_variant = ( not requested_variant or not loaded_variant @@ -2854,7 +2870,9 @@ async def delete_cached_model( from core.inference.diffusion import get_diffusion_backend diff_backend = get_diffusion_backend() - diff_status = diff_backend.status() + # include_internal=True so we can pair owned raw paths against + # the HF cache snapshot root (round 16 P1 #5). + diff_status = diff_backend.status(include_internal = True) if diff_status.get("is_loaded") or diff_status.get("is_loading"): needle = repo_id.lower() # Round 15 P1 #4: ALSO compare owned paths against the HF @@ -2879,10 +2897,22 @@ async def delete_cached_model( except Exception: pass except Exception as cache_scan_exc: - logger.debug( - "HF cache scan failed during diffusion delete guard: %s", + # Round 16 P1 #3: a transient cache-scan failure here + # used to silently fall through to repo-id-only + # matching, which misses local snapshot paths and + # let /delete-cached unlink an actively mmap'd + # snapshot. Fail-closed (503) so the user retries. + logger.warning( + "Could not scan HF cache during diffusion delete guard: %s", cache_scan_exc, ) + raise HTTPException( + status_code = 503, + detail = ( + "Could not verify diffusion cache ownership before " + "deleting. Try again." + ), + ) from cache_scan_exc # Pair each owned repo with the GGUF variant it actually # owns (active or pending) so a swap in progress does not diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index f9ab1b8968..584b462d34 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -174,6 +174,9 @@ def test_get_diffusion_backend_singleton(): def test_status_shape_unloaded(): + """Public status() (the browser-facing payload) must NOT contain + the guard-only ``active_*`` / ``pending_*`` fields (round 16 + P1 #5).""" from core.inference.diffusion import get_diffusion_backend s = get_diffusion_backend().status() @@ -185,12 +188,6 @@ def test_status_shape_unloaded(): "pipeline_class", "base_repo", "gguf_filename", - "active_repo_id", - "active_base_repo", - "active_gguf_filename", - "pending_repo_id", - "pending_base_repo", - "pending_gguf_filename", "device", "dtype", "loaded_at", @@ -198,10 +195,23 @@ def test_status_shape_unloaded(): "supported_families", } assert expected_keys.issubset(s.keys()) + # Guard-facing fields are gated behind include_internal=True. + for guard_key in ( + "active_repo_id", + "active_base_repo", + "active_gguf_filename", + "pending_repo_id", + "pending_base_repo", + "pending_gguf_filename", + ): + assert guard_key not in s, f"public status() must not expose {guard_key}" assert s["is_loaded"] is False assert s["repo_id"] is None - assert s["active_gguf_filename"] is None - assert s["pending_gguf_filename"] is None + + # Internal status() exposes the guard fields for delete/route use. + s_internal = get_diffusion_backend().status(include_internal = True) + assert s_internal["active_gguf_filename"] is None + assert s_internal["pending_gguf_filename"] is None # ── encode_png_base64 ─────────────────────────────────────────── @@ -1413,7 +1423,7 @@ def test_status_preserves_active_gguf_subdir(monkeypatch): aliases = (), ) - s = backend.status() + s = backend.status(include_internal = True) assert s["active_gguf_filename"] == "BF16/model.gguf" # UI-facing field still collapses to the basename. assert s["gguf_filename"] == "model.gguf" @@ -1504,6 +1514,82 @@ def test_detect_family_rejects_substring_collisions(): assert fam is not None and fam.name == "flux.2" +def test_detect_family_compact_aliases_with_owner_prefix(): + """Round 16 P2 #9: compact aliases must match when the repo has + an owner prefix. ``unsloth/Flux2Klein-GGUF`` -> flux.2-klein + via the ``flux2-klein`` alias's compact form. Embedded compact + matches (e.g. ``flux2`` inside ``flux20``) must NOT match.""" + from core.inference.diffusion import detect_family + + fam = detect_family("unsloth/Flux2Klein-GGUF") + assert fam is not None and fam.name == "flux.2-klein" + # 20 is a different number; must not collide with flux.2. + assert detect_family("unsloth/Flux20-GGUF") is None + + +def test_public_status_does_not_leak_local_path_via_active_fields(monkeypatch): + """Round 16 P1 #5: even the guard-facing active_*/pending_* keys + must be absent from the public status payload.""" + import core.inference.diffusion as d + + backend = d.DiffusionBackend() + backend._pipe = object() + backend._repo_id = "/home/alice/private-flux" + backend._base_repo = "/home/alice/base-private" + backend._family = d.DiffusionFamily( + name = "flux.2-klein", + pipeline_class = "Flux2KleinPipeline", + transformer_class = "Flux2Transformer2DModel", + base_repo = "black-forest-labs/FLUX.2-klein-4B", + aliases = (), + ) + + public = backend.status() + # UI-facing fields collapse to leaf and the guard-only fields are absent. + assert public["repo_id"] == "private-flux" + assert public["base_repo"] == "base-private" + for key in ( + "active_repo_id", + "active_base_repo", + "active_gguf_filename", + "pending_repo_id", + "pending_base_repo", + "pending_gguf_filename", + ): + assert key not in public + + internal = backend.status(include_internal = True) + assert internal["active_repo_id"] == "/home/alice/private-flux" + assert internal["active_base_repo"] == "/home/alice/base-private" + + +def test_generate_image_with_metadata_redacts_local_path(monkeypatch): + """Round 16 P1 #6: the generation response must not echo a raw + absolute path back to the browser.""" + import core.inference.diffusion as d + + backend = d.DiffusionBackend() + backend._pipe = object() + backend._repo_id = "/home/alice/private/secret-flux" + backend._family = d.DiffusionFamily( + name = "flux.2-klein", + pipeline_class = "Flux2KleinPipeline", + transformer_class = "Flux2Transformer2DModel", + base_repo = "black-forest-labs/FLUX.2-klein-4B", + aliases = (), + ) + + def _fake_unlocked(**kwargs): + from PIL import Image as _Image + + return _Image.new("RGB", (8, 8)) + + monkeypatch.setattr(backend, "_generate_image_unlocked", _fake_unlocked) + _, meta = backend.generate_image_with_metadata(prompt = "x") + assert meta["model"] == "secret-flux" + assert "/home/alice" not in meta["model"] + + def test_release_other_gpu_owners_raises_on_active_training(monkeypatch): """Round 15 P1 #3: direct backend callers must not bypass the route layer's training-active 409 guard."""