Round 16 reviewer aggregate (logs/review_round16_aggregate.md):
P1 fixes:
- routes/models.py /delete-cached llama guard pairs loading_id with
loading_hf_variant so deleting a different cached quant (Q8_0)
while another variant (Q4_K_M) is loading is no longer blocked.
- core/inference/diffusion.py load_model now calls
_release_other_gpu_owners_for_diffusion BEFORE
_release_chat_backend_for_diffusion. The other-owners helper
RAISES on active training/export, so a route -> worker race or
direct backend caller no longer drops the user's chat model
before the diffusion load is refused.
- routes/models.py /delete-cached diffusion guard fails CLOSED
(503) on HF cache scan failure instead of silently falling
through to repo-id-only matching, which could miss a loaded
local snapshot path.
- routes/inference.py _release_llama_for and
_release_safetensors_chat_for now raise 503 on actual unload
failure (exception or False return), so new GPU workloads do
not start while the old chat process still owns VRAM.
- core/inference/diffusion.py status() now takes
include_internal=False by default and only exposes the
guard-facing active_*/pending_* paths when callers opt in. The
public /api/inference/images/status route gets the redacted
payload; routes/models.py delete guards pass
include_internal=True so they still see the raw paths.
- core/inference/diffusion.py generate_image_with_metadata routes
the response model through _display_repo_id so /images/generate
cannot echo back an absolute local path.
P2 fixes:
- routes/inference.py /images/load now maps backend "Could not
verify training/export status" to 503 instead of 409, matching
the route-level pre-check.
- core/inference/diffusion.py _release_other_gpu_owners_for_diffusion
raises "Could not verify export status" when the
is_export_active() probe itself raises, instead of silently
treating it as active export.
- core/inference/diffusion.py detect_family compares compact family
spellings (Flux2Klein) against per-token compact strings so
unsloth/Flux2Klein-GGUF matches the flux.2-klein family without
matching the embedded substring inside flux.20.
- main.py installs a RequestValidationError handler that scrubs
hf_xxxxx tokens out of the 422 response body so a rejected
``repo_id`` containing a URL-embedded HF token does not echo it
back to the browser.
Tests:
- 3 new regression cases (Flux2Klein compact alias, public status
redaction, generate_image_with_metadata redaction).
- All 75 diffusion backend + route tests pass.
Round 15 split LlamaCppBackend.load_model into a thin wrapper that
publishes _loading_model_identifier + _loading_hf_variant under
_serial_load_lock and an inner _load_model_impl_locked body that
actually launches llama-server. The pre-existing source-inspection
regression tests inspected only load_model and broke because the
flag literals and _wait_for_vram_settle call now live in the inner
method:
- tests/test_llama_cpp_no_context_shift.py
test_no_context_shift_is_in_load_model
test_flag_sits_inside_the_base_cmd_list
- tests/test_llama_cpp_wait_for_vram_settle.py
test_load_model_calls_helper_outside_lock_and_uses_last_kill_timestamp
Update both helpers to concatenate the source of load_model AND
_load_model_impl_locked so the assertions still cover the launch
path without weakening their scope to the full module.
Round 15 reviewer aggregate (logs/review_round15_aggregate.md):
P1 fixes:
- core/inference/llama_cpp.py publishes loading_model_identifier +
loading_hf_variant AFTER acquiring _serial_load_lock; previously
a queued second load could overwrite or clear the identifier
currently in flight, breaking delete-safety and GPU handoff guards.
- routes/models.py /delete-finetuned compares the pending llama
load against loading_hf_variant (new), not the stale hf_variant
from the previous loaded model. Without this, a Q4-loaded
directory loading Q8 would still accept a Q8 delete.
- core/inference/diffusion.py _release_other_gpu_owners_for_diffusion
now also raises when training is active so direct backend callers
cannot bypass the route layer's 409 guard. Mirrors the
export-active check the same helper already enforces.
- routes/models.py /delete-cached diffusion guard compares owned
diffusion paths against the HF cache root for the target repo
via _all_hf_cache_scans + _is_path_under. Without this, loading
from a local models--owner--model/snapshots/<sha> path let the
cache delete proceed while the snapshot was still mmap'd.
- models/inference.py DiffusionLoadRequest refuses URL-embedded
hf_xxxxx tokens in repo_id / base_repo at the API boundary, so
the value never reaches self._repo_id and status() can never
echo it back to other authenticated sessions.
P2 fixes:
- core/inference/diffusion.py status() routes UI-facing repo_id /
base_repo through _display_repo_id, which collapses absolute
local paths to the leaf name (delete guards still see the full
path via active_*/pending_*).
- routes/inference.py /images/load maps backend RuntimeError that
reports an export/training conflict to HTTP 409 instead of 400.
- core/inference/diffusion.py detect_family now uses token-boundary
matching so owner/flux.20-model does not collide with flux.2.
P3 fixes:
- tests/test_diffusion_routes.py drops the partial routes.inference
module from sys.modules if exec_module() raises, so the real
ImportError surfaces instead of a misleading AttributeError on
follow-up tests.
Tests:
- 5 new regression cases (display_repo_id, token-boundary family
detection, training-active raise from backend helper, embedded HF
token rejection).
- All 72 diffusion backend + route tests pass.
Round 14 reviewer aggregate (logs/review_round14_aggregate.md):
P1 fixes:
- routes/export.py /load-checkpoint now runs the active-export 409
guard BEFORE the chat / diffusion unloads, so a rejected request
no longer tears down unrelated GPU state.
- core/inference/llama_cpp.py wraps the WHOLE load_model body in a
single try/finally that publishes loading_model_identifier across
download, metadata read, VRAM settle, process spawn, and health
check. Done via a thin load_model wrapper around the existing
body (renamed _load_model_impl) to avoid reindenting hundreds of
lines.
- routes/models.py /delete-finetuned now checks
loading_model_identifier so a pending HF GGUF download cannot
have its destination directory rmtree'd before llama-server
spawns.
- core/inference/diffusion.py stores the original caller-supplied
gguf_filename (e.g. ``BF16/model.gguf``) in a new self._gguf_filename
field and exposes it as active_gguf_filename. UI-facing
gguf_filename still collapses to basename for the panel.
- routes/models.py /delete-cached llama guard now allows safe
different-variant deletes when hf_variant differs, matching the
diffusion path's variant-aware behaviour.
- core/inference/diffusion.py tracks self._cpu_offload_enabled and
forces a CPU torch.Generator when offload is on, so seeded
generation no longer crashes on CUDA hosts with the default offload
enabled.
P2 fixes:
- core/inference/diffusion.py detect_family normalises mixed
separators (``Qwen_Image-Edit-GGUF``, ``Qwen-Image_Edit-GGUF``,
``QwenImageEdit-GGUF``) so every Qwen-Image-Edit spelling is
excluded from the base Qwen-Image family.
- core/inference/diffusion.py logger.info / logger.error in
load_model run repo_id and effective_base through _redact_hf_tokens
so URL-embedded ``hf_xxxxx`` tokens never reach structured-log
sinks.
- core/inference/diffusion.py _release_other_gpu_owners_for_diffusion
now raises RuntimeError when an export job is active instead of
logging and continuing, so direct backend callers cannot bypass
the route layer's 409 guard.
- core/inference/diffusion.py full-diffusers repo / base_repo paths
expand ``~`` via _expand_existing_local_path so
``repo_id="~/models/my-flux"`` no longer falls through to the Hub.
Tests:
- 5 new regression cases (mixed Qwen-Image-Edit separators, token
redaction, status full-filename, CPU offload generator device,
staging Windows leaf already-set sanity).
- All 68 diffusion backend + route tests pass.
Round 13 follow-up: on Windows Path('/etc/passwd').is_absolute()
returns False because POSIX absolute paths read as drive-relative,
which let the traversal check fall through to resolve(strict=True)
and crash with a raw FileNotFoundError instead of the friendlier
RuntimeError. Add a PurePosixPath check + explicit leading-separator
guard and wrap the resolve() in try/except so a missing path inside
the chosen repo is reported as 'Local repo path does not contain ...'
on every OS.
Pre-existing 59 diffusion backend + route tests still pass; staging
Windows Diffusion CI was failing on this exact case.
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.
Round 12 reviewer findings.
Backend correctness (P1)
* core/inference/diffusion.py load_model: GGUF branch now
handles an absolute local directory passed as repo_id by
joining Path(repo_id) / gguf_filename directly instead of
handing the path to hf_hub_download (which raises
HFValidationError because the path is not 'namespace/repo').
Closes round 12 review #1 -- the load request advertised
'local path' support but actually only worked for Hub repo ids.
Delete guard precision (P1)
* routes/models.py /delete-finetuned + /delete-cached:
diffusion guard now consults gguf_filename from status()
and ALLOWS per-variant deletes that target a different quant
than the one the loaded pipeline is reading. Loading
'Q4_K_S' no longer blocks deleting 'Q8_0' from the same
repo / export directory (round 12 reviews #3 and #4).
Accelerator (P2)
* core/inference/diffusion.py _drain_cuda_cache: also calls
torch.mps.empty_cache() when the MPS backend is the
active accelerator. Apple Silicon swaps now actually return
held VRAM instead of leaving it pinned in the Metal
allocator (round 12 review #10).
Smart base repo (P2)
* core/inference/diffusion.py _smart_base_repo: only inspects
the LAST segment of the repo id / path for the 'base' / '9b'
tokens. A namespace like baseorg/FLUX.2-klein-4B-GGUF or
a parent directory like /home/me/.cache/base/... no
longer falsely selects the Base variant (round 12 review #9).
Round 11 reviewer findings.
Backend lifecycle (P1)
* core/inference/diffusion.py _release_other_gpu_owners_for_
diffusion: now re-checks is_export_active() locally before
calling _shutdown_subprocess. The route layer already 409s on
active exports, but defence-in-depth means direct backend
callers (tests, scripts, future routes that forget the
higher-level guard) can no longer terminate an in-flight
export and corrupt the user's partial output.
* routes/inference.py standard chat-load path: the duplicate
inline 'if exp_backend.current_checkpoint -> _shutdown_subprocess'
block was removed. _release_export_for above already handles
settled checkpoints and skips active ones; the inline block
was the round 11 #2 asymmetric fix surface.
Routing / error mapping (P2)
* routes/training.py start_training: except HTTPException:
raise was inserted before the broad except Exception:
handler so the 409 raised by _raise_if_training_active /
_raise_if_export_active reaches the client intact instead of
being swallowed into a 500.
State publishing (P2)
* core/inference/diffusion.py load_model: success path now
clears _loading + _pending_* under _lock BEFORE returning
self.status(), so the response payload reports the resident
pipeline cleanly (no stale is_loading=true / pending_*). The
finally block remains idempotent for error / early-raise paths.
* core/inference/diffusion.py status(): nulls family /
pipeline_class while a swap is in flight (pending_repo set
and != active_repo). Previously the response paired pending
model B's repo_id with model A's family, producing a
combination that never existed.
Validation
* models/inference.py: DiffusionLoadRequest.repo_id and
base_repo length caps bumped from 256 to 1024; gguf_filename
bumped from 256 to 512. The earlier caps rejected realistic
Studio export paths (deeply nested outputs / exports
directories, especially on Windows).
Dependencies
* pyproject.toml huggingfacenotorch + studio/backend/
requirements/no-torch-runtime.txt: floor gguf at >=0.10.0
to match the diffusers requirement. Unconstrained pin allowed
a resolver to install older gguf releases that raise at
single-file load time.
Round 10's training-side _raise_if_export_active call broke
existing test mocks and older ExportBackend builds that only
expose current_checkpoint -- they raised AttributeError on
exp.is_export_active() and the outer guard converted that into
a 503, causing the prior Backend CI to fail
test_inference_route_returns_400_for_invalid_gpu_ids,
test_training_route_returns_400_for_invalid_gpu_ids, and
test_training_route_forwards_embedding_learning_rate.
Both _raise_if_export_active and _release_export_for now detect
the missing method with getattr(...) and treat absence as
'no async-job tracker available' (effectively 'not active').
The 503 fail-closed path still fires when the method exists but
the call itself raises, so production backends (the
ExportOrchestrator subclass that does expose is_export_active)
keep their stronger guard.
Round 10 reviewers found the round 9 export helpers had a
destructive bug: _release_export_for treated is_export_active=True
as a shutdown condition, so any caller (training, chat, diffusion)
could terminate an in-flight export and corrupt the user's output.
Conversely _raise_if_export_active raised 409 on a settled
checkpoint, blocking idle cleanup.
Backend (P1)
* routes/inference.py: split the export-active surface in two:
_raise_if_export_active() now ONLY raises when
is_export_active() is True. A settled current_checkpoint is
treated as held GPU memory, not an active job.
_release_export_for() now ONLY shuts down when
current_checkpoint is set AND is_export_active() is False
(i.e. a previously completed checkpoint just holding memory).
An unknown / unverifiable is_export_active is treated as
'might still be active' so the helper refuses to drop.
* routes/training.py: now calls _raise_if_export_active before
_release_chat_for / _release_export_for, mirroring the chat
and diffusion paths. The previous code went straight to
_release_export_for and would kill an in-flight export.
* routes/inference.py: split _release_chat_for into
_release_llama_for and _release_safetensors_chat_for so the
GGUF chat-load path can release only the OTHER chat backend
(round 10 review #4: the previous inline 'if active_model_name'
check skipped loading_models and let an in-flight safetensors
load race the new GGUF allocation).
* routes/inference.py: _raise_if_export_active now fails CLOSED
(503) when is_export_active() raises, not only when
get_export_backend() raises. Round 10 review #7.
Dependencies (P1)
* pyproject.toml huggingfacenotorch extra: pin gguf. The
Studio Images default curated picker is GGUF-only and
diffusers.GGUFQuantizationConfig + from_single_file require
the standalone gguf package at runtime; missing it would 500
on the first /api/inference/images/load with
'gguf>=0.10.0 is required'.
Round 9 reviewer flagged a pile of handoff asymmetries: every
GPU-owning lifecycle change (training, export, chat, images) needed
its own bespoke unload sequence and they had drifted out of sync.
Some skipped llama-server is_active; some missed safetensors
loading_models; export and training did not check is_export_active.
Backend handoff (P1)
* routes/inference.py: new _release_chat_for / _release_export_for
helpers. Both treat llama-server as held when is_loaded OR
is_active, safetensors as held when active_model_name OR
loading_models is non-empty, and export as held when
current_checkpoint OR is_export_active. Both helpers run their
unloads in worker threads so async routes do not block the
event loop.
* routes/training.py: replaces its bespoke inline llama / safe /
export unload sequence with await _release_chat_for / _release_
export_for.
* routes/export.py: same swap for the chat unload chain (export
still does NOT call _release_export_for on itself).
* routes/inference.py GGUF + standard chat-load paths: now use
_release_export_for to drop a settled export, and the standard
path's llama unload now also handles is_active=True (round 9
review #8).
Backend reject-on-active export (P1 #5)
* routes/inference.py: new _raise_if_export_active. Symmetric
with _raise_if_training_active: a long-running export is
refused with HTTP 409 instead of being silently killed when
/images/load or /load arrives. Diffusion / images load and
both chat-load paths call it.
* core/inference/diffusion.py _release_other_gpu_owners_for_
diffusion: no longer tears down an in-flight export job. Only
drops a SETTLED export checkpoint (current_checkpoint
populated, is_export_active False). Round 9 review #5 -- the
previous behavior could terminate an in-flight export and
leave a partial output artifact.
Token leak via logger.exception (P1 #6)
* core/inference/diffusion.py: load-failure logging now uses
logger.error(..., exc_msg) with the already-scrubbed string
and exc_info=False. logger.exception() with the raw Exception
would expose any hf_... token that diffusers / huggingface_hub
embedded in the message or traceback locals, defeating the
earlier in-flight scrub.
Dependency pinning (P1 #11)
* pyproject.toml: huggingfacenotorch optional extra now pins
diffusers>=0.37.0. Previously the floor was only set in
studio/backend/requirements/no-torch-runtime.txt, so a normal
pip install would resolve diffusers 0.36.0 (no
Flux2KleinPipeline) and the default curated FLUX.2 klein
Images model would fail at runtime.
Cache-delete exact match (P1 #14)
* routes/models.py /delete-cached: llama.cpp and safetensors
guards now match on exact repo-id (case-insensitive) instead
of prefix. Diffusion guard already does this; the chat guards
were the remaining surface where loading org/model-v2
blocked deleting org/model.
Round 8 reviewer surfaced event-loop stalls (blocking unload from
async routes), incomplete VRAM handoff coverage (is_active /
loading_models / is_export_active not checked), token leaks via
exception messages, /v1 exposure, and several fail-open paths.
Async / event-loop
* routes/inference.py /images/unload, GGUF chat-load handoff,
safetensors chat-load handoff: blocking DiffusionBackend.unload
pushed onto asyncio.to_thread. unload takes _load_lock +
_generate_lock and can block for the full duration of an
in-flight load / generation, which was freezing the FastAPI
worker, SSE stream, and hardware poller for minutes.
* routes/export.py + routes/training.py: same to_thread wrap on
diffusion unload during checkpoint / training start.
GPU-owner handoff completeness
* core/inference/diffusion.py _release_chat_backend_for_diffusion:
llama-server now also unloaded when is_active=True (mid-download
/ startup), not only when is_loaded; flushed in-flight
safetensors loads from loading_models too.
* core/inference/diffusion.py _release_other_gpu_owners_for_
diffusion: export shutdown now also fires when
is_export_active() returns True (checkpoint not yet assigned).
Security / scrubbing
* core/inference/diffusion.py: load failure paths now scrub
hf_token from both _last_error AND the raised RuntimeError
message (the previous scrub only cleared frame locals).
Falls back to a regex strip of hf_[A-Za-z0-9]{20,} to
catch tokens that came in via huggingface_hub default caching.
* routes/inference.py: image lifecycle endpoints moved from
router to studio_router so they no longer answer under
the /v1 OpenAI-compat prefix. Studio-only side effects
(download multi-GB GGUFs, unload chat, etc.) should not be
reachable via an OpenAI-compat client.
* models/inference.py: control-char validator now also rejects
tab. Some log sinks split fields on tab; allowing it left a
log-injection surface.
Fail-closed delete guards
* routes/models.py /delete-cached: llama.cpp and safetensors
branches now fail closed with 503 when their status check
raises (matches the diffusion-side guard added earlier).
* routes/export.py: split the try/except around the training
backend so import failure falls back to 'skip' (no
core.training in this build) while a runtime failure of
get_training_backend()/is_training_active() fails closed.
* routes/models.py /delete-finetuned: diffusion guard now also
compares against relative path candidates (Path.resolve() works
on relative input). Previously a load with a relative repo_id
bypassed the guard.
CUDA cleanup ordering
* core/inference/diffusion.py: split _release() (drops local +
gc.collect) from _drain_cuda_cache() (torch.cuda.empty_cache).
Callers now drain AFTER nulling every reference so the
allocator actually reclaims the freed slabs (previously
empty_cache ran while caller still held a local, which left
the cache pinned).
Generate response (P2 #16)
* routes/inference.py: response uses status()['active_repo_id']
instead of the UI-facing repo_id, so a queued /images/load
promoting a pending model cannot mislabel the just-rendered
image with the new model's identity.
Test wiring
* tests/test_diffusion_routes.py: mount inf.studio_router on the
test app so /images/* routes are reachable now that they live
on the Studio-only router.
Round 7 reviewer surfaced a handful of swap-window races, fail-open
guards, and seed precision mismatches. This commit closes them.
Lifecycle / state (P1)
* core/inference/diffusion.py: status() now emits active_repo_id,
active_base_repo, pending_repo_id, pending_base_repo, and
pending_gguf_filename alongside the existing UI-facing fields.
During a swap (model A loaded, model B loading) the previous
coalesced 'repo_id or pending_repo_id' hid the loading target
from delete guards. Splitting the fields lets guards block
deletion of either repo currently owned by the backend.
* core/inference/diffusion.py: generate_image() now takes
_generate_lock BEFORE snapshotting _pipe / _device. Snapshotting
outside the lock let a concurrent unload/load clear or replace
the backend between the snapshot and the forward, so the freed
or swapped pipeline would still run.
Symmetric handoffs (P1)
* routes/export.py: training-active check now runs BEFORE the
chat / inference / diffusion unload helpers, so a 409 does not
leave the user's chat session torn down for nothing. Also
explicitly fails CLOSED with 503 when is_training_active()
raises.
* routes/inference.py: _raise_if_training_active now fails closed
with 503 when the training backend is importable but its status
check raises. The previous best-effort log-and-continue could
let chat / diffusion loads collide with unverifiable training.
Delete guards (P1)
* routes/models.py /delete-cached: chat guard now also blocks
when llama-server is_active (i.e. mid-download) and when the
inference backend's loading_models set contains the target.
Round 7 review #7 flagged that the PR's diffusion-side loading
guard had no chat-side parallel, so deleting a chat repo while
it was downloading could still race the cache.
* routes/models.py /delete-cached: diffusion guard iterates the
new active_* + pending_* status fields so a delete during a
swap is refused on either repo.
* routes/models.py /delete-finetuned: same active_+ pending
handling, plus the guard now also refuses deletes of a parent
directory that contains the loaded pipeline (round 7 review #6:
rm -rf /exports/flux-model/ could unlink model_index.json that
the live pipeline is reading via mmap).
Seed precision (P2)
* models/inference.py + routes/inference.py: DiffusionGenerate-
Response now carries seed_str alongside the existing numeric
seed. Seeds above Number.MAX_SAFE_INTEGER are rounded by
JSON.parse in the browser; seed_str ships full decimal
precision for display and reproduction.
* frontend/api.ts: DiffusionGenerateResponse types seed_str;
images-page.tsx prefers seed_str over seed in the figure
caption so the displayed value reproduces the image.
* frontend/api.ts: stringifyWithBigInt no longer regex-replaces
sentinel strings over the full JSON output. It pulls the seed
BigInt out, JSON-serialises the remaining payload, and splices
the seed's decimal digits into the resulting object literal at
the known position. Avoids the round 7 #10 case where a
user-supplied prompt equal to '__bigint__:123' was rewritten
into a JSON integer and rejected as a non-string prompt.
Custom HF repo (P2)
* frontend/images-page.tsx: custom panel now exposes a 'Base
diffusers repo' input that maps to DiffusionLoadRequest.
base_repo. Required when a private / mirrored GGUF needs a
non-default base (e.g. a 9B Klein transformer would otherwise
fall back to the 4B base default).
Round 6 reviewers identified several races between load / unload /
generate and several fail-open delete guards. This commit closes
them by widening the lock scope, publishing the pending load
target through status(), and switching delete guards to
fail-closed.
Lifecycle (P1)
* core/inference/diffusion.py: load_model now also takes
_generate_lock. Previous behavior released and reallocated the
pipeline while a generation forward was still iterating
denoising steps, corrupting scheduler state and stacking VRAM.
The forward only briefly touches _lock, so taking it on the
load path does not introduce a deadlock.
* core/inference/diffusion.py: unload_model now also takes
_generate_lock. Without it, /images/unload returned
is_loaded=False while a slow forward was still running, which
let chat / training / export handoffs allocate VRAM on top of
the still-resident pipeline.
* core/inference/diffusion.py: previous pipeline release now
happens BEFORE from_single_file / from_pretrained. Switching
FLUX.2 klein 4B -> 9B on a 16-24 GB GPU was failing because
the new transformer allocation overlapped the old pipe's
residency.
* core/inference/diffusion.py: failed pipeline from_pretrained
now explicitly releases the just-loaded transformer; previously
its weights stayed pinned to GPU until GC and made the next
load more likely to OOM.
Pending-target / delete guards (P1)
* core/inference/diffusion.py: load_model now publishes
_pending_repo_id / _pending_base_repo / _pending_gguf_filename
under _lock at the start of the call (and refreshes
_pending_base_repo when the smart-base / repo defaults resolve).
status() exposes those as 'repo_id' / 'base_repo' /
'gguf_filename' during is_loading=True so delete guards can see
the target before _repo_id is set on success.
* routes/models.py /delete-cached + /delete-finetuned: diffusion
status check now fails CLOSED (HTTP 503) when status() raises.
Both guards previously logged and continued, which could let a
delete proceed against a repo whose status was unverifiable.
* routes/models.py: is_loading is also blocked on both guards
so a mid-download / mid-from_pretrained rmtree is refused.
Symmetric handoffs (P1)
* routes/export.py: /load-checkpoint now refuses with HTTP 409
when training is active instead of calling stop_training().
Chat and /images/load did the same after round 5; export was
the remaining asymmetry that would silently kill a long
training run.
* routes/training.py, routes/inference.py (GGUF and standard
chat), routes/export.py: diffusion handoff now treats
is_loading as is_loaded. The diffusion backend's unload waits
on _load_lock + _generate_lock so an in-flight load completes
first.
Requirements (P1)
* requirements/studio.txt: pin python-multipart explicitly. The
Studio routes package's eager router imports include
routes/datasets.py whose FastAPI UploadFile/File validation
crashes with RuntimeError without it in fresh test envs.
Frontend (P2)
* features/images/api.ts + images-page.tsx: seed handling now
accepts the full [-2^63, 2^64 - 1] range via BigInt. The
previous safe-integer cap rejected valid uint64 seeds the
backend accepts. A small stringify helper emits BigInts as JSON
integers without touching the rest of the payload.
Tests
* test_diffusion_routes.py: load routes/inference.py via
importlib.spec_from_file_location to avoid triggering
routes/__init__.py (which would pull in training / datasets /
data_recipe imports unrelated to diffusion tests).
* test_diffusion_backend.py: status() during is_loading shows
pending repo + base; unload waits for in-flight generation.
Round 5 reviewer findings, mostly symmetric-lifecycle and input
validation gaps the earlier rounds left open.
Backend lifecycle (P1)
* routes/training.py: training start now also unloads the GGUF
llama-server subprocess; was previously only unloading the
safetensors backend, so starting training while a GGUF chat
model was loaded kept the subprocess pinned to VRAM.
* routes/inference.py: new _raise_if_training_active helper. Both
GGUF and standard chat loads, plus /api/inference/images/load,
now refuse with HTTP 409 when training is active instead of
silently stopping training to free VRAM.
* core/inference/diffusion.py: _release_other_gpu_owners_for_
diffusion no longer stops active training. The route layer
refuses the request first, so reaching the helper with training
live would only happen from programmatic backend calls; better
to surface OOM than terminate a long training run.
* core/inference/diffusion.py: BF16 dtype is now gated on
torch.cuda.is_bf16_supported. Pascal/Turing GPUs report
is_available()=True but lack BF16 ALUs; FLUX kernels then fail
inside from_pretrained. Falls back to FP16 instead of refusing.
* core/inference/diffusion.py: GGUF transformer allocation and
pipeline allocation now run AFTER releasing chat/export GPU
owners; previously from_single_file ran first and could OOM
before the intended VRAM handoff happened.
* routes/models.py: /delete-cached now also blocks delete when
diffusion is_loading=True (not just is_loaded); concurrent
delete during hf_hub_download / from_single_file would have
raced the rmtree.
* routes/models.py: /delete-finetuned now also checks the
diffusion backend before unlinking a Studio outputs/exports
path. A user who exported a FLUX LoRA locally and loaded it via
/images/load could previously rmtree the directory the
diffusion backend was reading from.
Backend correctness / safety (P2)
* core/inference/diffusion.py: _FAMILY_EXCLUDE for qwen-image now
also covers qwen_image_edit / qwenimageedit underscore spellings
so '...qwen_image_edit-GGUF' no longer misdetects as Qwen-Image.
* core/inference/diffusion.py: detect_family now scans
_FULL_REPO_FAMILIES in addition to _FAMILIES, so SDXL repos
(stabilityai/stable-diffusion-xl-base-1.0) are auto-detected
instead of failing with 'Could not infer a diffusion family'.
* core/inference/diffusion.py: generate_image now uses a separate
_generate_lock for the pipeline forward instead of holding
_lock for the whole call. status() polls and concurrent unload
requests no longer block for the full minutes-long generation.
* routes/models.py: diffusion delete guard now uses exact repo-id
match instead of prefix match; previously loading 'org/model-v2'
would block deleting unrelated cached 'org/model'.
* models/inference.py: DiffusionLoadRequest now rejects ASCII
control characters in repo_id / gguf_filename / base_repo /
family via field_validator (closes log-injection surface from
authenticated callers). Also caps lengths at 256 chars.
* models/inference.py: DiffusionGenerateRequest seed is now
bounded to the int64/uint64 range; previously a huge seed
(e.g. 2**100) passed Pydantic then crashed inside
torch.Generator.manual_seed with 'Overflow when unpacking long
long'.
Frontend (P2)
* features/images/images-page.tsx: Custom HF repo panel now
exposes a Pipeline family override dropdown; previously the
backend supported it via DiffusionLoadRequest.family but the UI
had no way to send it, so custom repos whose names did not
contain a hard-coded substring failed to load.
* features/images/images-page.tsx: handleLoad now re-fetches
status on error. The backend clears its old pipeline before
allocating the replacement; a failed swap previously left the
UI showing 'Loaded:' with Generate enabled until manual
refresh.
Tests (10 new)
* underscore qwen-image-edit exclusion + SDXL full-repo detection
* BF16 fallback when is_bf16_supported() returns False
* status() does not block while generate_image holds _generate_lock
* route layer rejects control chars in repo_id
* route layer rejects 2**100 seeds (uint64-max boundary accepted)
* route layer happy-path with negative-prompt true_cfg_scale
forwarding (Qwen/Flux) and skip-when-no-neg (distilled CFG)
The Studio backend's no-torch-runtime.txt is installed via
pip --no-deps so the diffusion stack's transitive imports must
be pinned explicitly. huggingface_hub's blob downloader (used by
diffusers.GGUFQuantizationConfig and by every from_single_file
call) imports requests + urllib3 + charset_normalizer at module
load time; a fresh --no-deps install would 500 on the first
/api/inference/images/load with PackageNotFoundError: 'requests'.
Adds requests, urllib3, and charset_normalizer to the
transitive-deps block next to the existing httpx chain.
QwenImagePipeline and FluxPipeline treat guidance_scale as the
distilled CFG factor and expose true_cfg_scale as the real
classifier-free guidance knob. Negative prompts only steer the
output when true_cfg_scale > 1, so forwarding only guidance_scale
left Qwen-Image on the default true_cfg_scale=4.0 and the user's
slider value silently ineffective for negative prompts.
When the loaded pipeline accepts both negative_prompt and
true_cfg_scale and the caller supplies a non-empty negative
prompt, forward guidance_scale through both kwargs so the
negative prompt actually steers generation. When no negative
prompt is supplied, true_cfg_scale is left at the model default
to avoid switching distilled CFG models into real-CFG mode (which
would double inference cost and degrade quality).
Adds two regression tests covering the forward-when-negative and
skip-when-no-negative paths.
- DiffusionBackend.status() now takes _lock so frontend polling
cannot observe a torn snapshot mid-swap.
- Scrub hf_token / pipe_kwargs / single_file_kwargs from frame
locals before logger.exception() so rich tracebacks and structlog
formatters that render locals do not leak hf_... tokens into logs.
- routes/models.py delete_cached_repo: refuse to delete the cache
underlying a currently-loaded diffusion pipeline (both the GGUF
repo and the matching diffusers base_repo). Symmetric with the
existing chat-load + GGUF guard.
- Frontend seed validation: reject non-integer and out-of-safe-
integer-range inputs instead of silently rounding via Number(),
which would otherwise send a different seed than what the user
typed.
When a swap load fails after the previous pipeline is released,
status() previously reported is_loaded=false on top of the OLD
repo/family/base_repo metadata, which the frontend then rendered
as a misleading 'still loaded: X' label. Clear all metadata
atomically with the pipe drop so a failed swap reports a clean
empty status plus last_error. Add regression test.
- detect_family adds _FAMILY_EXCLUDE so 'stable-diffusion-3.5' no
longer matches the SD3 Medium family and 'qwen-image-edit' no
longer matches Qwen-Image. Both were misleading silent loads.
- from_single_file now forwards config=<effective_base>,
subfolder='transformer', and the HF token. Diffusers-format GGUFs
(FLUX.2 klein, Qwen-Image, SD3) need the matching base config or
the transformer load picks the wrong shapes; gated GGUFs need the
token both for download and config read.
- Move _release_chat_backend_for_diffusion + new
_release_other_gpu_owners_for_diffusion to AFTER the GGUF download
and pipeline class lookup so a typo or transient Hub error does
not kill the user's currently-loaded chat model. Peak VRAM still
stays at one model's worth because the releases run right before
from_pretrained.
- _release_other_gpu_owners_for_diffusion: shut down the export
subprocess and any active training subprocess before a diffusion
load. Symmetric with the export load path.
- routes/training.py: unload diffusion before starting training so
the new subprocess does not race FLUX/Qwen for VRAM.
- routes/export.py: also unload the GGUF llama-server before export
load (the existing inference-backend unload only covered the
safetensors path).
_release_chat_backend_for_diffusion was importing
get_inference_backend from core.inference.inference (the in-subprocess
class) and calling unload_model() without the required model_name
argument. The TypeError was swallowed and the active chat model
stayed resident, defeating the chat-to-diffusion lifecycle handoff.
Switch to the orchestrator's accessor at core.inference and pass
active_model_name through, mirroring the GGUF chat-load path. Add a
regression test that stubs both backends and verifies unload_model
is called with the active model name.
- routes/export.py load_checkpoint now unloads the diffusion
pipeline alongside the existing inference + training unloads, so
an export load after Images does not OOM the export subprocess.
- Remove the 'sd3.5' alias from the stable-diffusion-3 family.
SD3.5 needs its own family + base_repo (and its own smoke test);
pairing it with the SD3 Medium base produced a misleading load.
- _smart_base_repo: pick 9B base for unsloth/FLUX.2-klein-9B-GGUF
and -base- variants per the repo id, instead of always falling
back to the 4B family default.
- pipe_kwargs use_safetensors=True so diffusers refuses pickle .bin
weights at load time (defends against compromised base_repo).
- Release the previous pipeline BEFORE allocating the new one so
peak VRAM stays at one model's worth instead of two on swap.
- Reject empty gguf_filename when repo_id ends with -GGUF; the prior
behavior tried from_pretrained on a GGUF-only repo and 500'd deep
in diffusers with a confusing model-index error.
- Status returns gguf_filename (basename) instead of gguf_path so
the local cache path / username does not leak to authenticated
Studio sessions.
- requirements/no-torch-runtime.txt: pin diffusers>=0.37.0 so older
installs cannot resolve a version without Flux2KleinPipeline.
- Frontend curated distilled klein entries now point at the
matching non-base diffusers repos (FLUX.2-klein-4B / -9B) per
the published model cards. Update api.ts to mirror the renamed
status field.
_release_chat_backend_for_diffusion now unloads both the GGUF
chat backend (llama-server) and the safetensors / HF chat backend
(get_inference_backend) before a diffusion load. Mirror the
behaviour on the chat-load side: both the Unsloth/transformers
load path and the GGUF load path now unload the diffusion pipeline
before claiming GPU memory. Closes the OOM-on-swap path flagged
by reviewers in both directions.
- unload_model now takes _load_lock so it cannot race with an in-flight
load_model and have the load thread overwrite cleared state after
unload returned is_loaded=false.
- Move stable-diffusion-xl out of _FAMILIES into _FULL_REPO_FAMILIES.
SDXL uses a UNet (no transformer GGUF path is wired); listing it in
the GGUF families panel was misleading. SDXL full-repo loads still
work via family_override='stable-diffusion-xl'.
- Result gallery now uses h-auto + object-contain so portrait /
landscape outputs render at their true aspect ratio instead of
being cropped into a square thumbnail.
Backend
- Fix FLUX.2 klein family default base_repo: black-forest-labs/FLUX.2-klein
does not exist on the Hub. Point at the Apache 2.0 4B Base instead so
the from_pretrained call works out of the box for ungated users.
- Serialise concurrent load_model calls with a dedicated _load_lock so
two /images/load requests cannot both reach pipeline_cls.from_pretrained
at the same time (would double-spend VRAM and corrupt _pipe).
- When the caller passes a full diffusers repo (no gguf_filename),
use repo_id directly instead of silently substituting the family
default. Closes the load-the-wrong-model regression flagged by review.
- Drop negative_prompt from the pipeline call when the loaded pipeline
does not accept it (FLUX.2 / FLUX.2 klein). Inspect __call__ via
inspect.signature so we do not maintain a manual class list.
- Best-effort unload the chat backend (llama-server) before a diffusion
load so a 24 GB consumer GPU can swap between chat and diffusion
without manual unload steps.
Frontend
- Replace the four curated entries with the actual filenames published
on the Hub (lowercase flux-2-klein-Nb-Q4_K_S.gguf and flux2-dev*).
- Add an explicit base_repo per curated entry so the backend never
falls back to the family default for the curated picker.
- Add the Apache 2.0 FLUX.2 klein base 4B entry so first-time users
have an ungated, no-token-required default.
- Hide the negative prompt field for FLUX.2 / FLUX.2 klein and show a
small explanatory note instead.
Tests
- Add 6 new backend tests: base_repo override, full-repo (no GGUF)
no-substitution, concurrent serialise race, signature-based kwarg
filter, negative_prompt strip on FLUX.2, negative_prompt preserved
on supporting pipelines. 33 tests passing.
SectionCard requires an icon prop. Pass GpuIcon, PaintBrush02Icon,
and SparklesIcon for the three sections so tsc -b stops failing on
TS2741 'Property icon is missing'.
Backend
- core/inference/diffusion.py: DiffusionBackend singleton that loads
diffusion GGUFs from Hugging Face via diffusers.GGUFQuantizationConfig
and runs them on the active CUDA / MPS / CPU device. Supports FLUX.2,
FLUX.2 klein, FLUX.1, Qwen-Image, Stable Diffusion 3, and SDXL.
- routes/inference.py: POST /api/inference/images/load,
POST /api/inference/images/generate, POST /api/inference/images/unload,
GET /api/inference/images/status mirroring the llama-server lifecycle.
- models/inference.py: DiffusionLoadRequest, DiffusionGenerateRequest,
DiffusionGenerateResponse pydantic schemas with prompt / step / size
validation up front so callers get clear 422s rather than VAE crashes.
- requirements/no-torch-runtime.txt: pin gguf alongside the existing
diffusers entry so GGUFQuantizationConfig works out of the box.
- tests/test_diffusion_backend.py + tests/test_diffusion_routes.py:
27 unit tests covering family detection, validation, lifecycle, and
the full FastAPI round trip with the backend stubbed. No torch /
diffusers / GPU required to run.
Frontend
- features/images/: standalone images-page.tsx with curated model picker
(FLUX.2 klein 4B / 9B, FLUX.2 dev, FLUX.1 dev), HF token field,
prompt + negative prompt, resolution presets, steps + guidance
sliders, seed input, and a result gallery that renders base64 PNGs
inline.
- app/routes/images.tsx: lazy /images route wired into router.tsx.
- components/app-sidebar.tsx: PaintBrush02Icon nav item between
Recipes and Export, hidden in chat-only mode.
* Studio: strip orphan tool_call XML from streamed visible content
The speculative-buffer state machine in
`studio/backend/core/inference/llama_cpp.py` can slice a tool_call XML
block between the silent DRAINING path and the user-visible
content_accum, depending on when in the model's emission the BUFFERING
-> STREAMING -> DRAINING transitions fire. Three leak shapes were
observed in a 2026-05-22 sweep of 900 Qwen3.5 / Qwen3.6 GGUF runs:
Pre-fix XML leak rate: 20/900 (2.22%), concentrated 6.7% on the
larger Q8 / MTP configs:
Qwen3.6-35B-A3B Q8_0 4/60 (6.7%)
Qwen3.6-35B-A3B-MTP Q4 4/60 (6.7%)
Qwen3.5-35B-A3B Q8_0 3/60 (5.0%)
Qwen3.6-27B Q8_0 3/60 (5.0%)
The existing `_TOOL_XML_RE` only matched well-formed
`<tool_call>...</tool_call>` and `<function=...></function>` pairs, so
unterminated openings (close was DRAINED) and orphan closes (opening
was DRAINED) survived the strip and reached the user.
Fix relaxes the regex to also strip:
1. Orphan opening up to end-of-string: `(?:</tool_call>|\Z)`
2. Orphan closing tag: bare `</tool_call>` / `</function>`
Verified on the full sweep: 20/900 -> 0/900 (100% of detected leaks
eliminated). 16 unit tests in `test_tool_xml_strip.py` pin all three
leak shapes plus the well-formed cases, plus parametrised checks on
the 5 actual real-world leak samples from the sweep data.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: strip tail-only </parameter> orphan + tighten regex
The 2026-05-22 gdpval sweep surfaced a 4th XML-leak shape not caught
by the earlier regex: a bare `</parameter>\n\n` at end-of-buffer (7
of 192 trials, all Qwen3.5-27B + a few Qwen3.6-27B). The model emits
the full `<tool_call><function=...><parameter=...>...content...
</parameter></function></tool_call>` envelope, the speculative buffer
DRAINS the opening tags as intended, but EOS (max_tokens cutoff)
truncates the outer `</function></tool_call>` close, leaving just
`</parameter>` as the visible tail.
We strip this ONLY when end-anchored (`\s*\Z`) so legitimate
mid-text uses (user code samples, documentation discussing the
Qwen tool-call XML shape) survive. Verified on the 192-trial
gdpval corpus: before=7, after=0.
While at it, fold the five top-level alternations into three by
sharing tag-name and prefix subgroups:
<tool_call>... + <function=\w+>... + --> <(?:tool_call|function=\w+)>...
</tool_call> | </function> --> </(?:tool_call|function)>
Semantically identical (verified by replay over the 192-trial
corpus + adversarial inputs, 0 diffs) and 1.34x faster on real
workloads. Backtracking-safety pinned by two new perf guards
(256KB '<' spam, 1000x orphan opens).
Tests: 16 -> 28 (6 new functional + 4 well-formed-vs-orphan +
2 perf guards).
* Tighten comments in XML-strip regex and tests
Code says what it does; comments were repeating it. Strip the verbose
explanations down to the WHY-only bits (engine quirk, tail-anchor
rationale, real-world source of each test sample). No code changes.
inference.py: 21 -> 12 lines around _TOOL_XML_RE
test_tool_xml_strip.py: 343 -> 259 lines (-84)
Tests: 28/28 still pass.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
In full FT, AdamW weight decay shrinks the parameter directly so the
implicit prior is W -> 0. In LoRA the trained parameters are A and B
while the effective weight is W = W_init + (alpha/r) * B @ A; decaying
A and B separately drives BA -> 0, hence W -> W_init rather than 0.
The previous default of 0.01 inherited from full-FT recipes adds a
measurable pull on the merged adapter back toward the base model over
a few thousand steps. 0.001 keeps a small Frobenius-norm prior on
||A||^2 + ||B||^2 for numerical stability without meaningfully biasing
the merged weight toward init, and aligns with the value used across
the unsloth notebook templates.
