img2img and inpaint take their output size from the uploaded image and only
snap it to a multiple of 16, so an ordinary phone photo (up to the 4096/side
decode cap, 4x the txt2img 2048 ceiling and ~16x the area) drove an OOM-scale
latent and an opaque 500 on a normal card, while txt2img, upscale, edit, and
FLUX.2-klein inpaint are all already megapixel-bounded. Clamp the init longest
side to 2048 (the txt2img ceiling) before deriving width/height; edit is exempt
since its pipeline resizes to ~1MP internally.
_cast_nvfp4 quantized every nn.Linear with no filter, unlike the int8 and fp8
torchao text-encoder modes which exclude the VLM vision tower / lm_head / T5 wo.
On qwen-image / qwen-image-edit that 4-bit quantized the Qwen2.5-VL image tower,
degrading the edit/image conditioning the sibling schemes protect. Apply the
same make_filter_fn exclusion (require_bf16, mirroring _cast_fp8_dynamic).
union_control_mode() only matched the short catalog id, so a client naming a
curated union ControlNet by its bare HF repo id (the form resolve_controlnet
documents and accepts) got None and the generate path omitted control_mode,
which makes FluxControlNetModel.forward raise controlnet_mode cannot be None.
Fall back to matching repo_id against the curated entries so the bare repo id
resolves to the same union mode; non-union bare repos still return None.
The images page also wired Reapply for a resident single_file model with no
checkpoint filename (status carries none), but the backend rejects a
single_file/gguf load without a filename, so clicking Reapply 400'd. Narrow the
resident-Reapply wiring to pipeline (the one kind that needs no filename),
matching the existing GGUF handling, so single_file stays a no-op instead of
erroring.
Resolve the app-sidebar.tsx conflict: main refactored the chat-export dropdown to a
format-based CHAT_EXPORT_OPTIONS + dynamic-import exportConversationByFormat dispatcher,
which the merged body already uses. Keep main's dispatcher and drop the branch's static
export imports; keep TestTubeOutlineIcon imported from the shared @/lib/hugeicons-derived
module (also used by images-page) rather than main's duplicate inline definition. The
branch's Images and Video nav items are preserved.
The SDXL trainer drew min(train_batch_size, len(pairs)) indices, so a dataset with
fewer images than the batch trained at a smaller effective batch than configured
while the scheduler and samples-per-second still assumed the full batch. The shared
PermutationBatchSampler already refills across permutation cycles to return exactly k
indices, and the DiT trainer calls it with the full batch size, so drop the clamp and
pass train_batch_size through for parity and to honor the configured batch.
* Speed up Studio startup path
* Studio: recheck managed binary executability on preflight cache hit and ignore stale unauthenticated platform fetches
Preflight: a matching capability cache fingerprint no longer skips the
runnability check when the managed binary's executable bit was cleared
(size and mtime unchanged, since chmod bumps ctime not mtime). The cache
fast path now confirms the binary is still executable, otherwise it falls
back to the CLI help probe so preflight reports Stale and can repair,
instead of returning Ready and failing later at backend start. Adds a
regression test.
Frontend: now that first render is no longer gated on fetchDeviceType,
the initial unauthenticated health call can resolve after an
authenticated platform fetch. Guard the store so a late unauthenticated
or failed non-forced response cannot overwrite an already authoritative
device type, tunnel URL, or secure flag. Forced refreshes and the first
unauthenticated load are unaffected.
* Studio: use access(X_OK) for the preflight cache executability guard
A mode bitmask treats any execute bit as launchable, but the executable
bits can be set only for another owner or group, or be denied by an ACL,
so the current user could still hit PermissionDenied at launch and the
cached fast path would wrongly return Ready. access(X_OK) checks real
executability for the calling user, so an ownership or permission change
correctly falls back to the CLI help probe and the Stale repair path.
* Studio: ignore any stale non-forced platform fetch once authoritative
Extend the platform store guard so a non-forced health response never
overwrites an already authoritative result, not only unauthenticated
ones. With a saved token the post-render non-forced request can be
authenticated but older than a later forced refresh that already picked
up the tunnel URL and secure flag; if that earlier request resolves last
it would null those fields. Now any non-forced response is dropped once
the store holds a server-reported platform. Forced refreshes and the
first authoritative write are unaffected.
* Studio: run the managed CLI help probe before trusting the preflight cache
Restore running the managed CLI help probe before returning Ready from
the desktop capability cache, so a managed install whose venv interpreter
or a runtime dependency is broken (while path, size, mtime, and markers
are unchanged) is reported Stale for repair rather than proceeding to a
backend start that cannot spawn. The capability cache still skips the
heavier desktop-capabilities probe on a hit, so a warm cache runs one
probe instead of two. Removes the executable-access shortcut, which the
help probe now subsumes.
---------
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
A scanned On-Device model's id is its full filesystem path, and the family-token
matcher treats any path segment as a hint, so a folder under a family-named
parent (e.g. /models/qwen-image/misc) tagged an unrelated single-file as
text-to-image. That surfaced it in the Images picker, and a load would evict the
GPU owner before from_single_file failed on the unrelated weights. Match on
Path(model.id).name so only the leaf name counts, for both the image and video
task-tag loops (the video loop is gated on the same detection).
The native sd.cpp one-shot engine re-reads its companion VAE and text-encoder
files from the HF cache on every generation, but the delete-cached guard only
compared the loaded main diffusion repo_id. Deleting a companion repo such as
comfyanonymous/flux_text_encoders while a FLUX GGUF was loaded therefore
succeeded and broke the next generation. Add a loaded_repo_ids() accessor to the
native backend (mirroring loading_repo_ids(), reconstructed from the committed
family's VAE + text-encoder repos) and have the guard refuse those companions
too. Server mode and the diffusers engine hold companions in-process/VRAM, so
they are unaffected.
Video: a local FILE picked for a gguf/single_file load is handed straight to the
loader (the resolver returns the file itself, ignoring gguf_filename), so its own
suffix must match the kind. A .gguf picked as single_file (or a .safetensors picked
as gguf) slipped past the gguf_filename checks, evicting the resident GPU owner
before failing in from_single_file. Reject the mismatch in validate, before the handoff.
Engine router: publish the newly selected engine only after the old one finishes
unloading. The arbiter's diffusion evictor unloads the active engine, so flipping the
active name to the new (empty) engine first let a concurrent chat/video acquire evict
that empty engine and take the GPU while the old model was still freeing VRAM.
The two candidate-resolution tests exercised the real cache-disk gate, so a
small CI disk (< transient + 10 GiB free) dropped the candidate and the tests
failed non-deterministically across runners. Neutralize the disk probe in the
shared selector helper; the two disk-gate tests re-patch it to cover the gate.
Two evict/OOM fixes on the diffusion load paths:
- The video load moved a pipeline onto the GPU (apply_memory_plan) and
committed it while holding no lock, so an unload / GPU-arbiter eviction --
which bumps the load token and then barriers on _generate_lock before
freeing -- could hand VIDEO to chat/images and let the new owner allocate
concurrently with the in-flight placement, OOMing. Hold _generate_lock
across placement + the locked commit, mirroring the image backend, so an
evicting owner waits until this worker's placement is torn down or
committed. Lock order stays _generate_lock -> _lock (unload takes _lock
then releases it before the barrier), so there is no deadlock.
- resolve_local_single_file reinterpreted an On-Device folder as a base
single_file load whenever it held exactly one .safetensors, so a PEFT LoRA
adapter folder (adapter_config.json + adapter_model.safetensors) with a
family-token name was picked as a base checkpoint, evicting the resident
model before from_single_file failed on the adapter weights. Skip adapter
folders (adapter_config.json) and the adapter_model basename so the pick
stays a pipeline load and 400s in validation, before the GPU handoff.
Adds regression tests for both.
The video load-request preflight gated its local-pipeline check on
root.is_dir(), so a bare local file (e.g. /models/ltx-2.safetensors) sent
with model_kind=pipeline skipped it, passed validation, and the route then
evicted the resident GPU model before from_pretrained failed on the
non-directory path. Gate on root.exists() instead (mirroring the image
loader's diffusion.validate_load_request), so a local file is rejected up
front. Add a regression test.
* fix: match qwen3-thinking chat template double-newline in response pattern
The Qwen3-thinking chat template generates `<think>\n\n` (double newline)
after the think tag, but `train_on_responses_only` was looking for
`<think>\n` (single newline).
`\n\n` is token 271 while `\n` is token 198 -- different tokens, so the
pattern match in `train_on_responses_only` fails, masking ALL tokens and
dropping 100% of training samples.
Update the response pattern from `<think>\n` to `<think>\n\n` to match
what the actual qwen3-thinking template generates.
Fixes#6919
* fix qwen3 thinking response marker
---------
Co-authored-by: Ayushman Paul <ayushman@HP>
Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
Two evict/corruption fixes surfaced by review of the diffusion training path:
- krea-2 sets force_bf16 in the DiT trainer spec, but the route-level
_FORCE_BF16_FAMILIES preflight listed only qwen-image and z-image, so a
krea-2 start with mixed_precision=fp16 passed the route check, reserved
training and evicted resident GPU models, and only the child trainer then
raised. Add krea-2 to the set and a drift-guard test asserting it equals
the trainer specs whose force_bf16 is set.
- The dataset-upload same-stem duplicate check compared stems
case-sensitively, so on case-insensitive filesystems (Windows / default
macOS) sample.png and Sample.jpg both passed even though their caption
sidecars sample.txt / Sample.txt resolve to the same file, silently
sharing and corrupting one caption. Compare stems and the same-name guard
with casefold at both the on-disk and in-batch sites.
* show chat by by last activity
* Update chat thread updated_at logic and enhance sidebar chat item handling
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
The /delete-cached guard refused deleting a cached repo whose id is a bare
prefix of the loaded one: with Qwen/Qwen-Image-2512 loaded, deleting the
separate cached Qwen/Qwen-Image returned 400 "Unload the model before
deleting" because loaded_id.startswith(repo_id) matched across the repo
boundary. Both are real, independently cached catalog artifacts, so a
supported delete was wrongly blocked.
Add a /-boundary aware _loaded_id_matches_repo helper (mirroring
hub/services/models/deletion.py) and use it at all six guard sites (chat,
Images, Video, and their in-flight loading-id loops), so only the loaded
repo itself or a file within it blocks deletion. Add a regression test
covering the sibling-prefix pair.
Fold PR #6872's image-generation fixes into the branch, deduped against the
round-12 dataset-upload and gallery integrity work already on image-generation.
Fixes carried forward from #6872:
- fp8 single-file transformer memory estimate: an fp8 checkpoint loads with no
quantization_config and diffusers upcasts it to bf16 (~2x resident), so budget
it accordingly in _plan_memory and estimate_safetensors_dense_mib.
- dense-quant OOM-evict preflight: when the GGUF fits resident but the dense bf16
transformer this path materializes does not, skip the fast path up front rather
than evict the current pipeline and OOM in finalization. Combined with the
existing offload->resident candidate re-plan so both the family-table estimate
and the on-disk shard measurement gate engagement (unified on the
transformer_resident_override_mib plan override).
- ControlNet: evict the previous module and its from_pipe wrapper before loading a
new one so swapping ControlNets within a base-model load cannot accumulate to OOM.
- ControlNet union_control_mode: raise on an unknown control type instead of
silently defaulting to canny.
- edit-family mask rejection: raise instead of silently dropping a mask on an
image-editing model that has no inpaint pipeline.
- companion cache: walk the snapshot dir and exclude transformer/ so the
dense-quant prefetch's cached shards do not inflate the companion total and
wrongly force offload.
- training: drop piecewise_constant from the LR scheduler enum and force bf16 for
fp16-incompatible families.
- dataset upload: batch-atomic staging with the same-stem duplicate guard.
- images page: guard negative-prompt restore on guidance>0, clear stale ControlNet
selection on restore, and revert an optimistic quant label when a pipeline load
never starts.
- uninstall (sh + ps1): keep the owner-marker guard on sd.cpp removal.
Conflicts resolved in favour of image-generation's evolved memory system,
loadSpecFor catalog, and stop-and-save (lora_path) run detection; #6872's fp8 and
dense-preflight fixes carried forward on top. All affected backend tests pass
(test_diffusion_backend, test_diffusion_training, test_diffusion_lora_trainer,
test_video_gallery, test_diffusion_controlnet).
test_delete_keeps_sidecar_listable_when_mp4_unlink_fails patched os.unlink, but delete()
calls path.unlink(). On Python 3.10 Path.unlink dispatches through a cached _accessor bound
to os.unlink at import, so patching os.unlink had no effect there and the simulated mp4
unlink wrongly succeeded (the assert delete() is False failed on 3.10 while passing on
3.11+, where Path.unlink looks up os.unlink dynamically). Patch Path.unlink itself, which
delete() invokes directly on every supported Python version. Test-only; no behavior change.
union_control_mode fell back to control_mode=0 (the canny head) for ANY unmapped control
type, so a typo'd or unsupported value like 'detph' silently conditioned the map as canny
instead of failing. preprocess_control passes non-canny maps through unchanged, so that map
would be interpreted under the wrong mode with no error. Now only 'passthrough' (or an empty
type) keeps the deliberate mode-0 default; any other unknown type raises ValueError, which
the generate route maps to a 400. Known modes are unchanged.
Follow-up to removing piecewise_constant from the trainable scheduler allow-list: the
DiffusionTrainingStartRequest.lr_scheduler Literal still advertised it, so a client that
picked it straight from the schema passed request validation and then hit the 400 from
normalized(). Remove it from the enum too so the API only offers schedulers the trainers
can actually run, and add a test asserting the enum never advertises a scheduler outside
the validation allow-list.
* Studio: account for DeepSeek-V4 compute buffer in context auto-fit
DeepSeek-V4-Flash's lightning indexer plus compressed sparse attention reserve a
large context-scaling compute buffer that _compute_buffer_ctx_bytes did not model
(the KQ-mask and dequant-scratch rates both miss it, even with an f16 cache).
Measured on UD-Q4_K_XL at ub 512 it is about 65.5 GiB at 1M context, which the
mask estimate puts near 1.5 GiB, so the auto-fit kept the full 1M train context
and llama-server OOM'd allocating the ~70 GB buffer, then spilled to CPU (~4
tok/s). Add a deepseek4-gated flat plus per-token term so the fit caps the context
(about 256k on a B200) and the model stays fully on GPU.
* [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>
piecewise_constant is the only diffusers scheduler that needs a step_rules string, and
neither diffusion trainer passes one (get_scheduler is called with only warmup/training
steps, and there is no config field for it). Accepting it let /diffusion/start pass
normalized(), free the resident GPU workloads, spawn the trainer, and only then crash in
the subprocess (get_piecewise_constant_schedule does step_rules.split(",") on None) -- the
exact evict-then-fail the up-front validation exists to prevent. Reject it now with a clear
400. The remaining six schedulers all run with only warmup/training steps.
* Studio: add assistant response details panel
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Hide model badge by default, show on hover/focus
Wrap MessageResponseModelBadge in a span with hidden/group-hover visibility classes to reduce visual clutter. The badge now only displays when hovering or focusing on the assistant message, improving the UI presentation. Updated corresponding tests to verify the new CSS classes.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Add DeepSeek-V4-Flash-GGUF to Studio with none/high/max reasoning
Adds unsloth/DeepSeek-V4-Flash-GGUF as a default selectable model with the
recommended decoding defaults (temperature 1.0, top_p 1.0 from the official
generation_config.json) and its three tier reasoning control. The high/max
ladder is surfaced for deepseek-v4 model ids and flows through the existing
enable_thinking_effort reasoning style via chat_template_kwargs, so no
frontend changes are needed.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio DeepSeek-V4: segment-scope high, enable thinking for lone effort, render tests
Match deepseek-v4 on whole repo-name segments so a future deepseek-v40 or
deepseek40 cannot false-match the synthetic 'high'. In _request_reasoning_kwargs,
emit enable_thinking when a named effort level is sent without it, so the
newly exposed High mode renders thinking-on over the API (the UI already sent
it explicitly). Add a none/high/max render-path test file (jinja behind
importorskip) with a lone-high regression.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Promote the example-import staging dir into the dataset folder in one atomic
same-filesystem rename instead of a per-file move loop. A hard process death between
two moves left the folder with SOME images, and the image_count>0 idempotency check
then accepted that truncated dataset as complete on retry. The folder is created
empty on this path, so a single rename is atomic; a folder holding unrelated
non-image files falls back to a per-file move rather than abort.
Reject a second uploaded image whose stem matches an existing image but differs by
extension (sample.png vs sample.jpg): both map to one <stem>.txt caption sidecar
(the kohya/diffusers convention), so keeping both would silently corrupt captions
during training. Exact-name overwrites and .txt caption sidecars stay allowed.
list_videos globs *.mp4 but requires a readable sidecar, so a video whose sidecar
is gone is skipped. delete()/clear() dropped the sidecar first, so if the mp4
unlink then failed (a Windows lock from a concurrent stream/transcode), the
still-present mp4 vanished from the gallery with no way to retry the delete. Unlink
the mp4 first and only then best-effort the sidecar: the worst case is now an
orphaned sidecar, which list_videos already ignores.
Two overlapping /diffusion/start requests can interleave between the is_active()
check and the reservation, so reserve() itself must reject a second reservation
atomically. Otherwise both callers reserve, both free the GPU's resident chat or
image model, and the loser only 409s after the eviction -- the evict-then-fail the
reservation exists to prevent. reserve() now raises under the lock if a start is
already reserved or a job is already running.
* Run the malware gate on the RAG embedding model before it loads
Setting the RAG embedding model through PUT /api/settings/embedding-model
persisted an arbitrary repo and later handed it straight to
SentenceTransformer, which deserializes pickle weights. Unlike the normal
model-load paths, this route never ran evaluate_file_security, and force
skipped verification entirely, so a repo Hugging Face flags as unsafe (or
any repo under force) could be downloaded and loaded in the backend
process without a scan.
Run the malware/pickle scan at both ends: the settings endpoint now scans
before persisting and returns 409 on a flagged repo even under force
(force still only skips the is-embedding-model type check for offline or
local repos), and the embedder scans again at the load sink so a name that
arrives via env or default is covered too. Local paths and unreachable
scans fail open inside evaluate_file_security, and the sink never bricks
the embedder on a gate error.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Thread the load token into the embedding scan and hard-fail on a block
The load-sink scan ran without a token, so evaluate_file_security (which
passes token=False when none is given) could not reach a gated or private
repo and failed open for exactly the model SentenceTransformer would still
load. Resolve the loader's own token (HF_TOKEN env or the cached login)
and pass it to the sink scan, and fall back to it in the settings endpoint
when the request omits one.
The sink previously raised a plain RuntimeError, which the llama-server
fallback in encode() and _build_st_backend_or_fallback() swallowed as a
routine ST failure, silently switching backends instead of blocking. Raise
a distinct UnsafeEmbeddingModelError that both fallback paths re-raise, so
a flagged model hard-fails.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Scan sentence-transformers module dirs and scope the embedding pickle gate to the ST backend
Extend the RAG embedding malware gate so a poisoned pickle under a SentenceTransformer
module dir (for example 0_Transformer/pytorch_model.bin) blocks. Those dirs are read
from the repo's modules.json and passed as load roots to evaluate_file_security at both
the settings endpoint and the load sink, so such a pickle is treated as root-level there
instead of an unreferenced nested shard that was previously allowed.
Scope the ST pickle scan to the sentence-transformers backend. On the llama-server
backend the embedder loads GGUF files (inert) from the -GGUF companion repo, never the
ST repo's pickle, so a custom ST repo with a flagged pickle and a clean GGUF companion
is no longer rejected. The existing GGUF availability checks already cover that path.
Return 403 for the hard security block instead of 409. The settings UI routes every 409
into the forceable save-anyway flow, but this block cannot be bypassed by force, so it
now uses a distinct status the client treats as non-forceable.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Base the embedding pickle scan on the actual backend, not just the resolver
_llama_backend_active only consulted the auto resolver, so on a GPU box
where auto resolves to sentence-transformers but the process already fell
back to the llama-server backend at runtime (a torch or CUDA load/encode
failure), it returned False and the settings endpoint hard-blocked a save
whose ST pickle is flagged even though the process loads only inert GGUF.
Add active_backend_is_llama, which reflects the actual built backend (True
when the cached backend is a LlamaServerBackend, including a runtime
fallback) and otherwise defers to the resolver as a fresh process would,
and delegate _llama_backend_active to it.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Report the cached embedding backend verbatim, not the resolver
active_backend_is_llama() fell through to the config resolver whenever a
backend was already built but was not llama-server, so a live
sentence-transformers backend could report llama=True once the resolver
picked llama (GPU heuristic or a runtime config change) and wrongly skip
its pickle scan. Once a backend exists, return isinstance(backend,
LlamaServerBackend) directly; only defer to the resolver before any
backend is built.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Guard RoPE scaling against the transformers v5 buffer blank; honor extended factor
Add a family-agnostic guard that builds each rotary from a scaled config,
blanks its non-persistent buffers (what transformers v5 does on load), runs
loader._fix_rope_inv_freq, and asserts every buffer is restored to its scaled
value (llama3 and longrope). This catches the whole bug class, not just the
one call site, and is validated to fail on the pre-fix repair.
Also make LlamaExtendedRotaryEmbedding read the llama3 factor from the config
instead of hardcoding 8 (wrong for Llama-3.2, factor 32), falling back to the
Llama-3.1 defaults when built without a config.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Pass config into extended rotary codegen; skip v5 round-trip on transformers 4.x
- patch_llama_rope_scaling now builds the llama3 extended rotary with
config=self.config so it reads the real factor (32 for Llama-3.2) instead
of falling back to 8; the template already references self.config.
- test_v5_blank_repair_roundtrip now skips when loader._NEEDS_ROPE_FIX is
False, since _fix_rope_inv_freq is a no-op on transformers 4.x and cannot
restore the blanked buffers there.
* Raise stream deadlock-guard timeouts from 0.2s to 5.0s in passthrough tests
These asyncio.wait_for guards bound test setup and cross-task event
signaling that complete near-instantly on success; the 0.2s budget is a
latency assertion in disguise and times out under CI scheduling load
(seen on the 3.11 matrix leg while 3.10/3.12/3.13 pass the same commit).
5.0s matches the timeout used elsewhere in the suite and still fails fast
on a real hang. No test relies on the guard expiring.
* Extended rotary reads rope_parameters as well as rope_scaling
transformers v5 stores llama3 scaling under config.rope_parameters and
exposes rope_scaling only as a back-compat property. Reading that property
works on 5.0-5.13 (verified: factor resolves to 32 for Llama-3.2), but a
future release may drop the shim, after which the subclass path would fall
back to factor 8. Read either field so the factor survives the rename.
Adds test_extended_rotary_reads_rope_parameters_v5 (fails on the old
single-field read: rope_parameters-only config resolves to 8, not 32).
* [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>
start_diffusion_training freed resident GPU models and only then called
service.start(config), which is where is_active() first flips true. During that
free-then-spawn window a concurrent /images/load or /video/load saw training as
inactive, passed its training guard, acquired the GPU, and began a background load,
so the trainer and that pipeline both allocated VRAM. Add reserve()/unreserve() to
the training service (is_active() also reports the reservation) and reserve BEFORE
the free, in a try/finally so a failed start rolls the reservation back. An
overlapping load's guard now refuses during the window. Regression tests: the route
reserves before the free (and the free sees an active service), and the service
reservation marks active then rolls back.
Three gates the image backend has were missing on the video path:
- kind/extension mismatch: a model_kind 'single_file' with a .gguf name (or 'gguf'
with a non-.gguf) passed the video preflight, so the route acquired VIDEO and
evicted the resident GPU owner before the wrong single-file loader failed in the
background. Reject the mismatch up front, mirroring the image loader.
- Windows-shaped missing local pick: the missing-path check only matched POSIX
prefixes (/ ~ ./ ../), so a missing Windows path (C:\ / C:/ or any backslash path)
was treated as a Hub repo and only failed after the GPU handoff. Use the image
loader's is_absolute()/backslash path-shaped check.
- speed=off auto-quant: a full-pipeline load with Speed=off but Precision=auto still
promoted the unset precision to auto-quant, so on a dense-capable GPU it engaged
torchao quantization and forced the speed back to default -- silently breaking the
user's bit-exact request. Suppress the auto promotion when speed_mode is off, as
the image loader does.
Regression tests for each.
The dataset upload took Path(filename).name, which on POSIX does not split on a
backslash, so a Windows client sending a backslash path in the multipart filename
stored the name verbatim. The caption/thumbnail/delete endpoints then run it
through _safe_dataset_image_path, which rejects backslashes and '..', so the
labeling grid could list an image it could never preview, caption, or delete (an
orphan). Fold backslashes to forward slashes before taking the basename so the
true name is stored, and reject a basename that still contains '..' at upload
rather than persisting an unmanageable entry. Regression test covers a Windows
backslash path (stored and served under the clean basename) and a '..' rejection.
The multipart upload passthrough was a bare /api/train/diffusion/dataset prefix,
so its JSON sub-routes (PUT .../{name}/caption/{filename}, POST .../import-example)
also bypassed the default JSON body cap and inherited the far larger upload limit.
A large caption/import body would then be buffered and parsed up to the upload cap
before the route-level length checks ran. Match the upload route by EXACT path
instead: the JSON sub-routes fall through to the normal small-JSON cap. Adds an
upload_passthrough_exact_paths param to MaxBodyMiddleware, with regression tests
that the sub-routes keep the default cap and a large sub-route body is 413'd.
start_diffusion_training is async but called _preflight_gated_base inline; it does
a blocking urllib urlopen HEAD to Hugging Face (up to a 5s timeout) to detect a
gated/unauthorized base repo. On a slow or unreachable network that stalled the
FastAPI event loop, freezing every concurrent status/progress/cancel request until
it returned or timed out. Wrap it in asyncio.to_thread, matching the dataset
preflight and GPU cleanup just below it. Regression test asserts it runs off the
coroutine thread.
Several image/video/training preflights ran before the route acquires the GPU or
frees resident models, but let a doomed local pick through and only failed deep in
the background load, after the user's chat/Images/Video model was already evicted.
- Local base_repo / base_model: _is_trusted_diffusion_repo accepts any existing
local path, but the base loads via from_pretrained (needs model_index.json). A
local dir that is not a diffusers pipeline passed the trust gate, evicted the
resident model, then failed. Add a shared _assert_local_base_is_pipeline check
and call it in the image, video, and training preflights.
- Dataset images: discover_image_caption_pairs only checked filenames, so a
corrupt or zero-byte upload passed the start-route preflight, freed the GPU, then
crashed the spawned trainer in PIL. Add an opt-in verify_images decode probe
(cheap PIL header check) that the start route enables; the trainers leave it off
since they decode every image anyway.
- Local single-file safetensors: the On-Device scanner advertises a bare
.safetensors directory (no model_index.json) as a text-to-image model, but the
picker starts it as a pipeline with no filename, so every click 400s. Reinterpret
such a pick as a single_file load of the sole checkpoint (resolve_local_single_file)
so the advertised model is actually loadable.
Regression tests for each: local non-pipeline base (image/video/training), the
verify_images decode gate, and resolve_local_single_file.
start_diffusion_training is an async route, but it called the blocking
_free_gpu_for_diffusion_training() inline. That teardown waits on generation
locks and joins the export subprocess, so it can block for seconds and freeze the
FastAPI event loop, stalling every concurrent status/progress/cancel request
until it finishes. Wrap it in asyncio.to_thread, mirroring the dataset-preflight
call just above it and the inference routes' load/unload offloading. Add a
regression test asserting the cleanup runs off the coroutine thread.
_reset_step_cache looked up reset_stateful_hooks on the transformer, but on a
diffusers CacheMixin transformer (Flux, QwenImage) that method lives only on the
HookRegistry; the transformer-level entry point is _reset_stateful_cache. So with
FBCache engaged on an image model the reset was a silent no-op, and the next
generation reused the previous request's first-block residual: a tensor-shape
mismatch (crash) when the resolution or batch changed, or stale cached output
otherwise. Prefer _reset_stateful_cache and fall back to reset_stateful_hooks,
matching the video backend. Update the tests to the real hook name.
The images and video unload routes dropped their arbiter claim whenever the
backend was not committed-loaded, but a concurrent /load re-acquires the owner
and starts a background load that is not is_loaded for its whole download and
finalize window. Releasing during that window cleared the newer load's claim, so
a later chat/image load saw no owner, skipped eviction, and could allocate a
second heavy model on the GPU. Gate the release on loading_repo_ids() too (both
diffusion engines and the video backend expose it), not just the committed
state, so an overlapping load keeps ownership. Add regression tests for the
in-flight case on both routes.
The three train_precision_modes capability tests patched is_available and
get_device_capability but not is_bf16_supported, so on the CPU-only CI runner
(where the real probe is False) the dense modes collapsed to nf4 and the fp8 /
mxfp8 assertions failed; they only passed on a bf16 GPU dev box. An Ada or
Blackwell GPU is by definition bf16-capable, so the helper must stub it True to
exercise the capability gate the tests target.
* Studio: re-exec to prepend torch's bundled CUDA libs to LD_LIBRARY_PATH
On Linux the dynamic linker reads LD_LIBRARY_PATH before the RUNPATH baked into
torch's .so files, so a pre-existing LD_LIBRARY_PATH pointing at a system CUDA
(conda, a Docker base image, /usr/local/cuda-*/lib64) shadows torch's bundled
nvidia/*/lib libraries and causes undefined-symbol errors when the Studio backend
imports torch. Detect torch's lib dirs without importing torch, prepend them to
LD_LIBRARY_PATH, and re-exec once (LD_LIBRARY_PATH is only read at process start).
Linux-only, sentinel-guarded against re-exec loops, and called only from run.py's
__main__ so library/embedder imports (e.g. Colab's `from run import run_server`)
are never re-exec'd.
* [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>
When the images page mounts with a diffusion model already loaded (a prior
session or another route left it resident), refreshStatus only set status; the
steps/guidance sliders kept the unrecognised-model fallback (few-step, no CFG),
so a resident full model such as flux.1-dev generated at 9 steps guidance 0 and
produced garbage until the user re-picked it. Add a one-shot effect that seeds
the sliders from defaultsFor(status.repo_id) when the page discovers a resident
model it did not load itself, and wires Reapply to it for non-GGUF pipelines.
Guarded by lastLoad and a per-repo ref so a manual slider edit is never
clobbered.
The start route's precision preflight folded bf16/int8/fp8 into the CUDA
requirement but omitted mxfp8, so an mxfp8 request on a GPU-less host (or an
older CUDA GPU without Blackwell) passed the preflight, evicted resident image
and chat models, then raised only in the spawned trainer child. Mirror
_resolve_base_precision: require CUDA for mxfp8 and re-check the Blackwell
(sm100+) capability up front, so a doomed run is rejected before teardown.
- The companion base for a GGUF/single-file image load is resolved from the GGUF
repo's base_model card tag when no base_repo is passed, and that value loads via
from_pretrained. The explicit base_repo is already trust-gated, but the card tag is
attacker-controlled metadata on any remote repo, so it now clears the same
unsloth/allowlist/local trust bar; an untrusted tag is dropped in favour of the
curated family default and never reaches from_pretrained. This closes a pickle
deserialization vector on the normal GGUF load path (a user loading an attacker's
GGUF repo whose card points base_model at a malicious pipeline), matching the
trust discipline the ControlNet path already applies via evaluate_file_security.
The allowlist already contains every legitimate variant base, so variant
resolution for the supported unsloth GGUFs is unchanged.
- The images/unload route ran the slow VRAM-freeing unload on a thread and then
released the DIFFUSION arbiter owner unconditionally. release() is owner-guarded
and identity-less, so a concurrent /images/load that re-acquired DIFFUSION while
the unload ran would have its ownership cleared by the trailing release, and a
later chat load would then see no owner, skip eviction, and OOM against the newly
resident pipeline. The route now releases only when nothing is resident again.
- _apply_group_offload placed the resident companions before attaching the
transformer's group-offload hooks so a companion OOM returns with no hooks and the
whole-module fallback stays valid, but the streamed loop itself installs hooks on
each DiT in turn. On a dual-DiT pipeline where the second tower failed after the
first got its hooks, it returned False with hooks already installed, and the
caller's enable_model_cpu_offload fallback then crashed (diffusers rejects it on a
partially group-offloaded pipe). It now propagates the real failure once any hook
is installed, so the load fails with its actual cause instead of a misleading crash.
Adds regression tests: the untrusted card tag dropped to the family default (trusted
tag still honoured, explicit base still wins), unload keeping ownership when a model
is still resident (and releasing when not), and the partial dual-DiT hook set
propagating rather than falling through to a crashing whole-module offload.
* Studio: apply presence_penalty on the safetensors and MLX inference paths
The safetensors and MLX generate paths resolved the inference config and
then dropped presence_penalty before generation, so the same model applied
the configured value under GGUF and 0 under safetensors/MLX. Thread the
already-resolved presence_penalty through the orchestrator command, worker
gen_kwargs, and the safetensors/MLX generate calls, and apply it with a
small logits processor (subtract once per distinct completion token,
prompt excluded, presence not frequency, zero is a no-op, negatives raise).
Backwards compatible: presence_penalty defaults to 0.0 (byte-identical
output when unset) and the GGUF path is unchanged. Also forward min_p on
the legacy /generate/stream route and add the missing min_p field to
GenerateRequest.
* Studio: bound presence_penalty generated ids to valid vocab range on both paths
The presence-penalty logits processors index by generated token ids. The
torch path filtered only the upper bound (seen < vocab_size), so a negative
id would silently wrap to the wrong row; the MLX path had no bound at all,
and MLX out-of-bounds indexing is documented undefined behavior (crash or
memory corruption on Apple Silicon), unlike torch's harmless negative wrap.
Bound generated ids to [0, vocab) consistently on both paths:
- torch: seen[(seen >= 0) & (seen < vocab_size)] (zero-regression safety net;
real completion tokens are always in range).
- MLX: route out-of-range/negative ids to a discarded scratch slot via
mx.where and a (vocab + 1)-wide scatter-assign mask, then subtract. MLX has
no boolean-mask filtering (data-dependent output shape), so this keeps a
fixed shape, stays on-device, and preserves once-per-distinct-token
semantics without any torch/numpy dependency.
Add torch tests for out-of-range and negative ids (only in-range distinct
ids penalized, stray ids ignored, no wrong-index wrap) and a bound-documenting
MLX test that runs on the arm64 macOS CI.
* [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>
* Studio: serialize the compare-mode dispatcher lifecycle to fix a start race
_generate_dispatched (compare mode) bypasses _gen_lock so two concurrent
compare requests can both reach _start_dispatcher. The check-then-spawn there
had no lock, so both could observe no live dispatcher and each spawn one. The
extra dispatcher is orphaned (self._dispatcher_thread tracks only the last) and
during a later unload it can consume the 'unloaded' reply off _resp_queue before
unload_model's _wait_response, hanging the unload on its timeout.
Add _dispatcher_lifecycle_lock and take it around the whole body of both
_start_dispatcher and _stop_dispatcher, so start/stop cannot interleave and the
second concurrent starter sees the dispatcher alive and returns. _start_dispatcher
now returns whether it actually spawned the thread, and _generate_dispatched
derives dispatcher_preexisting from that atomic result instead of a separate
unlocked is_alive() read.
No call site holds _mailbox_lock when calling start/stop, so joining the
dispatcher (which takes _mailbox_lock) under the new lock cannot deadlock; the
lock order is always _gen_lock then _dispatcher_lifecycle_lock and is never
inverted.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: refuse dispatcher start queued behind an unload's stop
A compare request could pass the early _unload_pending check, then block in _start_dispatcher on _dispatcher_lifecycle_lock behind an unload's _stop_dispatcher. When the unload released the lock the start spawned a fresh dispatcher, which became the resp_queue reader and consumed the worker's unroutable 'unloaded' reply before unload_model's _wait_response saw it, hanging the unload for 300s.
Gate _start_dispatcher on _unload_pending under the lifecycle lock, and set _unload_pending under the same lock ahead of the stop, so any start queued behind the stop observes the unload and refuses. Ordering stays _gen_lock -> _dispatcher_lifecycle_lock. Adds a regression test forcing the queued-behind-stop interleaving.
* [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>
- Image load now trust-gates a client-supplied base_repo. validate_load_request
rejects a base_repo that is not an unsloth/* repo, an allowlisted official base, or a
local path, mirroring the repo_id gate and the video loader. The route passes
base_repo into that pre-eviction validation, so an authenticated client can no longer
keep model_path on a trusted GGUF while pointing base_repo at an arbitrary remote repo
that the server would download and deserialize (a from_pretrained pickle/config path),
and no resident model is evicted for the rejected load.
- The keepwarm middleware now tracks the image and video generation routes
(/images/generate, /images/generations, /video/generate), so
other_inference_request_count() sees an in-flight generation and an API-key training
start is refused (409) before its unload would cancel that generation. endswith keeps
the GET *-progress and */cancel variants untracked.
- The OpenAI-compatible /v1 surface is now blanket body-capped like /api/inference,
instead of only /v1/chat/completions and /v1/completions. Every /v1 POST route
(images/generations, audio, embeddings, responses, messages, ...) buffers a JSON body
and none is a multipart-upload passthrough, so an unbounded ImageGenerationRequest
prompt on /v1/images/generations can no longer be buffered outside the request limit.
Adds regression tests: the base_repo trust gate at both the backend (untrusted remote
raises, local passes) and the route (untrusted base_repo returns 400 with no load), the
keepwarm tracking of the image/video generation paths (and not the progress/cancel
variants), and the /v1 surface being body-protected.
