From ae41bfdbfdaa21551d2a427b7755f2b63e65eed2 Mon Sep 17 00:00:00 2001 From: Daniel Han-Chen Date: Mon, 25 May 2026 05:57:32 +0000 Subject: [PATCH] Fix/adjust diffusion: round 13 P1+P2 batch for PR #5754 Round 13 reviewer aggregate (logs/review_round13_aggregate.md): P1 fixes: - routes/export.py load_checkpoint refuses (409) when an export job is currently active, mirroring the chat/diffusion/training handoff guards. ``is_export_active`` absence is tolerated for older / mocked backends. - core/inference/diffusion.py local-path GGUF loader now accepts relative directories (Studio exports surface as ``exports/my-flux``) and confines ``gguf_filename`` to the chosen repo via ``_resolve_local_gguf_child``: absolute filenames, ``..`` segments, and Windows separators are rejected before any file is opened. - core/inference/diffusion.py status() exposes ``active_gguf_filename`` alongside the pending variant so delete guards can pair each owned repo with the GGUF variant it actually owns. - routes/models.py cache delete + finetuned delete adopt a shared ``_diffusion_owned_targets`` + ``_variant_delete_is_safe_for_owned_gguf`` helper. Per-variant deletes during a swap-in-flight cannot remove the active variant while the pending variant is loading. - core/inference/llama_cpp.py publishes ``loading_model_identifier`` before ``_download_gguf`` starts and clears it in ``finally``. Cache delete (routes/models.py) and the cross-workload release helpers (routes/inference.py::_release_llama_for and diffusion.py::_release_chat_backend_for_diffusion) consult it so a multi-GB HF download cannot be rmtree'd or be ignored by /images/load while still in flight. P2 fixes: - core/inference/diffusion.py adds ``generate_image_with_metadata`` + ``async_generate_with_metadata``; /images/generate uses it so the response model/family reflect the pipeline that actually produced the image even if an unload races the route. - core/inference/diffusion.py: ``base_repo`` only applies when picking a GGUF quant. Filling Base diffusers repo while loading a full diffusers repo no longer silently swaps the load target. - core/inference/diffusion.py: failed device placement / offload now drops pipe + transformer references explicitly before drain so partial allocations cannot keep VRAM around. - core/inference/diffusion.py: torch/diffusers imports surface as a clear RuntimeError naming the missing dependency. - core/inference/diffusion.py: _smart_base_repo splits on both POSIX and Windows separators so ``C:\\Users\\me\\base\\FLUX.2-klein-4B-GGUF`` no longer picks the Base 4B variant via the parent dir. Tests: - 6 new regression cases (Windows leaf, traversal/backslash rejection, relative-dir local load, metadata snapshot, lock serialisation). - All 59 diffusion backend + route tests pass. --- studio/backend/core/inference/diffusion.py | 374 +++++++++++++----- studio/backend/core/inference/llama_cpp.py | 71 +++- studio/backend/routes/export.py | 34 ++ studio/backend/routes/inference.py | 35 +- studio/backend/routes/models.py | 164 ++++---- .../backend/tests/test_diffusion_backend.py | 225 +++++++++++ studio/backend/tests/test_diffusion_routes.py | 16 + 7 files changed, 715 insertions(+), 204 deletions(-) diff --git a/studio/backend/core/inference/diffusion.py b/studio/backend/core/inference/diffusion.py index 93af5d2fd6..335e93c957 100644 --- a/studio/backend/core/inference/diffusion.py +++ b/studio/backend/core/inference/diffusion.py @@ -35,10 +35,11 @@ from __future__ import annotations import asyncio import gc import io +import re import threading import time from dataclasses import dataclass, field -from pathlib import Path +from pathlib import Path, PurePosixPath from typing import Any, Optional from loggers import get_logger @@ -159,11 +160,15 @@ def _smart_base_repo(fam: DiffusionFamily, repo_id: str) -> str: 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). + variant (round 12 review #9). Splits on BOTH ``/`` and ``\\`` so + Windows local paths like ``C:\\Users\\me\\base\\FLUX.2-klein-4B`` + do not get scored as "base" via the parent directory either + (round 13 P2 #13). """ if fam.name != "flux.2-klein": return fam.base_repo - last_segment = (repo_id or "").rstrip("/").rsplit("/", 1)[-1].lower() + cleaned = (repo_id or "").rstrip("/\\") + last_segment = re.split(r"[\\/]+", cleaned)[-1].lower() if cleaned else "" is_9b = "9b" in last_segment is_base = "base" in last_segment if is_9b and is_base: @@ -177,6 +182,47 @@ def _smart_base_repo(fam: DiffusionFamily, repo_id: str) -> str: return "black-forest-labs/FLUX.2-klein-4B" +def _resolve_local_gguf_child(repo_root: Path, gguf_filename: str) -> Path: + """Resolve a GGUF filename inside a local repo directory safely. + + Returns the resolved absolute path or raises ``RuntimeError`` if: + - ``gguf_filename`` is absolute (``/etc/passwd``) or contains a + Windows separator (``..\\..\\secret.gguf``); + - the parts contain ``""`` / ``.`` / ``..`` (``../other.gguf``); + - the resolved candidate escapes ``repo_root`` after symlinks / + ``..`` collapse; + - the resolved candidate is not a regular file. + + This is the only path that bridges a user-supplied ``gguf_filename`` + string into ``Path``s the loader opens, so confining it to the + chosen repo here protects the delete-ownership guards downstream + (round 13 P1 #2). ``hf_hub_download`` already enforces the same + invariant for Hub repos. + """ + if Path(gguf_filename).is_absolute() or "\\" in gguf_filename: + raise RuntimeError( + "gguf_filename must be a relative file path inside repo_id." + ) + rel = PurePosixPath(gguf_filename) + if any(part in ("", ".", "..") for part in rel.parts): + raise RuntimeError( + "gguf_filename must not contain empty, '.', or '..' segments." + ) + root = repo_root.expanduser().resolve(strict = True) + candidate = (root / Path(*rel.parts)).resolve(strict = True) + try: + candidate.relative_to(root) + except ValueError as exc: + raise RuntimeError( + "gguf_filename must stay inside the local repo_id directory." + ) from exc + if not candidate.is_file(): + raise RuntimeError( + f"Local repo path '{repo_root}' does not contain '{gguf_filename}'." + ) + return candidate + + # Negative substrings that disqualify a candidate family even when its # name appears as a substring of the repo id. Prevents # "stable-diffusion-3" matching SD3.5 and "qwen-image" matching @@ -335,6 +381,7 @@ class DiffusionBackend: # user just clicked. active_repo = self._repo_id active_base = self._base_repo + active_gguf = gguf_basename 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 @@ -357,11 +404,15 @@ class DiffusionBackend: "family": ui_family, "pipeline_class": ui_pipeline_class, "base_repo": pending_base or active_base, - "gguf_filename": pending_gguf or gguf_basename, - # Guard-facing fields: every repo / path the backend - # owns RIGHT NOW. Delete routes iterate both. + "gguf_filename": pending_gguf or active_gguf, + # 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, @@ -431,9 +482,24 @@ class DiffusionBackend: keep peak VRAM bounded; status() reports is_loaded=false with last_error set so the caller can react. """ - from huggingface_hub import hf_hub_download - import diffusers - import torch + # Surface a friendly load error when the no-torch / partial + # install path is active: the user clicked Load on the Images + # page but the runtime never installed torch + diffusers (round + # 13 P2 #12). Without this wrapper the import surfaces as a + # raw ``ModuleNotFoundError`` -> 500 instead of a 400 the UI + # can display. + try: + from huggingface_hub import hf_hub_download + import diffusers + import torch + except ModuleNotFoundError as exc: + missing = exc.name or str(exc) + raise RuntimeError( + "Diffusion image generation requires the torch / diffusers " + f"runtime. Missing dependency: {missing}. Install the Studio " + "torch runtime (re-run setup.sh / install.ps1) before " + "loading an image model." + ) from exc fam = detect_family(repo_id, override_family = family_override) if fam is None: @@ -485,19 +551,21 @@ class DiffusionBackend: # Resolution rules for the "what repo to call # from_pretrained on" question: - # 1. caller-supplied base_repo wins - # 2. if no GGUF file was requested the user is loading a - # full diffusers repo; use repo_id directly so we do + # 1. no GGUF file -> caller is loading a full + # diffusers repo; use repo_id directly so we do # not silently substitute the family default - # 3. otherwise use the family + repo_id heuristic so a - # 9B GGUF picks the 9B base, not the 4B fallback - if base_repo: - effective_base = base_repo - # Refresh pending so delete guards see the actual - # base, not just caller-supplied None. - with self._lock: - self._pending_base_repo = effective_base - elif not gguf_filename: + # AND ignore any base_repo input (it is only + # meaningful as a GGUF companion override). The + # old order let ``base_repo`` swap a fine-tuned + # ``owner/my-flux.1-finetune`` for + # ``black-forest-labs/FLUX.1-dev`` while status + # still advertised the user's repo (round 13 + # P2 #10). + # 2. otherwise prefer caller-supplied base_repo for + # the missing VAE / text encoder components. + # 3. otherwise use the family + repo_id heuristic so + # a 9B GGUF picks the 9B base, not the 4B fallback. + if not gguf_filename: # Guard: a repo that ends in "-GGUF" (the unsloth # convention) is GGUF-only and will 500 on # from_pretrained; surface a clear error instead of @@ -507,12 +575,18 @@ class DiffusionBackend: raise RuntimeError( f"'{repo_id}' looks like a GGUF-only repo. " "Either provide gguf_filename to pick a quant, " - "or pass base_repo to override the full-repo " - "load target." + "or load a full diffusers repo (base_repo only " + "applies when picking a GGUF quant)." ) effective_base = repo_id with self._lock: self._pending_base_repo = effective_base + elif base_repo: + effective_base = base_repo + # Refresh pending so delete guards see the actual + # base, not just caller-supplied None. + with self._lock: + self._pending_base_repo = effective_base else: effective_base = _smart_base_repo(fam, repo_id) with self._lock: @@ -535,21 +609,26 @@ class DiffusionBackend: "path wired in this build; load the full repo instead." ) # 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". + # accept either a Hub repo id OR a local path + # (Studio export, downloaded HF snapshot, etc.). + # We accept BOTH absolute and relative local + # directories: Studio exports surface as relative + # paths like ``exports/my-flux`` and earlier + # versions only accepted absolute paths, falling + # through to ``hf_hub_download`` which then + # raised HFValidationError on the relative path + # (round 13 P1 #2). For local paths we route the + # gguf_filename through ``_resolve_local_gguf_child`` + # so traversal (``../secret.gguf``) and absolute + # filename escapes (``/etc/passwd``) are rejected + # BEFORE the file is opened, which also keeps the + # delete-ownership guards aligned with what was + # actually loaded. 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) + if repo_id_path.is_dir(): + local_gguf_path = str( + _resolve_local_gguf_child(repo_id_path, gguf_filename) + ) else: local_gguf_path = hf_hub_download( repo_id = repo_id, @@ -638,24 +717,31 @@ class DiffusionBackend: if hf_token: pipe_kwargs["token"] = hf_token + pipe = None try: pipe = pipeline_cls.from_pretrained(effective_base, **pipe_kwargs) + # Device placement / offload can ALSO raise after + # from_pretrained succeeded (OOM at the .to(device) + # copy, accelerate offload hook misconfigured, etc.). + # If we let the exception escape now, the local + # ``pipe`` lives on the traceback frame until the + # caller drops it, holding multi-GB of VRAM behind + # 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": + pipe.enable_model_cpu_offload() + else: + pipe.to(device) except Exception: - # If from_pretrained fails after the transformer was - # already loaded, the transformer object holds GPU - # weights that would only be freed at GC. Drop the - # local reference and force a collect so the next - # load attempt does not stack VRAM with a phantom - # transformer. + if pipe is not None: + _release(pipe) + pipe = None if transformer is not None: _release(transformer) transformer = None - _drain_cuda_cache() + _drain_cuda_cache() raise - if enable_model_cpu_offload and device == "cuda": - pipe.enable_model_cpu_offload() - else: - pipe.to(device) with self._lock: self._pipe = pipe @@ -775,6 +861,39 @@ class DiffusionBackend: requests for the entire (minutes-long) generation, which made the UI feel frozen. """ + # Take _generate_lock FIRST so a concurrent unload/load that + # observes us holding it will queue behind this generation + # (and `unload_model` then waits its turn before clearing + # state). Snapshotting `self._pipe` outside the lock and then + # taking the lock let a load/unload race in between, so the + # forward could run against a freed or swapped pipeline. + with self._generate_lock: + return self._generate_image_unlocked( + prompt = prompt, + negative_prompt = negative_prompt, + num_inference_steps = num_inference_steps, + guidance_scale = guidance_scale, + width = width, + height = height, + seed = seed, + ) + + def _generate_image_unlocked( + self, + *, + prompt: str, + negative_prompt: Optional[str] = None, + num_inference_steps: int = 24, + guidance_scale: float = 3.5, + width: int = 1024, + height: int = 1024, + seed: Optional[int] = None, + ) -> "Any": + """Inner body of ``generate_image`` that ASSUMES the caller + already holds ``_generate_lock``. Lets + ``generate_image_with_metadata`` snapshot metadata under the + same lock without deadlocking on a non-reentrant + ``threading.Lock`` (round 13 P2 #9).""" if not prompt or not prompt.strip(): raise ValueError("prompt is empty") if num_inference_steps < 1 or num_inference_steps > 200: @@ -789,64 +908,81 @@ class DiffusionBackend: import torch - # Take _generate_lock FIRST so a concurrent unload/load that - # observes us holding it will queue behind this generation - # (and `unload_model` then waits its turn before clearing - # state). Snapshotting `self._pipe` outside the lock and then - # taking the lock let a load/unload race in between, so the - # forward could run against a freed or swapped pipeline. - with self._generate_lock: - with self._lock: - if self._pipe is None: - raise RuntimeError("No diffusion model is loaded.") - pipe = self._pipe - device = self._device or "cpu" - 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" + with self._lock: + if self._pipe is None: + raise RuntimeError("No diffusion model is loaded.") + pipe = self._pipe + device = self._device or "cpu" + 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" + ) + generator = torch.Generator(device = gen_device).manual_seed(int(seed)) + + call_kwargs: dict[str, Any] = { + "prompt": prompt, + "num_inference_steps": int(num_inference_steps), + "guidance_scale": float(guidance_scale), + "width": int(width), + "height": int(height), + } + # FLUX.2 / FLUX.2 klein pipelines do NOT accept + # negative_prompt and 500 if you pass it in. Inspect the + # signature and only forward when supported; warn otherwise + # so the UI can disable the field for incompatible families. + if negative_prompt is not None and negative_prompt.strip(): + if _pipe_accepts_kwarg(pipe, "negative_prompt"): + call_kwargs["negative_prompt"] = negative_prompt + # QwenImagePipeline and FluxPipeline treat + # guidance_scale as distilled CFG and use + # true_cfg_scale as the real classifier-free + # guidance knob; the negative prompt is only + # effective when true_cfg_scale > 1. Forward the + # user-supplied guidance_scale through both so the + # negative prompt actually steers generation. + if _pipe_accepts_kwarg(pipe, "true_cfg_scale"): + call_kwargs["true_cfg_scale"] = float(guidance_scale) + else: + logger.info( + "Dropping negative_prompt: %s does not accept it", + type(pipe).__name__, ) - generator = torch.Generator(device = gen_device).manual_seed(int(seed)) + if generator is not None: + call_kwargs["generator"] = generator - call_kwargs: dict[str, Any] = { - "prompt": prompt, - "num_inference_steps": int(num_inference_steps), - "guidance_scale": float(guidance_scale), - "width": int(width), - "height": int(height), - } - # FLUX.2 / FLUX.2 klein pipelines do NOT accept - # negative_prompt and 500 if you pass it in. Inspect the - # signature and only forward when supported; warn otherwise - # so the UI can disable the field for incompatible families. - if negative_prompt is not None and negative_prompt.strip(): - if _pipe_accepts_kwarg(pipe, "negative_prompt"): - call_kwargs["negative_prompt"] = negative_prompt - # QwenImagePipeline and FluxPipeline treat - # guidance_scale as distilled CFG and use - # true_cfg_scale as the real classifier-free - # guidance knob; the negative prompt is only - # effective when true_cfg_scale > 1. Forward the - # user-supplied guidance_scale through both so the - # negative prompt actually steers generation. - if _pipe_accepts_kwarg(pipe, "true_cfg_scale"): - call_kwargs["true_cfg_scale"] = float(guidance_scale) - else: - logger.info( - "Dropping negative_prompt: %s does not accept it", - type(pipe).__name__, - ) - if generator is not None: - call_kwargs["generator"] = generator + out = pipe(**call_kwargs) + images = getattr(out, "images", None) or [] + if not images: + raise RuntimeError("Diffusion pipeline returned no images.") + return images[0] - out = pipe(**call_kwargs) - images = getattr(out, "images", None) or [] - if not images: - raise RuntimeError("Diffusion pipeline returned no images.") - return images[0] + def generate_image_with_metadata( + self, + **kwargs: Any, + ) -> tuple[Any, dict[str, Any]]: + """Generate a single image AND snapshot its identifying metadata. + + Returns ``(pil_image, {"model": , "family": })`` + where the metadata reflects the pipeline that produced the + image. Snapshotted under ``_generate_lock + _lock`` so a + queued unload / load that promotes a different pipeline + cannot replace ``self._repo_id`` / ``self._family`` between + the forward returning and the route reading status (round + 13 P2 #9). The route uses these values directly in the + response instead of re-calling ``status()``. + """ + with self._generate_lock: + image = self._generate_image_unlocked(**kwargs) + with self._lock: + meta = { + "model": self._repo_id, + "family": self._family.name if self._family else None, + } + return image, meta def _pipe_accepts_kwarg(pipe: Any, name: str) -> bool: @@ -895,21 +1031,24 @@ def _release_chat_backend_for_diffusion() -> None: """ # 1. GGUF chat backend (llama-server subprocess). We unload when # EITHER is_loaded is True (resident model) OR is_active is - # True (mid-download / startup); the latter case is the - # "llama-server is currently starting" race where weights are - # being downloaded and the diffusion load would otherwise - # double-spend GPU memory. + # True (mid-download / startup) OR loading_model_identifier is + # populated (HF GGUF download in progress, before is_active / + # is_loaded flip). The last case is what round 13 P1 #8 flagged: + # a multi-GB HF download from one workload + a diffusion load + # racing on the same GPU would otherwise both end up live. try: from routes.inference import get_llama_cpp_backend # type: ignore backend = get_llama_cpp_backend() is_loaded = bool(getattr(backend, "is_loaded", False)) is_active = bool(getattr(backend, "is_active", False)) - if is_loaded or is_active: + is_loading = bool(getattr(backend, "loading_model_identifier", None)) + if is_loaded or is_active or is_loading: logger.info( - "Unloading llama-server (loaded=%s active=%s) before diffusion load", + "Unloading llama-server (loaded=%s active=%s loading=%s) before diffusion load", is_loaded, is_active, + is_loading, ) backend.unload_model() except Exception as exc: @@ -1086,3 +1225,20 @@ async def async_generate( do not block the event loop for the 5-30 s a diffusion step takes.""" loop = asyncio.get_event_loop() return await loop.run_in_executor(None, lambda: backend.generate_image(**kwargs)) + + +async def async_generate_with_metadata( + backend: DiffusionBackend, + **kwargs: Any, +) -> tuple[Any, dict[str, Any]]: + """Run ``generate_image_with_metadata`` in the default executor. + + Used by the /images/generate route so the response model / family + fields reflect the pipeline that actually produced the image, even + if an unload races the route between the forward returning and the + response being assembled (round 13 P2 #9).""" + loop = asyncio.get_event_loop() + return await loop.run_in_executor( + None, + lambda: backend.generate_image_with_metadata(**kwargs), + ) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index bf8a3c04df..57cda178a9 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -612,6 +612,13 @@ class LlamaCppBackend: self._process: Optional[subprocess.Popen] = None self._port: Optional[int] = None self._model_identifier: Optional[str] = None + # Pending-load identifier: set BEFORE _download_gguf starts and + # cleared after the load finishes (success or failure). Delete + # guards and cross-workload handoff helpers read it via + # ``loading_model_identifier`` so a multi-GB HF download cannot + # have its cache rmtree'd or be ignored by /images/load, + # /training/start, /export/load while it is still resolving. + self._loading_model_identifier: Optional[str] = None self._gguf_path: Optional[str] = None self._hf_repo: Optional[str] = None self._hf_variant: Optional[str] = None @@ -713,6 +720,19 @@ class LlamaCppBackend: def model_identifier(self) -> Optional[str]: return self._model_identifier + @property + def loading_model_identifier(self) -> Optional[str]: + """Identifier of a load currently in progress, or None. + + Populated while ``_download_gguf`` is fetching the GGUF for a + new ``load_model`` call. Cleared in the surrounding + ``finally`` block, so a failed load leaves it None. Delete + guards in ``routes/models.py`` and handoff helpers in + ``routes/inference.py`` consult this so a long HF download + cannot have its destination rmtree'd or be ignored by a + concurrent /images/load that thinks llama-server is idle.""" + return self._loading_model_identifier + @property def is_vision(self) -> bool: return self._is_vision @@ -2673,25 +2693,44 @@ 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. - 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( + # + # 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( hf_repo = hf_repo, + hf_variant = hf_variant, 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") + # 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 # Set identifier early so _read_gguf_metadata can use it for DeepSeek detection self._model_identifier = model_identifier diff --git a/studio/backend/routes/export.py b/studio/backend/routes/export.py index aabcd1920b..11c5004ced 100644 --- a/studio/backend/routes/export.py +++ b/studio/backend/routes/export.py @@ -142,6 +142,40 @@ async def load_checkpoint( 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. diff --git a/studio/backend/routes/inference.py b/studio/backend/routes/inference.py index c4e247c2a1..d376a289e1 100644 --- a/studio/backend/routes/inference.py +++ b/studio/backend/routes/inference.py @@ -359,18 +359,25 @@ def _raise_if_export_active(workload: str) -> None: async def _release_llama_for(workload: str) -> None: """Unload the llama-server (GGUF) chat backend if it owns the - GPU. Treats ``is_loaded`` OR ``is_active`` as held (the latter - is mid-download / mid-startup, before health probes pass). + GPU. Treats ``is_loaded`` OR ``is_active`` OR + ``loading_model_identifier`` as held: the third covers an HF GGUF + download that has not yet flipped ``is_active`` to True (round + 13 P1 #7). Without it, /images/load, /training/start, and + /export/load could start while a long ``_download_gguf`` was in + flight; llama-server would then come up afterwards and double-own + the GPU. """ try: llama = get_llama_cpp_backend() is_loaded = bool(getattr(llama, "is_loaded", False)) is_active = bool(getattr(llama, "is_active", False)) - if is_loaded or is_active: + 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) before %s load", + "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) @@ -1985,9 +1992,16 @@ async def diffusion_generate( start = time.time() try: - from core.inference.diffusion import async_generate, encode_png_base64 + from core.inference.diffusion import ( + async_generate_with_metadata, + encode_png_base64, + ) - image = await async_generate( + # ``async_generate_with_metadata`` snapshots ``model`` / + # ``family`` under the same ``_generate_lock`` that owns the + # forward, so a queued unload/load cannot replace them between + # generation end and response assembly (round 13 P2 #9). + image, meta = await async_generate_with_metadata( backend, prompt = payload.prompt, negative_prompt = payload.negative_prompt, @@ -2006,7 +2020,6 @@ async def diffusion_generate( raise HTTPException(status_code = 500, detail = str(exc)) duration_ms = int((time.time() - start) * 1000) - status = backend.status() return DiffusionGenerateResponse( image_b64 = encode_png_base64(image), image_mime = "image/png", @@ -2022,12 +2035,8 @@ async def diffusion_generate( # browser side. seed_str = str(payload.seed) if payload.seed is not None else None, duration_ms = duration_ms, - # Use ``active_repo_id`` (the pipeline that just ran the - # forward) rather than the UI-facing ``repo_id`` so a - # queued /images/load promoting a new pending model cannot - # leak that model's identity into our response. - model = status.get("active_repo_id") or status.get("repo_id"), - family = status.get("family"), + model = meta.get("model"), + family = meta.get("family"), ) diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index 570f024bf8..a573f88233 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -1709,6 +1709,53 @@ def _is_path_under(path: Path, root: Path) -> bool: return False +def _diffusion_owned_targets(diff_status: dict) -> list[tuple[str, str | None]]: + """Return ``(owned_repo_or_path, owned_gguf_filename)`` pairs for + every diffusion target the backend currently holds. + + Pairs the active / pending repo with the active / pending GGUF + filename (not the UI-facing collapsed ``gguf_filename``) so the + per-variant delete guards know which quant is actually owned by + each repo. Without this pairing, a swap in progress (active + ``Q4_K_S``, pending ``Q8_0``) collapsed both to the pending + variant and the active ``Q4_K_S`` GGUF could be deleted while + still mmap'd by the resident pipeline (round 13 P1 #3-5). + + Base repos are paired with ``None`` for the GGUF: the base / + component repo is loaded whole via ``from_pretrained`` and has no + per-variant delete to take advantage of. + """ + return [ + ( + diff_status.get("active_repo_id") or "", + diff_status.get("active_gguf_filename"), + ), + (diff_status.get("active_base_repo") or "", None), + ( + diff_status.get("pending_repo_id") or "", + diff_status.get("pending_gguf_filename"), + ), + (diff_status.get("pending_base_repo") or "", None), + ] + + +def _variant_delete_is_safe_for_owned_gguf( + requested_variant: str | None, + owned_gguf_filename: str | None, +) -> bool: + """True iff a per-variant delete for ``requested_variant`` against + a repo that owns ``owned_gguf_filename`` cannot remove the owned + file. + + Returns False (i.e. unsafe -> block the delete) when either + argument is missing so a NULL owned filename or a full-repo delete + (no variant) does not accidentally pass the guard.""" + if not requested_variant or not owned_gguf_filename: + return False + loaded_label = (_extract_quant_label(owned_gguf_filename.lower()) or "").lower() + return bool(loaded_label and loaded_label != requested_variant.lower()) + + def _is_path_under_lexically(path: Path, root: Path) -> bool: """Check containment without resolving the final path's symlink target.""" try: @@ -1991,27 +2038,17 @@ async def delete_finetuned_model( diff_backend = get_diffusion_backend() diff_status = diff_backend.status() if diff_status.get("is_loaded") or diff_status.get("is_loading"): - candidates: list[str] = [] - for key in ( - "active_repo_id", - "active_base_repo", - "pending_repo_id", - "pending_base_repo", - ): - v = diff_status.get(key) or "" - 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: + # Pair each owned repo / path with the GGUF variant it + # actually owns (round 13 P1 #5). For a swap in flight + # (active Q4_K_S, pending Q8_0) the active variant must + # NOT be deleted just because the pending variant uses + # a different quant. + for candidate, owned_gguf in _diffusion_owned_targets(diff_status): + if not candidate: + continue try: - candidate_path = Path(candidate).expanduser() + candidate_resolved = Path(candidate).expanduser().resolve() except Exception: continue # Relative paths (the user can do @@ -2019,27 +2056,23 @@ async def delete_finetuned_model( # legitimate path candidates; resolve against the # backend cwd so they can be compared with the # absolute ``target_path``. Round 8 review #11. - try: - candidate_resolved = candidate_path.resolve() - except Exception: - continue - if ( + overlaps = ( candidate_resolved == target_path or str(candidate_resolved) == target_str or _is_path_under(candidate_resolved, target_path) or _is_path_under(target_path, candidate_resolved) + ) + if not overlaps: + continue + if export_type == "gguf" and _variant_delete_is_safe_for_owned_gguf( + gguf_variant, + owned_gguf, ): - # 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", - ) + continue + raise HTTPException( + status_code = 400, + detail = "Unload the diffusion image model before deleting", + ) except HTTPException: raise except Exception as e: @@ -2695,12 +2728,25 @@ async def delete_cached_model( llama_backend = get_llama_cpp_backend() loaded_id = (llama_backend.model_identifier or "").lower() + loading_id = ( + getattr(llama_backend, "loading_model_identifier", 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. + needle = repo_id.lower() + if loading_id == needle: + 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 # guard fixed in round 5. - wants = loaded_id == repo_id.lower() - if wants and ( + if loaded_id == needle and ( llama_backend.is_loaded or getattr(llama_backend, "is_active", False) ): raise HTTPException( @@ -2775,35 +2821,21 @@ async def delete_cached_model( diff_status = diff_backend.status() if diff_status.get("is_loaded") or diff_status.get("is_loading"): needle = repo_id.lower() - owned = { - (diff_status.get("active_repo_id") or "").lower(), - (diff_status.get("active_base_repo") or "").lower(), - (diff_status.get("pending_repo_id") or "").lower(), - (diff_status.get("pending_base_repo") or "").lower(), - } - owned.discard("") - if needle in owned: - # 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", - ) + # Pair each owned repo with the GGUF variant it actually + # owns (active or pending) so a swap in progress does not + # collapse both quants into the pending one (round 13 + # P1 #4). Per-variant delete is still allowed if the + # requested variant differs from the variant that owns + # the matched repo. + for owned_id, owned_gguf in _diffusion_owned_targets(diff_status): + if not owned_id or owned_id.lower() != needle: + continue + if _variant_delete_is_safe_for_owned_gguf(variant, owned_gguf): + continue + raise HTTPException( + status_code = 400, + detail = "Unload the diffusion image model before deleting", + ) except HTTPException: raise except Exception as e: diff --git a/studio/backend/tests/test_diffusion_backend.py b/studio/backend/tests/test_diffusion_backend.py index 44386fc8ea..43918fb17e 100644 --- a/studio/backend/tests/test_diffusion_backend.py +++ b/studio/backend/tests/test_diffusion_backend.py @@ -24,6 +24,7 @@ import base64 import io import sys import types +from types import SimpleNamespace from typing import Any import pytest @@ -184,6 +185,12 @@ 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", @@ -193,6 +200,8 @@ def test_status_shape_unloaded(): assert expected_keys.issubset(s.keys()) 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 # ── encode_png_base64 ─────────────────────────────────────────── @@ -1192,3 +1201,219 @@ def test_bf16_falls_back_to_fp16_on_old_cuda(monkeypatch): device, dtype = backend._pick_device_and_dtype() assert device == "cuda" assert dtype is fake_torch.float16 + + +# ── round 13 regressions ────────────────────────────────────────── + + +def test_smart_base_repo_uses_windows_leaf_only(): + """Round 13 P2 #13: a Windows path whose PARENT directory contains + 'base' must not be misclassified as the Klein Base 4B variant.""" + from core.inference.diffusion import _smart_base_repo, detect_family + + repo = r"C:\Users\me\base\FLUX.2-klein-4B-GGUF" + fam = detect_family(repo) + assert fam is not None and fam.name == "flux.2-klein" + assert _smart_base_repo(fam, repo) == "black-forest-labs/FLUX.2-klein-4B" + + +def test_resolve_local_gguf_child_rejects_traversal(tmp_path): + """Round 13 P1 #2: gguf_filename must not escape the repo root.""" + from core.inference.diffusion import _resolve_local_gguf_child + + repo_root = tmp_path / "my-flux" + repo_root.mkdir() + (repo_root / "model.gguf").write_bytes(b"x") + sibling = tmp_path / "other.gguf" + sibling.write_bytes(b"y") + + assert _resolve_local_gguf_child(repo_root, "model.gguf").name == "model.gguf" + + # ``./model.gguf`` is normalised by PurePosixPath to ``model.gguf`` + # and stays inside the repo, so it is intentionally accepted. + for bad in ("../other.gguf", "", "sub/../model.gguf"): + with pytest.raises(RuntimeError): + _resolve_local_gguf_child(repo_root, bad) + with pytest.raises(RuntimeError): + _resolve_local_gguf_child(repo_root, "/etc/passwd") + + +def test_resolve_local_gguf_child_rejects_backslash(tmp_path): + """Round 13 P1 #2: a Windows-style separator inside gguf_filename + must be rejected even on POSIX so it never becomes a literal name.""" + from core.inference.diffusion import _resolve_local_gguf_child + + repo_root = tmp_path / "my-flux" + repo_root.mkdir() + (repo_root / "model.gguf").write_bytes(b"x") + + with pytest.raises(RuntimeError): + _resolve_local_gguf_child(repo_root, r"..\\other.gguf") + + +def test_load_model_accepts_relative_local_dir(monkeypatch, tmp_path): + """Round 13 P1 #2: relative directory paths (Studio exports) must + NOT be routed through hf_hub_download.""" + import core.inference.diffusion as d + + repo_root = tmp_path / "exports" / "my-flux" + repo_root.mkdir(parents = True) + gguf_file = repo_root / "model.gguf" + gguf_file.write_bytes(b"x") + + # cwd so the relative path resolves to repo_root + monkeypatch.chdir(tmp_path) + + fake_transformer = object() + fake_pipe = SimpleNamespace( + to = lambda *a, **kw: None, + enable_model_cpu_offload = lambda: None, + ) + + class _FakeQuantConfig: + def __init__(self, **_): + pass + + class _FakeTransformerCls: + from_single_file_calls: list[tuple[str, dict]] = [] + + @classmethod + def from_single_file(cls, path, **kwargs): + cls.from_single_file_calls.append((path, kwargs)) + return fake_transformer + + class _FakePipeCls: + @classmethod + def from_pretrained(cls, base, **kwargs): + return fake_pipe + + fake_diffusers = SimpleNamespace( + __version__ = "0.99", + GGUFQuantizationConfig = _FakeQuantConfig, + Flux2Transformer2DModel = _FakeTransformerCls, + Flux2KleinPipeline = _FakePipeCls, + ) + + fake_torch = SimpleNamespace( + cuda = SimpleNamespace( + is_available = lambda: False, + is_bf16_supported = lambda: False, + empty_cache = lambda: None, + ), + bfloat16 = "bf16", + float16 = "fp16", + float32 = "fp32", + backends = SimpleNamespace( + mps = SimpleNamespace(is_available = lambda: False), + ), + ) + + def _boom(**_): + raise AssertionError("hf_hub_download must not run for a local dir") + + fake_hub = SimpleNamespace(hf_hub_download = _boom) + monkeypatch.setitem(sys.modules, "huggingface_hub", fake_hub) + monkeypatch.setitem(sys.modules, "diffusers", fake_diffusers) + monkeypatch.setitem(sys.modules, "torch", fake_torch) + + backend = d.DiffusionBackend() + backend.load_model( + repo_id = "exports/my-flux", + gguf_filename = "model.gguf", + family_override = "flux.2-klein", + enable_model_cpu_offload = False, + ) + + assert _FakeTransformerCls.from_single_file_calls + resolved_path = _FakeTransformerCls.from_single_file_calls[0][0] + assert str(gguf_file.resolve()) == resolved_path + + +def test_generate_image_with_metadata_returns_active_pipeline(monkeypatch): + """Round 13 P2 #9: meta returns the resident pipeline's identity.""" + import core.inference.diffusion as d + + backend = d.DiffusionBackend() + fake_fam = d.DiffusionFamily( + name = "flux.2-klein", + pipeline_class = "Flux2KleinPipeline", + transformer_class = "Flux2KleinTransformer3DModel", + 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)) + + backend._pipe = object() + backend._repo_id = "unsloth/FLUX.2-klein-4B-GGUF" + backend._family = fake_fam + monkeypatch.setattr(backend, "_generate_image_unlocked", _fake_unlocked) + + _, meta = backend.generate_image_with_metadata(prompt = "x") + assert meta == { + "model": "unsloth/FLUX.2-klein-4B-GGUF", + "family": "flux.2-klein", + } + + +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.""" + import threading + import core.inference.diffusion as d + + backend = d.DiffusionBackend() + fake_fam = d.DiffusionFamily( + name = "flux.2-klein", + pipeline_class = "Flux2KleinPipeline", + transformer_class = "Flux2KleinTransformer3DModel", + base_repo = "black-forest-labs/FLUX.2-klein-4B", + aliases = (), + ) + + started = threading.Event() + finish = threading.Event() + + def _fake_unlocked(**kwargs): + from PIL import Image as _Image + + started.set() + # Hold long enough for the unload thread to race the metadata + # snapshot if the lock were released too early. + finish.wait(timeout = 2.0) + return _Image.new("RGB", (8, 8)) + + backend._pipe = object() + backend._repo_id = "unsloth/FLUX.2-klein-4B-GGUF" + backend._family = fake_fam + monkeypatch.setattr(backend, "_generate_image_unlocked", _fake_unlocked) + + result: list = [] + + def _gen(): + result.append(backend.generate_image_with_metadata(prompt = "x")) + + gen_thread = threading.Thread(target = _gen) + gen_thread.start() + assert started.wait(timeout = 2.0) + + def _unload(): + backend.unload_model() + + un_thread = threading.Thread(target = _unload) + un_thread.start() + # The unload must NOT have completed yet; it queues behind the + # generation's _generate_lock. + un_thread.join(timeout = 0.2) + assert un_thread.is_alive() + finish.set() + gen_thread.join(timeout = 5.0) + un_thread.join(timeout = 5.0) + + assert result + _, meta = result[0] + assert meta["model"] == "unsloth/FLUX.2-klein-4B-GGUF" + assert meta["family"] == "flux.2-klein" diff --git a/studio/backend/tests/test_diffusion_routes.py b/studio/backend/tests/test_diffusion_routes.py index fc99228a63..d3659b89ee 100644 --- a/studio/backend/tests/test_diffusion_routes.py +++ b/studio/backend/tests/test_diffusion_routes.py @@ -80,6 +80,14 @@ class _FakeBackend: "pipeline_class": "Flux2KleinPipeline" if self._loaded else None, "base_repo": "black-forest-labs/FLUX.2-klein" if self._loaded else None, "gguf_filename": None, + "active_repo_id": self._repo, + "active_base_repo": ( + "black-forest-labs/FLUX.2-klein" if self._loaded else None + ), + "active_gguf_filename": None, + "pending_repo_id": None, + "pending_base_repo": None, + "pending_gguf_filename": None, "device": "cpu", "dtype": "torch.bfloat16", "loaded_at": 0, @@ -102,6 +110,14 @@ class _FakeBackend: self.calls.append({"op": "generate", **kw}) return Image.new("RGB", (kw["width"], kw["height"]), color = (123, 45, 67)) + def generate_image_with_metadata(self, **kw): + image = self.generate_image(**kw) + meta = { + "model": self._repo, + "family": "flux.2-klein" if self._loaded else None, + } + return image, meta + @pytest.fixture def app_with_stub(monkeypatch):