ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix/adjust diffusion: round 13 P1+P2 batch for PR #5754

Browse source 
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.
This commit is contained in:
Daniel Han-Chen 2026-05-25 05:57:32 +00:00
commit ae41bfdbfd
7 changed files with 716 additions and 205 deletions
Download patch file Download diff file Expand all files Collapse all files

374
studio/backend/core/inference/diffusion.py
View file

@ -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": <repo_id>, "family": <name>})``
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),
)

71
studio/backend/core/inference/llama_cpp.py
View file

@ -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

34
studio/backend/routes/export.py
View file

@ -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.

35
studio/backend/routes/inference.py
View file

@ -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"),
)

164
studio/backend/routes/models.py
View file

@ -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:

225
studio/backend/tests/test_diffusion_backend.py
View file

@ -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"

16
studio/backend/tests/test_diffusion_routes.py
View file

@ -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):