* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add SDXL diffusion family (U-Net pipeline support)
SDXL is the first U-Net family in the diffusion backend: its denoiser is
pipe.unet (UNet2DConditionModel), not a DiT pipe.transformer, and a single-file
.safetensors is the whole pipeline rather than a transformer-only file. The
backend previously assumed a DiT transformer everywhere, so add the two hooks a
U-Net family needs and register SDXL.
DiffusionFamily gains denoiser_attr ("transformer" for DiT, "unet" for SDXL) and
single_file_is_pipeline (SDXL loads a single file via pipeline_class.from_single_file
with the base repo as config, instead of transformer_class.from_single_file plus a
companion assembly). _align_vae_dtype now reads the denoiser generically so img2img
and inpaint keep the VAE and U-Net dtypes aligned.
The non-GGUF trust gate is extended with a short, exact-match, safetensors-only
allowlist of official base repos (the SDXL base/refiner and sdxl-turbo), because
SDXL ships only as a full pipeline and has no unsloth-hosted GGUF. Local paths stay
trusted as before; a random repo, even one that detects as SDXL, is still rejected.
The image-conditioned and ControlNet workflows are the standard SDXL pipelines,
built around the resident modules via from_pipe like every other family, so SDXL
gets txt2img, img2img, inpaint, outpaint, upscale, LoRA and ControlNet. There is no
native sd.cpp mapping yet, so the no-GPU route falls back to diffusers.
Frontend catalog gains SDXL Base 1.0 and SDXL Turbo entries with SDXL step/guidance
defaults (Turbo: few steps, no CFG; base: ~30 steps, real CFG).
Tests: new test_diffusion_sdxl.py (family shape, detection, trust allowlist, model
kind, U-Net VAE-dtype alignment, LoRA gate) plus loader-branch tests in
test_diffusion_backend.py (pipeline-kind from_pretrained, single-file whole-pipeline
from_single_file, allowlist accept/reject). Verified live on GPU: sdxl-turbo loads
both as a pipeline and as a single file and generates coherent txt2img + img2img.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Images: LoRA free-text Hugging Face entry + recipe round-trip
The backend has always accepted a bare Hugging Face repo id (owner/name, or
owner/name:weight-file.safetensors) as a LoRA, downloading and applying it. But the
picker only rendered when the curated catalog had entries, and the catalog is empty,
so there was no UI path to apply any LoRA. Show the LoRA section whenever the loaded
model supports LoRA, and replace the curated-only dropdown with a text input: type a
Hub repo id, or pick a discovered adapter from a datalist of suggestions when the
catalog is populated.
Also restore LoRAs when loading a recipe. restoreSettings now parses the recipe's
"id:weight" strings (splitting on the last colon, since the id itself may contain one
for a specific weight file) back into the selection, so replaying a saved image
reproduces its adapters. The generate payload trims hand-typed ids and drops empty /
zero-weight rows, and a model swap clears the selection (a LoRA is family-specific)
without discarding a free-text pick that is not in the curated list.
* Add diffusion LoRA training (SDXL text-to-image)
First diffusion training path in Studio: train a LoRA on the SDXL U-Net from an
image + caption dataset and export it as a diffusers .safetensors that the existing
diffusion LoRA loader (and any diffusers pipeline) can load.
core/training/diffusion_lora_trainer.py:
- DiffusionLoraConfig with validation/defaults (rank, alpha, targets, lr, steps, grad
accumulation, resolution, min-SNR gamma, gradient checkpointing, lr scheduler, seed,
mixed precision).
- discover_image_caption_pairs: captions from metadata.jsonl / captions.jsonl, per-image
.txt/.caption sidecars, or a dreambooth instance_prompt fallback (pure, unit-tested).
- run_diffusion_lora_training: the loop -- freeze base, PEFT-wrap the U-Net attention
projections, VAE-encode (fp32 VAE to avoid the SDXL fp16 overflow), sample noise +
timesteps, predict, MSE loss with optional min-SNR weighting (epsilon / v-prediction),
AdamW + get_scheduler + grad accumulation + grad clipping, then export via
save_lora_weights. Emits worker-protocol events (model_load_*, progress, complete) and
polls should_stop for a clean stop with a partial save.
- run_diffusion_training_process: mp.Queue subprocess adapter (event_queue / stop_queue),
so the training worker can spawn it; plus a CLI entry point.
Only SDXL (U-Net) is trained here; DiT families and the Studio UI form + route wiring are
follow-ups. The trainer is decoupled and worker-ready.
Tests: test_diffusion_lora_trainer.py covers caption discovery (metadata / sidecar /
instance prompt / skip-uncaptioned / errors), config normalisation + validation, the SDXL
add-time-ids, and the dict->config adapter. Verified live on GPU: a 60-step SDXL LoRA run
lowers the loss, exports a ~45 MB adapter, and loading it back shifts generation from
baseline (mean abs pixel diff ~55/255).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* diffusion trainer: emit learning_rate in progress events (Studio pump compatibility)
The Studio training pump reads 'learning_rate' from progress events; the diffusion
trainer emitted 'lr'. Rename the field (and the CLI reader) so the trainer's events are
directly consumable by the existing training status/SSE machinery when it is wired into
the worker, without a translation shim.
* Wire diffusion LoRA training into the Studio API
Make the SDXL LoRA trainer reachable from the app with a small, self-contained job
service and JSON routes, deliberately separate from the LLM TrainingBackend (whose
lifecycle -- LLM config build, per-run SQLite rows, matplotlib plots, transfer-to-chat-
inference -- is text-training specific and would mis-handle a diffusion run).
core/training/diffusion_training_service.py: DiffusionTrainingService runs one job at a
time -- validate the config cheaply (before any spawn), spawn the trainer subprocess
(spawn context, parent-lifetime bound), pump its events (model_load_* / progress /
complete / error) into an in-memory status snapshot, and support a clean stop. The
subprocess context and target are injectable so the full start -> pump -> status ->
complete path is unit-tested without real multiprocessing or torch.
routes/training.py: POST /api/train/diffusion/start (400 on a bad config, 409 when a job
is already running), POST /api/train/diffusion/stop, GET /api/train/diffusion/status
(JSON poll). models/training.py: DiffusionTrainingStartRequest + response schemas
mirroring DiffusionLoraConfig, so model_dump() passes straight through.
Tests: test_diffusion_training.py -- service happy path, bad-config-before-spawn,
concurrent-job rejection, clean stop, crash-without-terminal-event, event transitions;
plus route wiring via the FastAPI TestClient (start / 422 / 400 / 409 / status / stop)
with a mocked service. The diffusion trainer's progress events already use the field
names this path expects.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Images: add a Train LoRA (SDXL) dialog
Surface the diffusion training API in the Images page. A "Train LoRA" button in the top
bar opens a self-contained dialog to fine-tune an SDXL LoRA on a folder of images: pick
the base model, dataset folder, output folder, an optional instance prompt, and the core
hyperparameters (steps, rank, resolution, batch, learning rate), then Start. The dialog
polls the training status while open and shows a progress bar, step count, live loss, and
the saved adapter path, with a Stop button for a clean stop.
The dialog is independent of the loaded generation model (training runs in its own
subprocess), and prefills the base model with the loaded checkpoint when it is SDXL, else
the SDXL base. api.ts gains startDiffusionTraining / stopDiffusionTraining /
getDiffusionTrainingStatus plus their types, matching the /api/train/diffusion routes.
* Import diffusion training schemas from models.training directly
The import-hoist lint flags newly re-exported names in the models/__init__.py hub as
unused (it does not treat __all__ membership as a use). Import the three diffusion
training schemas straight from models.training in routes/training.py, where they are
used in the route annotations and calls, and drop the __init__ re-export.
* Remove stray async task scratch outputs committed by mistake
* ControlNet: reject filesystem-like ids and do not cache a model past an unload race
Two review findings on the ControlNet path:
- resolve_controlnet's bare-repo fallback accepted any id with a slash, so a
path-shaped id (/tmp/x, ../x) reached from_pretrained as a local directory.
Restrict the fallback to a strict owner/name HF repo id shape.
- _controlnet_pipe now re-checks the cancel event after the blocking
from_pretrained: an unload that raced the download had already cleared the
caches, so caching the late module would pin it past the unload.
* Pipeline prefetch: fetch only the default torch weights
A full-pipeline prefetch kept every repo file outside assets/, so an official
repo that ships multiple formats (SDXL Base: fp16 variants, ONNX, OpenVINO,
Flax, a top-level single-file twin) downloaded tens of GB from_pretrained never
loads. Skip non-torch exports and dtype-variant twins in
_pipeline_file_downloaded, and drop a component .bin when the same directory
carries a picked safetensors weight (diffusers' own preference).
* Diffusion LoRA training: fall back to fp16 when CUDA lacks bf16
The default mixed_precision=bf16 hard-fails on pre-Ampere GPUs (T4 / V100 /
RTX 20xx) which have no bf16 compute; check torch.cuda.is_bf16_supported()
and drop to fp16 there.
* Diffusion training service: join the old pump outside the lock
start() joined a finished job's pump thread while holding the service lock,
but the pump's final state writes need that same lock, so the join always
burned its full timeout and a stale pump could then overwrite the new job's
state. Join outside the lock (with a re-check after), and fence _apply_event
and the exit handler by process identity so a superseded pump can never touch
the current job's state. Adds regression tests for both.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion LoRA training: harden config handling, cancellation, SDXL conditioning, and safety
Addresses review findings on the SDXL LoRA trainer:
- Gate the base model with the same trust check as inference (unsloth/*, allowlisted
official bases, or a local path) before from_pretrained, so an untrusted remote repo
is never fetched or deserialised.
- Check the stop signal before the (slow) model load, not only between steps, so a
cancel during download is honoured; a stop may carry save=False to cancel without
leaving a partial adapter.
- Per-sample SDXL add_time_ids from the actual crop (original size + crop offset, with
the offset mirrored on horizontal flip) instead of a fixed uncropped-square tensor.
- Apply EXIF orientation before resize/crop so rotated photos train upright.
- Skip gradient clipping when max_grad_norm <= 0 (the Studio 'disable' value) instead
of scaling every gradient to zero.
- Coerce Studio config strings/blanks: learning_rate string to float, blank hf_token to
anonymous, gradient_checkpointing 'none'/'true'/'unsloth' to bool; reject a zero/negative
lora_alpha or learning_rate.
- Alias the generic Studio training payload keys (model_name/max_steps/batch_size/lora_r/
lr_scheduler_type/random_seed) onto the diffusion field names.
- Mirror the trained adapter into loras/diffusion so the Images LoRA picker discovers it.
- Report worker exceptions in both message and error keys so the failure is not lost.
Adds regression tests for the config coercion/validation and aliasing.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* ControlNet: address review findings on the diffusers path
- resolve_controlnet enforces catalog family compatibility so a direct API call
cannot load a ControlNet built for another family through the wrong pipeline.
- Unknown ControlNet ids now surface as a 400 (call site maps FileNotFoundError
to ValueError) instead of a generic 500.
- strength 0 disables ControlNet entirely, so a no-op selection never pays the
download / VRAM cost; the control image is decoded and validated BEFORE the
ControlNet is resolved or built, so a malformed image fails fast for the same reason.
- ControlNet loads use the base compute dtype (state.dtype is a display string,
not a torch.dtype, so it silently fell back to float32) and honor the base
offload policy via group offloading instead of forcing the module resident.
- Empty/malformed HF token coerced to anonymous access.
- Flux Union ControlNet control_mode mapped from the selected control type.
- resolve_controlnet drops the unused hf_token/cancel_event params.
- ControlNetSpec validates guidance_start <= guidance_end (clean 422).
- Images UI ControlNet Select shows its placeholder when nothing is selected.
Adds regression tests for family enforcement and the union control-mode map.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion training API: LLM interlock, pre-spawn VRAM free, path containment, no dropped knobs
Four review findings on the diffusion training start path:
- It spawned the SDXL trainer without checking the LLM TrainingBackend, so a
start while an LLM run was active put two trainers on the same GPU. Add a
symmetric interlock: diffusion start returns 409 when LLM training is active,
and LLM start refuses while a diffusion job is active.
- It went straight to service.start() without freeing GPU residents. Add a
pre-spawn free of the export subprocess, the resident Images pipeline (with an
arbiter release), and chat models, mirroring the LLM start path.
- data_dir / output_dir were passed through unresolved, so Studio-relative names
failed and absolute paths bypassed containment. Resolve them with
resolve_dataset_path / resolve_output_dir before spawn (400 on an uncontained
path).
- The request model dropped max_grad_norm and lora_target_modules, so runs that
set them trained with defaults. Add both fields.
The gemini pump-join deadlock was already fixed earlier (join outside the lock +
proc-identity fence). Note: honoring a stop DURING model load is a trainer-loop
change owned by the diffusion training engine PR (should_stop polled before the
first optimizer step). Adds route + model regression tests.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion LoRA: harden resolution, native tag precedence, and diffusers teardown
Address review findings on the LoRA path:
- resolve_one: normalise a blank/whitespace hf_token to None (anonymous access)
and reject a client-supplied weight file with traversal / absolute path.
- resolve_specs: convert FileNotFoundError from an unknown/stale id to ValueError
so the route returns 400 instead of a generic 500.
- _scan_local: disambiguate local adapters that share a stem (foo.safetensors vs
foo.gguf) so each is uniquely addressable.
- inject_prompt_tags: the backend-validated weight now wins over a user-typed
<lora:ALIAS:...> for a selected adapter; unselected user tags are left alone.
- diffusers _apply_loras: reject a .gguf adapter with a clear error before touching
the pipe (diffusers loads safetensors only).
- _unload_locked: drop the explicit unload_lora_weights() on teardown; the pipe is
dropped wholesale (freeing adapters), so the previous call could race an in-flight
denoise on the same pipe.
- Images page: use a stable LoRA key and clear the selection (not just the options)
when the catalog refresh fails.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion: guard trust check against OSError and validate conditioning inputs
- _is_trusted_diffusion_repo: wrap Path.exists() so a repo id with invalid
characters (or a bare owner/name id) can't raise OSError; treat any failure as
not-a-local-path and fall through to the unsloth/ allowlist. validate_load_request
still raises the clear FileNotFoundError for a genuinely missing local pick.
- generate(): reject mask_image / upscale / reference_images supplied without an
input image, and reject reference_images on a family that does not support
reference conditioning, instead of silently degrading to txt2img / img2img.
* SDXL: reject GGUF up front, skip unused base weights, drop refiner, and harden helpers
Addresses review findings on the SDXL family:
- Reject a GGUF load for single_file_is_pipeline families (SDXL) in validate_load_request,
before the route evicts the current model; SDXL has no transformer-only GGUF variant.
- Skip base-repo weight files when a whole-pipeline single file is loaded: from_single_file
(config=base) needs only the base config/tokenizer/scheduler, so a local .safetensors no
longer triggers a multi-GB base download.
- Remove the SDXL refiner from the non-GGUF trust allowlist: it is an img2img-only pipeline
but this backend loads every sdxl repo as the base txt2img pipeline.
- Normalize a blank/whitespace hf_token to None once in load_pipeline so every load branch
degrades to anonymous instead of erroring on a malformed token.
- Read the denoiser dtype from a parameter (compile-wrapped modules may lack .dtype) and
access state.family.denoiser_attr directly.
Adds/updates regression tests for the trust allowlist, GGUF rejection, and base-config filter.
* Images: preserve restored LoRAs through model load and never send hidden LoRAs
- The LoRA effect cleared the selection on every load->capable transition, which
wiped adapters restored from a gallery recipe before the model finished loading.
Track the previously-loaded family in a ref and clear only on a real family swap;
keep the selection on the initial load and on unload.
- Gate the generate payload's loras on loraCapable so a restored selection that is
hidden (loaded model does not support LoRA) is never sent to the backend.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Images Train LoRA dialog: token, validation, precision, base-repo prefill, gating, refresh
Nine review findings on the SDXL training dialog:
- Forward the saved Hub token so a gated/private SDXL base can be trained (the
image load flow already sends it).
- Re-seed the base-model field from the current default each time the dialog
opens; the keep-alive dialog otherwise kept its mount-time default after a
model loaded.
- Prefill from base_repo (the diffusers pipeline) rather than repo_id, which for
a GGUF/single-file SDXL load is the checkpoint path from_pretrained can't open.
- Add client-side validation of steps/rank/resolution/batch/learning-rate before
the request.
- Expose a precision selector (bf16/fp16/fp32) so non-bf16 GPUs can train from
the UI, not only the API.
- Gate the dialog on the active Images route (active && trainOpen) so switching
tabs closes it and stops its polling.
- Rescan the LoRA picker when a run completes, so a freshly-trained adapter
appears without a model reload.
- Cap the dialog height and scroll the body so the Start/Stop footer stays
reachable on short viewports.
- Correct the copy to not over-promise picker auto-discovery.
Freeing the resident Images pipeline before training is handled backend-side in
the diffusion training start route.
* Merge diffusion-sdxl into diffusion-lora-ux; keep options-only LoRA catch
The catalog-refresh .catch from the lower branch clears the selected adapters
too, which is right for its catalog-only picker but wrong here: this picker
holds free-text HF repo ids that are valid without being in the catalog, so a
transient refresh failure must not wipe them. Family swaps still clear the
selection and hidden LoRAs are never sent.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Train LoRA dialog: stop suggesting absolute paths the backend rejects
The dataset and output placeholders showed /path/to/... examples, but the
training routes resolve those fields inside the Studio home and reject
absolute paths outside the approved roots, so following the placeholder
produced a 400. Use folder-name placeholders and say in the labels and the
dialog description where each folder resolves.
* Align the VAE to the denoiser's first FLOATING dtype, not its first parameter
A GGUF-quantized transformer's leading parameters are packed uint8 storage,
so reading next(parameters()).dtype handed nn.Module.to() an integer dtype
and every image-conditioned generation on a GGUF model (Qwen-Image-Edit)
failed with a 500. Probe the parameters for the first floating dtype, treat
an all-integer module as a no-op, and also catch TypeError so an unexpected
dtype can never break generation. Regression test included.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Count LR scheduler warmup/decay in optimizer steps, not micro-steps
lr_sched.step() runs once per outer optimizer step (after the gradient
accumulation inner loop), for train_steps total. The scheduler was
configured with num_warmup_steps and num_training_steps multiplied by
gradient_accumulation_steps, so with accumulation > 1 a warmup or
non-constant schedule stretched past the run and never reached the
intended decay. Count both in optimizer steps.
* Address Codex review findings on the image-workflows PR
Keep diffusion.py importable without torch: the compile/arch patch modules
import torch at module level, so import them lazily at their load/unload
call sites instead of at module load. This restores the torchless contract
so get_diffusion_backend() works on a CPU/native sd.cpp install.
Match family reject keywords and aliases as whole path/name segments, not
raw substrings, so an unrelated word like edited, edition, or kontextual no
longer misroutes or hides a valid base image model, while supported edit
families (Qwen-Image-Edit, FLUX Kontext) still resolve. Mirror the same
segment matching in the picker task filter.
Route FLUX.2-dev native guidance through --guidance like the other FLUX
families rather than --cfg-scale. Reject native upscale requests that have
no input image. Read image header dimensions and reject over-limit inputs
before decoding pixels, so a crafted small-payload image cannot spike
memory. Reject an upscale that would shrink the source below its input
size. Validate the model_kind against the filename extension before the
GPU handoff. Estimate a local diffusers pipeline's size from its on-disk
weights so auto memory planning does not skip offload and OOM. Report
workflows: [txt2img] from the native backend status so the Create tab
stays enabled for a loaded native model. Clamp the outpaint canvas to the
backend's 4096px decode limit.
Adds regression tests for segment matching and kind/extension validation.
* Guard inference loads and worker lifetime against diffusion training
Teach the chat and image load guards about an active diffusion (SDXL) LoRA
job: a chat load is refused (its footprint cannot be fit-checked against the
trainer) and an image load is refused outright, mirroring the existing LLM
training guards, so a load can no longer allocate GPU memory alongside the
trainer and undo the pre-start cleanup.
Bind the diffusion trainer subprocess to the parent's lifetime and scrub the
native path lease secret from it by running the child through
run_without_native_path_secret, matching the inference/export/LLM workers, so
a Studio crash or kill no longer leaves the trainer holding the GPU.
Reset in_model_load on the complete and error terminal events: a stop or
failure during model loading otherwise leaves the status reporting a stale
loading indicator after the job has ended.
* Harden diffusion LoRA handling on the diffusers and native paths
Reject LoRA on a torch.compile'd diffusers transformer (Speed=default/max):
diffusers requires the adapter loaded before compilation, so applying one to
the already-compiled module fails with adapter-key mismatches. The status
gate now hides the picker and generate raises a clear message instead.
Convert a cancelled Hub LoRA download (RuntimeError Cancelled) to the
diffusion cancellation sentinel in resolve_specs, so an unload/superseding
load during resolution maps to a 409 instead of a generic server error.
Drop weight-0 LoRA rows before the native support gate so a request carrying
only disabled adapters stays a no-op on families where native LoRA is
unsupported, matching the diffusers path.
Reject duplicate LoRA ids in the request model: both apply paths suffix
colliding names, so a repeated id would stack the same adapter past its
per-adapter weight bound.
Strip all user-typed <lora:...> prompt tags on the native path (only the
selected adapters are materialized in the managed lora-model-dir, so an
unselected tag can never resolve), and restore saved LoRA selections from a
gallery recipe so restore reproduces a LoRA image.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Harden ControlNet resolve, gallery metadata, and the control-type picker
Check cancellation immediately after a ControlNet from_pretrained and before
any device placement, so an unload/eviction that raced the download does not
allocate several GB onto the GPU after the load was already cleared.
Require a loadable weight or shard index (not just config.json) before a local
ControlNet folder is advertised, so an interrupted copy is hidden instead of
failing deep in from_pretrained as a generic 500.
Do not record a strength-0 ControlNet in the gallery recipe: it is treated as
disabled and skipped, so the image is unconditioned and the metadata must not
claim a ControlNet was applied.
Build the control-type picker from the selected ControlNet's advertised
control_types instead of a hardcoded passthrough/canny pair, so a union model
with a precomputed depth or pose map sends the correct control_mode.
* Address further Codex findings on the image-workflows PR
- Persist the actual output image size in the gallery recipe instead of the
request sliders: Transform/Inpaint/Edit derive the size from the uploaded
image, Extend grows the canvas, and Upscale resizes it, so the sliders
recorded (and later restored) the wrong dimensions for those workflows.
- Reject a remote '*-GGUF' repo loaded as a full pipeline (no single-file
name) in validate_load_request, so the unloadable pick fails before chat is
evicted rather than deep in from_pretrained.
- Only publish an image-conditioned from_pipe wrapper to the shared aux cache
when the load is still current: from_pipe runs under the generate lock but
not the state lock, so an unload racing its construction could otherwise
cache a wrapper over torn-down modules that a later load would reuse.
- Verify the Windows CUDA runtime archive checksum before extracting it, like
the main sd-cli archive, so a corrupt or tampered runtime is rejected rather
than extracted next to the binary.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Refuse non-SDXL base models at diffusion training start
The trainer only supports the SDXL U-Net, but a FLUX / Qwen-Image / Z-Image
repo or a GGUF filename passed as base_model was accepted and then failed
minutes later inside StableDiffusionXLPipeline.from_pretrained with an
unrelated-looking error. Add a name-based guard in normalized() so known
DiT-family names and .gguf checkpoints are rejected up front, which the API
start route surfaces as an immediate 400 with a message that says exactly
which bases are trainable. Unrecognisable names still pass through so custom
local SDXL checkpoints keep working.
* Add diffusion dataset upload and training info endpoints
Training an image LoRA required knowing the Studio home layout and copying
files onto the server by hand, which is the most confusing step of the whole
flow. Two small endpoints fix that:
- GET /api/train/diffusion/info reports the datasets and outputs roots plus
every dataset folder that contains images (with image/caption counts), so
the UI can offer a picker instead of a blind free-text path.
- POST /api/train/diffusion/dataset uploads images and optional caption
.txt / metadata.jsonl files into a named folder under the datasets root,
creating it on first use and accumulating on repeat uploads so large sets
can arrive in batches. Names are validated to a single path component and
files stream to disk under the same per-upload size cap as LLM dataset
uploads. The returned name is a valid data_dir for /diffusion/start.
* Rework the Train LoRA dialog into a guided SDXL flow
The dialog assumed users knew the Studio home layout and that only SDXL is
trainable, and hid both facts behind free-text fields. Restructure it around
the three real decisions:
- Base model is a dropdown of the trainable SDXL picks (Base 1.0, Turbo, the
loaded SDXL pipeline when there is one) with a custom repo/path escape
hatch, instead of a bare text field defaulting to a repo id.
- Training images come from an in-browser upload (new dataset endpoints) or
a picker over existing dataset folders with image/caption counts. No shell
access or knowledge of the datasets root is needed any more, and the
captioning rules are explained inline.
- The output field is now Adapter name and the instance prompt is labelled
as the trigger prompt, with a no-captions warning wired to the selected
dataset's actual caption count.
Hyperparameters collapse behind a training settings toggle since the
defaults suit a first run. A completed run says where the adapter went and
offers Done / Train another, and the top-bar button gets an icon and a
plainer description. The dialog title states the SDXL-only scope and that
other families load LoRAs but cannot train them yet.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Validate diffusion training config before freeing the GPU
The start route freed resident GPU workloads (export, Images pipeline, chat)
before the service validated the config, so a start that was then refused,
now including a non-SDXL base model, tore down the user's loaded model for
nothing. Run the same cheap normalise pass first; the LLM path already
follows this rule via its before_spawn hook.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Refactor diffusion LoRA training into a family-aware platform
Split the SDXL trainer into a shared, architecture-agnostic layer so more model
families can be trained without duplicating the plumbing:
- New core/training/diffusion_train_common.py holds the config + validation, dataset
discovery, event emission, stop protocol, adapter publishing, and a lazy trainer
registry (get_trainer). diffusion_lora_trainer.py keeps the SDXL-specific loop and
re-exports the moved names so existing imports are unchanged.
- The SDXL-only base-model blocklist becomes a positive check: the family is resolved
from the base model (or an explicit model_family) via the diffusion family registry,
and a known-but-not-yet-trainable family is refused with a clear message. Unknown
custom names still default to the SDXL trainer.
- DiffusionFamily gains a trainable flag and train_base_repos; SDXL is marked trainable.
DiT families flip on when their trainers land.
- Trained adapters now write a <name>.json metadata sidecar (family, base model, rank,
trigger prompt, ...) that the LoRA scanner reads to family-gate the adapter in the
picker instead of showing it as unknown for every model.
- The training base-model trust allowlist adds the official FLUX.1-dev, Z-Image-Turbo,
and Qwen-Image repos (safetensors-only, no remote code).
* Retain diffusion training loss history and expose it in status
The training service kept only the latest loss, so a live loss chart could show a
single point. Fold each progress event into bounded (step, loss, lr) history arrays
(capped at 4000 points, decimated when full) plus the latest throughput and peak VRAM,
and record the family / base model / catalog path on completion. The status endpoint
returns these as a nested metric_history object the UI can chart directly, and the
start request accepts an optional model_family override.
* Tests for the diffusion training platform
Cover the trainer registry (get_trainer resolves SDXL, unknown family raises),
family resolution (explicit model_family validation, resolved_family on the config),
the metadata sidecar write + scan read with family gating, and the service loss-history
folding (append, bad-point skipping, decimation at cap, family/perf fields) plus the
status route nesting metric_history.
* Add diffusion dataset labeling and example-import endpoints
The Train tab needs to let users caption small datasets in the browser and
pull in a ready-made set to see training work end to end, neither of which
the upload-only endpoint supported.
Add, under /api/train/diffusion/dataset:
- GET {name}/images lists every image with its resolved caption (metadata
beats a per-image sidecar, matching the trainer's discovery order) so
uncaptioned images are visible and flaggable.
- GET {name}/image/{filename} serves an image, with ?thumb=<px> returning a
cached downscaled JPEG kept in a hidden .thumbs subdir (regenerated when
the source is newer) so the labeling grid stays light.
- PUT {name}/caption/{filename} writes, or when blank clears, the .txt
sidecar; DELETE {name}/image/{filename} removes the image plus its
sidecars and thumbnails.
- GET dataset-examples lists a curated, license-labelled registry, and
POST dataset/import-example materializes one into a dataset folder as
numbered images + .txt captions. Two loaders cover the shapes seen in the
wild: streaming rows from datasets.load_dataset (dog-example, Tuxemon) and
a snapshot + jsonl walk for imagefolder repos whose captions live in a
non-standard *.jsonl (the public-domain tarot set). Imports are idempotent
and cap the image count.
Filenames and dataset names are validated against path traversal and pinned
inside the datasets root.
* Test diffusion dataset labeling and example-import endpoints
Cover caption precedence, thumbnail generation and .thumbs exclusion,
caption write/clear, image delete cleanup, path-traversal rejection on
names and filenames, and example import with a mocked datasets.load_dataset
(files plus sidecars written, idempotent second call, cap respected, load
failure mapped to 502).
* Add flow-matching DiT LoRA trainers (FLUX.1-dev, Qwen-Image, Z-Image)
Extends diffusion LoRA training beyond SDXL to the three popular DiT families
via a single shared flow-matching loop parameterised by small per-family specs
(loading, prompt/latent encoding, transformer forward, save). Verified against
diffusers 0.38.0:
- FLUX.1-dev: 2x2 latent packing + image ids, guidance-embed forward, on-the-fly
nf4 QLoRA of the 12B transformer (the dev repo is gated, so training needs the
user's HF token).
- Qwen-Image: 5D VAE latents normalised by the per-channel latents_mean/std,
img_shapes forward, prequant nf4 base by default (on-the-fly nf4 for the bf16
base).
- Z-Image: list I/O with the reversed timestep convention and a negated
prediction, bf16 only.
The registry (get_trainer) and DiffusionFamily.trainable / train_base_repos now
route these families to the DiT trainer; the SDXL blocklist guard is replaced by
a positive family resolution that also rejects GGUF repos (inference-only) and
still-unsupported families. Per-family defaults + labels + VRAM notes are exposed
via family_train_infos for the Train UI.
Memory: caption embeddings are precomputed once and the text encoders freed
before the loop; gradient checkpointing (non-reentrant, required for bnb 4-bit)
and 8-bit AdamW are on by default.
* Speed up + shrink SDXL LoRA training (precompute text embeds, 8-bit AdamW)
SDXL re-encoded every caption with both CLIP text encoders on every step (pure
waste, since captions are constant) and kept the encoders resident. Precompute
each unique caption's embeddings once, then free the text encoders before the
loop: numerically identical (embeddings are deterministic and this consumes no
torch RNG, so the noise/timestep stream is unchanged) but faster and ~1.5 GB
lighter. Default the optimizer to 8-bit AdamW (bitsandbytes) with an fp32
fallback, halving optimizer state with no meaningful LoRA quality cost. Env
toggles (UNSLOTH_DIFFUSION_NO_PRECOMPUTE / _FP32_OPTIM) let the accuracy guard
A/B the paths.
* Expose trainable families in /diffusion/info and preflight gated bases
The training info endpoint now returns the trainable model families (name,
label, default + allowed base repos, recommended defaults, and a VRAM/access
note) so the Train UI can offer a base picker with realistic guidance. The start
route preflights a gated base repo (HEAD model_index.json with the user's token)
BEFORE freeing resident GPU workloads, so a missing FLUX.1-dev license/token
fails fast with an actionable 400 instead of evicting the loaded model and then
hitting a confusing mid-load 401.
* Tests for DiT trainers, family resolution, info families, gated preflight
Cover the DiT spec table, the QLoRA prequant heuristic, the Z-Image bf16-only
guard, the gated-repo name check, family resolution now that FLUX/Qwen/Z-Image
are trainable (and GGUF repos are rejected as inference-only), the families list
in /diffusion/info, and the gated-base 400 preflight that leaves the GPU
untouched.
* Add diffusion training API client: metrics, families, dataset labeling, examples
Extends the Images training client for the Train tab: the status type now carries
metric_history (step/loss/lr) plus catalog_path/family/base_model/samples_per_second/
peak_memory_gb; the start request gains model_family; and info gains an optional
families list (per-family bases + defaults). Adds typed calls for the dataset
labeling and one-click example endpoints: list images with captions, thumbnail URL,
write/clear a caption, delete an image, list example datasets, and import an example.
* Add diffusion Train panel: config, dataset labeling, live charts, deploy
New full-page training workspace for the Images tab. Left column configures the run:
model family (FLUX.1-dev, Qwen-Image, Z-Image, SDXL in popularity order, with per-family
VRAM/license notes and defaults, backfilled from the backend families list when present),
base repo, dataset (existing folder, browser upload, or one-click example import), an
in-browser caption labeling grid (per-image thumbnail + caption saved on blur, delete,
uncaptioned highlight), adapter name, trigger prompt, and collapsed training settings.
Right column shows the live run: progress + loss/avg/speed/peak-VRAM readouts, the reused
training loss/LR charts fed from metric_history, and a completion card that deploys the
adapter into Create or starts another run.
* Wire Create/Train tab switch into the Images page and deploy flow
Replaces the Train LoRA dialog with a top-bar Create | Train segmented control next to
the model selector. Create renders the existing generation workspace unchanged; Train
renders the full-page training panel (unmounted in Create so its polling stops while the
backend run and its retained metric history survive a tab switch). Adds a deploy handler:
loading the trained adapter's base as a pipeline, queueing the adapter so the LoRA
discovery effect applies it once the base is loaded and LoRA-capable for the matching
family (with a mismatch warning), seeding the prompt with the trigger, and switching back
to Create. Removes the now-unused dialog.
* Wrap the DiT training forward in bf16 autocast
The fp32 LoRA parameters and the bnb 4-bit base matmuls need a single
compute dtype during the forward, exactly like the diffusers dreambooth
scripts run under accelerator.autocast. Without it the 4-bit backward on
FLUX.1-dev fails with an illegal-address CUBLAS error partway into the
first step. Z-Image and Qwen-Image smokes are unaffected and the SDXL
path (its own trainer) is untouched.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix Train tab example cards and Create/Train tab layout
The example-dataset cards used a two-column grid in the ~340px config
column, which wrapped titles one word per line and let the long license
text overrun into the neighbouring card. Switch to one card per row with a
horizontal layout: title with a compact truncated license badge (full text
in the tooltip), a two-line clamped description, and the Import button on
the right.
The Create/Train switch had an icon inside the Train trigger that overhung
the pill corner. Drop the icon, make both triggers a fixed equal width so
the active pill sits flush in the top bar.
* Show only loss and learning-rate charts for diffusion training
The Train tab reused the LLM charts section, which also rendered an empty
Grad Norm card and an Eval Loss card showing an Evaluation not configured
placeholder with a red smear. Neither applies to diffusion LoRA training.
Add a diffusion-only two-card view that reuses the loss and learning-rate
cards directly with fixed presentation defaults, and note under the loss
chart that per-step loss is noisy by design so users read the smoothed
line for the trend rather than the raw jitter.
* Add a dataset preview strip to the Train tab
When a dataset with images is selected, show a strip of up to 8 sampled
thumbnails with a +N more tile, so users can see what is in the folder
before training. Clicking the strip opens the existing caption review
grid. Samples are drawn evenly across the folder and refresh on dataset
change or after an upload/import.
* Stop example cards from overflowing the Train config column
The example-dataset cards still overran the ~340px config column: the
license used the Badge component whose baked-in w-fit and whitespace-nowrap
ignored the max-width and truncate, and the grid children had the default
min-width auto so wide content pushed past the column edge and clipped the
Import buttons. Replace the badge with a plain truncating pill span, and
give the config column min-w-0 with overflow-x-hidden so nothing escapes
its width.
* Add Smithsonian Butterflies and Nouns example datasets
Two permissive ~100-image sets for the Train tab: huggan/smithsonian_butterflies_subset
(CC0, the classic diffusers-docs training set, imported as a subject set with a trigger
prompt since its metadata columns are species names not captions) and m1guelpf/nouns
(CC0, captioned pixel-art avatars via the text column). Both cap at 100 images.
* Paginate the Train tab caption grid with prev/next controls
Large example datasets (100+ images) rendered every tile at once, so the
caption review grid grew unbounded. Show 24 images per page with < >
chevrons and an x-y of N indicator; a new dataset or refresh resets to
the first page.
* Offer example datasets in the Train dropdown with previews
Add an Examples group to the training-images dropdown that imports a
curated dataset in one pick, alongside the existing cards. Cards now show
up to three preview thumbnails pulled from the public HF datasets-server
so the set is visible before download. Hide the trigger prompt when every
image already has a caption (a captioned style set needs no trigger), and
turn the training-settings toggle into a ghost button with a rotating
chevron.
* Clamp the training base repo to the selected family
The base-model select's state could briefly hold the previous family's
repo after a family switch (the reseed effect runs a beat later, and a
value with no matching option makes the browser display the first option
anyway). The request then carried the stale repo: picking Qwen or Z-Image
still sent black-forest-labs/FLUX.1-dev and surfaced FLUX's gated-repo
error under the wrong family. Derive an effectiveBase clamped to the
current family's repos and use it for the select value, the start
request, and the deploy fallback.
Also move the Trigger prompt above Adapter name: the trigger describes
the dataset, the name only labels the output.
* Speed up diffusion LoRA training and cut DiT peak VRAM by a third
Perf core for the diffusion trainers, defaults preserving the training math:
- Phased model loading: the pipeline now loads without its transformer
(conditioning only), captions are encoded and the text encoders freed,
the VAE latent cache is built and the VAE freed, and only then does the
transformer load. The multi-GB denoiser never shares VRAM with the
encoders, cutting measured peak VRAM on B200: FLUX 17.1 -> 10.4 GB,
Qwen-Image 19.1 -> 12.8 GB, Z-Image 7.3 -> 4.7 GB.
- Latent cache (cache_latents, default on): per-image crop/flip variants
(cache_variants, default 4 vs the single frozen variant of the diffusers
--cache_latents) store the VAE posterior's affine parameters, so every
step still draws a fresh VAE sample; a cached center-crop Z-Image run
matches the uncached one at the bf16 nondeterminism floor.
- True batching: train_batch_size now actually batches the transformer
forward (it was silently 1). nf4 dequant dominates the step cost, so
batch 4 lands near batch-1 step time: 4.0x samples/s on Qwen-Image,
3.1x on FLUX, 2.1x on Z-Image, with multi-seed loss envelopes
overlapping batch-1.
- LR scheduler support in the DiT loop (lr_scheduler / lr_warmup_steps
were accepted but ignored); progress events now report the real
per-step LR.
- TF32 + high fp32 matmul precision under enable_tf32 (default on),
snapshot/restored around the run. cudnn.benchmark is scoped to a
caller opt-in only: autotuning the fp32 VAE convs doubled peak VRAM
on the DiT families for zero steady-state gain.
- Vectorized sigma gathering (drops a per-step Python search loop),
cached FLUX img_ids/guidance, fused torch AdamW fallback, steady-state
samples_per_second (excludes the first-step warmup).
- Regional torch.compile plumbing (compile_transformer off/on/auto with
eager fallback): auto stays off over a bitsandbytes base where compile
is a net loss (27 s warmup, slightly slower steady on Z-Image); it
arms automatically for the dense/quantized speed modes that follow.
- Stop parity with the LLM trainer: /api/train/diffusion/stop accepts an
optional {save} body and the service forwards save=False as a
no-save cancel; a new preparing event surfaces cache-build progress.
- SDXL trainer gets the same latent cache, perf flags, and fused
fallback; its batching, LR schedule, and min-SNR stay as they were.
Verified: 83 backend tests green; per-family 30-40 step runs with
adapter round-trip generation through the normal LoRA path (FLUX,
Qwen-Image, Z-Image all pass).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add base_precision speed modes to DiT training: bf16 2.3-2.6x, int8, fp8
New base_precision config for the DiT trainers: nf4 (unchanged default) |
bf16 | int8 | fp8 | auto, advertised per family + per machine through
/api/train/diffusion/info (precision_modes, recommended_precision,
supports_compile) so the UI can gate the selector.
- bf16: dense transformer + regional torch.compile (auto-armed). The
measured speed mode: 2.3x nf4 on FLUX (1.81 -> 4.12 steps/s), 2.6x on
Z-Image (2.5 -> 6.38 steps/s) on B200, at dense-weight VRAM
(FLUX 24.7 GB / Z-Image 13.6 GB peak vs 10.4 / 4.7 for nf4).
- int8: torchao weight-only int8 on the frozen base, quantized AFTER
add_adapter (quantizing first trips peft 0.18's TorchaoLoraLinear,
which is incompatible with the torchao 0.16 config API). Runs eager:
inductor rejects the int8 subclass training graph (aliased subclass
outputs), so compile is force-disabled for it.
- fp8: torchao convert_to_float8_training on the frozen linears
(filter skips lora_ modules, proj_out, non-divisible-by-16 dims,
pad_inner_dim), applied after add_adapter, compile auto-armed.
Works and round-trips, but measured SLOWER than compiled bf16 at
LoRA-training shapes (FLUX 3.15 vs 4.12 steps/s; Z-Image similar),
so it is an explicit opt-in and auto never picks it.
- auto: free VRAM (measured before load) + dense-size table -> bf16
when it fits with headroom, int8 in the middle band, else nf4.
Prequant bnb repos always resolve to nf4; dense modes on them are
rejected at validation with a pointer to the family's dense base.
Two crashes found and fixed along the way:
- The cuDNN SDPA backend's training graph fails on the FLUX attention
shapes (torch 2.10 + cu130, B200): mha_graph.execute errors, then the
context degrades into illegal memory accesses. The perf-flag guard now
pins flash/mem-efficient SDPA for the run (mathematically equivalent,
snapshot/restored). nf4 escaped it by routing attention differently.
- Regional compile now uses dynamic=True (the inference layer's proven
default): dynamic=False specialisation fused a gemm_and_bias epilogue
that failed with CUBLAS_STATUS_EXECUTION_FAILED on the FLUX training
graph; dynamic=True is also faster (Z-Image 3.84 -> 6.38 steps/s).
Verified: 98 backend tests green (new test_diffusion_base_precision.py:
validation, auto policy table, fp8 filter, compile gating, /info fields);
per-mode 40-step runs on FLUX + Z-Image with loss means inside the nf4
envelope and adapter round-trip generation through the normal LoRA path
for bf16-, fp8-, and int8-trained adapters.
* Clear TF32 flags when enable_tf32 is off so the opt-out is strict fp32
* Address review: auto int8 requires the dense-load transient to fit, dense modes are CUDA-only, auto respects bf16 compute, exact cudnn SDPA restore
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: fp32 latent cache stats + strict-JSON-safe progress floats
- Latent caches (DiT + SDXL) now hold the posterior mean/std in fp32 and draw the
per-step sample in fp32, casting only the result to the training dtype. This
matches the in-loop path (encode fp32 -> sample fp32 -> cast) exactly instead of
sampling in bf16; the cache is tiny so the doubled RAM is negligible.
- The training service nulls non-finite floats (NaN/Inf loss, avg_loss,
learning_rate) at its single ingestion point so status snapshots and persisted
run records stay strict-JSON serializable; the metric history skips non-finite
loss points. Test covers NaN/Inf progress followed by a finite point.
* Address review: gate auto int8 on torchao, scope dense validation to DiT
- base_precision="auto" only picks int8 when torchao is importable (the int8
quantize has no runtime fallback, unlike fp8); otherwise the middle band falls
back to nf4. Threaded as a parameter so the policy stays pure.
- The dense-mode validation (prequant base / bf16 compute) now applies only to
DiT families: sdxl ignores base_precision entirely, so a leftover value can no
longer fail an SDXL run. The mode-name validity check still runs everywhere.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Gate int8 and fp8 on a functional torchao import, not find_spec
The Windows ROCm torchao import stub satisfies find_spec and even lets
from torchao.quantization import quantize_ succeed, but its quantize_ is a
no-op: auto would pick int8, leave the transformer dense, and disable
compile as if it were quantized. has_functional_torchao imports the exact
symbols the int8 path uses and rejects the stub via its sentinel; both the
auto picker and the /info advertised modes now use it
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Skip the sigma-gather test when diffusers is not installed
CI runs the backend suite without diffusers; the test checks our index math against
the scheduler's own gather, so it skips rather than fails there.
* Coerce cache_latents and enable_tf32 string flags in the config dict path
The generic Studio config dict path can deliver these flags as strings, and a
non-empty string like "false" is truthy, so an opt-out silently no-ops (the
latent cache still builds, TF32 stays on). Coerce them the same way
gradient_checkpointing already is.
* Remove committed runtime scratch artifacts and ignore their dirs
logs/ (a 1.3 MB ComfyUI object_info dump plus stale PID files), temp/ (PR body
and commit message scratch), and async_task_outputs/ (agent task transcripts)
are environment specific runtime artifacts that were committed by accident and
carry stale local state into every checkout. Remove them and gitignore the
directories so they cannot be re-added.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Restore pre-Ampere bf16 fail-fast in the DiT trainer
The perf rewrite dropped the bf16 capability guard, so a pre-Ampere CUDA
device (T4/V100/RTX 20xx) would die deep in model load with an opaque dtype
error instead of a clear message. Restores parity with the SDXL trainer.
* Size-gate the automatic diffusion latent cache
The latent cache holds two fp32 posterior tensors per crop/flip variant per
image, pinned on CUDA hosts, so datasets with thousands of images can exhaust
host or pinned memory with no fallback. Estimate the cache size from the first
real encoded latent and fall back to per-step VAE encoding when it exceeds a
4 GiB budget. UNSLOTH_DIFFUSION_FORCE_LATENT_CACHE bypasses the gate; the
existing UNSLOTH_DIFFUSION_NO_LATENT_CACHE opt-out is unchanged.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Gate DiT training precision: deny fp8 for Qwen, gate explicit int8 on torchao, gate advertised dense modes + route on bf16
- normalized() + family_train_infos() mirror the inference fp8 deny for
Qwen-Image (activation outliers exceed fp8's range and corrupt the trained
result); int8 stays allowed and the UI no longer advertises fp8 for it.
- _resolve_base_precision() gates an explicit int8 on a FUNCTIONAL torchao, the
same gate auto and /info already apply, so a missing/stub torchao fails fast
instead of silently loading dense with compile disabled.
- train_precision_modes() gates the dense modes (bf16/int8/fp8/auto) on
torch.cuda.is_bf16_supported(), so a non-bf16 CUDA GPU (T4/V100/RTX 20xx) is
offered only nf4 instead of a start that evicts resident models and then fails.
- start_diffusion_training preflights bf16 support for the DiT families BEFORE
_free_gpu_for_diffusion_training(), so any DiT start (nf4 included, since the
trainer requires bf16 unconditionally on CUDA) fails fast without eviction.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Gate DiT training on functional torchao for explicit int8; hide always-400 DiT modes on non-bf16 GPUs
The start route preflight only rejected non-bf16 GPUs; an explicit int8 request on
a host with a missing or stub torchao passed the preflight, evicted resident GPU
workloads, then died in the trainer child (its int8 base quantizer has no fallback).
Fold both gates into training_precision_preflight_error so int8-without-torchao fails
fast before eviction. Also empty the advertised DiT precision_modes (and surface the
reason in vram_note, drop compile) whenever the bf16 preflight would reject the family,
so /info never offers an nf4 DiT option the route always 400s.
* Reject dense DiT precisions on a CUDA-absent host before eviction; stabilize family-info tests
The start-route preflight caught the bf16-GPU and int8-torchao requirements but not the dense
precisions' CUDA requirement: on a GPU-less host bf16_unsupported_reason exempts CPU-only, so a
bf16/fp8 (or int8-with-torchao) DiT request passed the preflight, evicted resident workloads, then
raised only in the trainer child. Add the dense-mode CUDA gate mirroring _resolve_base_precision so
the doomed run is rejected up front. Also pin bf16_unsupported_reason in the two positive-path
family-info tests so they are deterministic across GPU types (a non-bf16 CUDA box would otherwise
empty every DiT family's advertised modes).
* [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>
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
- Trainers emit the pre-clip gradient norm; the service keeps a bounded
grad_norm history and the Train tab renders a Grad Norm chart next to
Loss and LR
- Completed runs show 'Training complete' with a celebratory marker in
the success color instead of a plain status word
- metadata.jsonl caption keys now match on Windows (as_posix relative
paths) in both the trainer discovery and the dataset image records
- RMSNorm eager patch skips installation on torch builds without
F.rms_norm instead of failing at forward time
- GGUF compute description no longer says the GGUF is dequantised: the
INT8/FP8/FP4 modes load the base model's bf16 transformer and quantise
that directly; label no longer wraps in the Advanced panel
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 9): pre-quantized transformer loading
The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on
the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs
13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline
(scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton
on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True)
the quantized weights, so the dense bf16 never touches the GPU.
Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching
GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same
torchao config + min_features filter the runtime path uses, applied ahead of time.
New core/inference/diffusion_prequant.py (resolve_prequant_source +
load_prequantized_transformer, best-effort, lazy imports). diffusion.py
_load_dense_quant_pipeline tries the pre-quant source first and falls back to the
dense materialise+quantise path, then to GGUF, so the default is unchanged.
DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty
prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for
the resolver, the meta-init+assign loader, and the backend branch selection +
fallbacks; GPU verification via scripts/verify_prequant_backend.py.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 10): attention-backend selection
Add a selectable attention kernel via the diffusers set_attention_backend
dispatcher. Attention is memory-bandwidth bound, so a better kernel is an
end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV
matmuls torchao never touches) and composes with torch.compile.
auto picks the best exact backend for the device: cuDNN fused attention
(_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x
end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default
(below the compile/quant noise floor); native SDPA elsewhere and when speed=off
(so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/
xformers/aiter are honored, and an unavailable kernel falls back to the default
rather than failing the load.
New core/inference/diffusion_attention.py (normalize + per-device select + apply,
best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline;
attention_backend threads through begin_load / load_pipeline / status like the other
load knobs. New request field attention_backend + status field. Hermetic CPU tests
for normalize / select policy / apply fallback, plus route threading + 422. Measured
via scripts/perf_levers_probe.py.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder
Consumer / workstation GPUs halve fp8 (and fp16/bf16) FP32-accumulate tensor-core
throughput, while int8 runs at full rate (int32 accumulate is not nerfed). Public
benchmarks (SDNQ across RTX 3090/4090/5090, AMD, Intel) confirm int8 via torch._int_mm
is as fast or faster than fp8 on every consumer part, and the only path on pre-Ada
consumer cards without fp8 tensor cores. So when transformer_quant=auto, reorder the
arch tier to put int8 first on a consumer/workstation GPU (detected by the existing
_is_consumer_gpu name heuristic), while data-center HBM parts keep fp8 first.
Pure ladder reorder via _prefer_consumer_scheme; no new flags. Verified non-regression
on a B200 (still picks fp8). Hermetic tests for consumer Blackwell/Ada/workstation
(-> int8) and data-center Ada/Hopper/Blackwell (-> fp8).
* Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT
Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across
denoise steps a DiT's output settles, so once the first block's residual barely
changes the remaining blocks are skipped and their cached output reused. diffusers
ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the
standalone apply_first_block_cache hook as a fallback).
Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 ->
2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar.
OFF by default and a per-load opt-in: the win scales with step count, so it is for
many-step models (Flux / Qwen-Image) and pointless for few-step distilled models
(e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction
of the trajectory. It composes with regional compile only with fullgraph=False (the
cache's per-step decision is a torch.compiler.disable graph break), which the speed
layer now switches to automatically when a cache is engaged. Best-effort: a model
whose block signature the hook does not recognise is caught and the load proceeds
uncached.
- new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache
(enable_cache / apply_first_block_cache fallback; threshold auto-raised for a
quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import).
- diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph
when a cache is engaged.
- diffusion.py: apply_step_cache before compile; thread transformer_cache /
transformer_cache_threshold through begin_load -> load_pipeline and report the
engaged mode in status().
- models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and
transformer_cache_threshold request fields, engaged mode in the status response.
- hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold
selection, and best-effort failure handling, plus route threading + validation.
- scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM /
LPIPS vs the compiled no-cache baseline).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 14): fix int8 dense quant on Flux / Qwen (skip M=1 modulation linears)
The opt-in dense int8 transformer path crashed on Flux.1 and Qwen-Image with
'torch._int_mm: self.size(0) needs to be greater than 16, but got 1'. int8 dynamic quant
goes through torch._int_mm, which requires the activation row count M > 16. A DiT's AdaLN
modulation projections (Flux norm1.linear 3072->18432, Qwen img_mod.1 / txt_mod.1, Flux.2
*_modulation.linear) and its timestep / guidance / pooled-text conditioning embedders are
computed once from the [batch, dim] conditioning vector (M = batch = 1), not per token, so
they hit _int_mm at M=1 and crash. Their feature dims are large, so the existing
min_features filter did not exclude them.
Fix: the int8 filter now also skips any Linear whose fully-qualified name matches a
modulation / conditioning-embedder token (norm, _mod, modulation, timestep_embed,
guidance_embed, time_text_embed, pooled). These layers run at M=1 once per block and are a
negligible share of the FLOPs, so int8 keeps the full speedup on the attention / FFN layers
(M = sequence length). fp8 / nvfp4 / mxfp8 use scaled_mm, which has no M>16 limit and
quantises these layers fine, so the exclusion is int8-only. Sequence embedders
(context_embedder / x_embedder / txt_in, M = seq) are deliberately not excluded -- note
'context_embedder' contains the substring 'text_embed', which is why the token is the
specific 'time_text_embed', not 'text_embed'.
Measured on a B200 (1024px, transformer_quant=int8 + speed=default), int8 now runs on every
supported model and is the fastest dense path on Flux/Qwen (int8 runs full-rate vs fp8's
FP32-accumulate): FLUX.1-dev 9.62s eager -> 1.98s (4.86x, vs fp8 2.15s), Qwen-Image -> 1.87s
(5.57x, vs fp8 2.09s), FLUX.1-schnell -> 0.41s (3.59x). Z-Image and Flux.2-klein (already
working) are unchanged.
- diffusion_transformer_quant.py: add _INT8_EXCLUDE_NAME_TOKENS; make_filter_fn takes
exclude_name_tokens; quantize_transformer passes it for int8 only.
- hermetic test that the int8 filter excludes the modulation / embedder linears (and keeps
attention / FFN / sequence-embedder linears), while fp8 keeps them.
- scripts/int8_linear_probe.py: the meta-device probe used to enumerate each transformer's
Linear layers and derive the exclusion list.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 15): build int8 pre-quantized checkpoints (skip M=1 modulation linears)
The prequant-checkpoint builder applied the dense quant filter without the int8-only
M=1 modulation / conditioning-embedder exclusion the runtime path uses, so a built int8
checkpoint baked those projections as int8 and crashed (torch._int_mm needs M>16) at the
first denoise step on Flux / Qwen. Factor the scheme->exclusion decision into a shared
exclude_tokens_for_scheme() used by both the runtime quantise path and the offline builder
so they can never drift, and apply it in build_prequant_checkpoint.py. int8 prequant now
produces a working checkpoint on every supported model, giving int8 (the consumer-preferred
scheme) the same ~2x load-VRAM and download reduction fp8 already had.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 16): route no-GPU loads to the native sd.cpp engine
When no CUDA/ROCm/XPU GPU is available, route diffusion load/generate to the
native stable-diffusion.cpp engine instead of diffusers, with diffusers as the
guaranteed fallback. On CPU sd.cpp is 1.4-2.8x faster and uses 1.5-2.2x less RAM.
- diffusion_engine_router: centralised engine selection (built on the existing
select_diffusion_engine), env opt-outs, MPS gating, recorded fallback reason.
- sd_cpp_backend (SdCppDiffusionBackend): the diffusers backend method surface
backed by sd-cli, with lazy binary install, registry-driven asset fetch,
step-progress parsing, and cancellation.
- diffusion_families: per-family single-file VAE + text-encoder asset mapping.
- sd_cpp_engine: cancellation support (process-group kill + SdCppCancelled).
- routes/inference + gpu_arbiter: drive the active engine via the router; the
API now reports the active engine and any fallback reason.
- tests for the backend, router, route selection, and cancellation.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Phase 16 review fixes: engine-switch unload, sd.cpp error mapping, per-image seeds, Qwen sampler
Address review feedback on #6724:
- engine router: unload the engine being deactivated on a switch, so the old
model is not left resident-but-unreachable (the evictor only targets the active
engine).
- generate route: sd.cpp execution errors (nonzero exit / timeout / missing
output) now map to 500, not 409 (which only means not-loaded / cancelled).
- native batch: return per-image seeds and persist the actual seed for each image
so every batch image is reproducible.
- Qwen-Image native path: apply --sampling-method euler --flow-shift 3 per the
stable-diffusion.cpp docs; other families keep sd-cli defaults.
- honor speed_mode (native --diffusion-fa) and, off-CPU, memory_mode/cpu_offload
offload flags on the native load instead of hardcoding them off.
- fail the load when the sd-cli binary is present but not runnable (version()
now returns None on exec error / nonzero exit).
- size estimate: only treat the transformer asset as a possible local path.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in
load_prequantized_transformer ends in torch.load(weights_only=False), which executes
arbitrary code from the pickle. The transformer_prequant_path load-request field reached
that unpickle for any local file an authenticated caller named, so a request could trigger
remote code execution. Refuse the source.kind=='path' branch unless the operator sets
UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted
and unaffected. Document the requirement on the API field and add gate tests.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa
- apply_attention_backend now restores the native default when no backend is requested or a
kernel fails. diffusers keeps a process-wide active attention backend that
set_attention_backend updates, and a fresh transformer's processors follow it, so a load
that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile
load pinned, breaking the bit-identical/off guarantee.
- select_attention_backend drops flash3/flash4 up front when the CUDA capability is below
Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit
request on the wrong card set fine then crashed mid-generation; it now falls back to native.
- Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a
valid alias of native) is accepted instead of 422-rejected by Pydantic.
- Drop the dead replace('-','_') normalization (no alias uses dashes/underscores).
- perf_levers_probe.py output dir is now relative to the script, not a hardcoded path.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF
- apply_step_cache now engages only via the transformer's native enable_cache (the diffusers
CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a
cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin
transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported
transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model
now runs uncached per the best-effort contract.
- GGUF transformers are quantized (the default Studio load path), so they now use the higher
quantized FBCache threshold when the caller leaves it unset, instead of the dense default
that could keep the cache from triggering.
- fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so
fullgraph=True failed warmup and silently measured an eager cached run); output dir is now
relative to the script, not a hardcoded path.
* Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder
_is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as
consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead
of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already
classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the
same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate.
Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0).
* Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config
Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just
TypeError on older torchao), which would break the mxfp8 config helper instead of falling
back to the default. Catch both so the fallback is robust.
quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds
CUDA memory during the per-row VRAM probe; output dir relative to the script.
* Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load
- snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform
missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the
whole snapshot. restore_backend_flags restores each flag independently so one failure can't
leave the others leaked process-wide.
- load_pipeline restores the flags (and clears the GPU cache) when the build fails after
apply_speed_optims mutated the process-wide flags but before _state captured them for unload
to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and
contaminated later off generations.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output
Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs
without producing output (or without closing stdout) would never reach proc.wait and the
wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the
PROCESS, so the main thread always enforces the timeout and kills a hung process (which
closes the pipe and ends the reader). Add a test that times out even when stdout blocks,
and make the no-binary test hermetic so a host-installed sd-cli can't leak in.
* Studio diffusion (Phase 14): guard the int8 exclusion filter against a None fqn
The filter callback can be invoked without a module name, so fqn.lower() would raise
AttributeError on None. Fall back to an empty name (nothing matches the exclusion tokens,
so the linear is kept) instead of crashing the quantise pass.
* Studio diffusion (Phase 16) review fixes: native engine robustness
- sd_cpp_backend: stop truncating explicit seeds to 53 bits (mask to int64);
a large requested seed was silently collapsed (2**53 -> 0) and distinct seeds
aliased to the same image. Random seeds stay 53-bit (JS-safe).
- sd_cpp_backend: sanitize empty/whitespace hf_token to None so HfApi/hf_hub
fall back to anonymous instead of failing auth on a blank token.
- sd_cpp_backend: a superseding load now cancels the in-flight generation, so the
old sd-cli can no longer return/persist an image from the previous model.
- diffusion_engine_router: run the previous engine's unload() OUTSIDE the lock so a
slow 10+ GB free / CUDA sync does not block engine selection.
- diffusion_engine_router: probe sd-cli runnability (version()) before committing to
native, so a present-but-unrunnable binary falls back to diffusers at selection.
- diffusion_device: resolve a torch-free CPU target when torch is unavailable, so a
CPU-only install can still reach the native sd.cpp engine instead of failing load.
- tests updated for the runnability probe + a not-runnable fallback case.
* Studio diffusion (Phase 9) review fixes: prequant safety + validation
- SECURITY: a request-supplied local pre-quant path is now unpickled only when it
resolves inside an operator-configured ALLOWLIST of directories
(UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in,
once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on
any path a load request named (arbitrary code execution). realpath() blocks symlink
escapes; a bare on/off toggle is no longer a wildcard.
- Validate the checkpoint's min_features against the runtime Linear filter, so a
checkpoint that quantised a different layer set is rejected instead of silently
loading a model that mismatches the dense path while reporting the same scheme.
- Tolerant base_model_id compare (exact or same final path/repo segment), so a local
path or fork of the canonical base is accepted instead of falling back to dense.
- _has_meta_tensors uses any(chain(...)) (no intermediate lists).
- prequant verify/probe scripts use repo-relative paths (+ env overrides), not the
author's absolute /mnt paths.
- tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts
- diffusion_memory: when group offload is unavailable and the plan falls back to
whole-module offload, enable VAE tiling (the group plan left it off, but the fallback
is the low-VRAM path where the decode spike can OOM). Covers both the group and
sequential fallback branches.
- perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a
balanced bit-identity regression actually fails the check instead of exiting 0.
- compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and
the load-progress poll has a 30 min deadline instead of looping forever on a hang.
- test for the group->model fallback enabling VAE tiling.
* Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path
- diffusion: a torchao-quantized transformer is committed only compiled. A dense model
resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF
it replaced), so when transformer_quant engaged and speed resolved to off, promote to
default (regional compile); warn loudly if compile still does not engage.
- diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the
CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK);
otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF.
- nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path).
- test asserts the eager-quant -> default-compile promotion.
* Studio diffusion (Phase 10) review fixes: attention gating + probe isolation
- diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere
NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation,
so auto now stays on native SDPA there.
- diffusion_attention: _active_attention_backend handles get_active_backend() returning an
enum/None (not a tuple); the old unpack always raised and was swallowed, so
the native-restore short-circuit never fired.
- perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS
on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul
so the inductor_flags variant doesn't leak into later compiled rows.
- tests for the SM80 cuDNN gate.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening
- install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit
timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket);
extract through a per-member containment check (Zip-Slip guard); expanduser the
--install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch
the separately-published cudart runtime DLL archive so sd-cli.exe can start.
- sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the
installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start
sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend
crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie.
- tests: Zip-Slip rejection, normal extraction, studio-home discovery.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs
Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes
the numbered files <stem>_<idx><suffix> (base_0.png, base_1.png, ...) instead of
the literal --output path. SdCppEngine.generate checked only the literal path, so
a batch generation would exit 0 and then raise 'no image' (or return a stale
file). generate now returns the literal path when present and otherwise falls
back to the numbered siblings; single-image behavior is unchanged.
Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected
without error.
* Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats
Codex review on the native engine arg builder:
- build_sd_cpp_command emitted --width/--height unconditionally, so an
img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of
the input. width/height are now Optional (None = unset): an image-conditioned
run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives
the size from the input image (set_width_and_height_if_unset); a plain txt2img
run with unset dims keeps the prior 1024x1024 default; explicit dims are always
honored. width/height are read only by the builder, so the type change is local.
- build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...)
that silently swallowed repeats=0 into sd-cli's default of one pass, turning an
explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError
and emits the flag for any explicit value != 1.
Tests: img2img unset dims omit width/height (init_img and ref_images), explicit
dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at
the default. (Two pre-existing binary-discovery tests fail only because a real
sd-cli is installed in this dev environment; unrelated to this change.)
* Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc
Codex review: the transformer_prequant_path field description still told operators
to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the
prior security fix made that variable a directory allowlist -- _allowed_prequant_roots
deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator
following the documented =1 would have every transformer_prequant_path request
silently refused. The description now states it must name one or more allowlisted
directories and that a bare on/off value is not accepted.
Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does
not say =1, and describes an allowlist/directory (guards against doc drift).
* Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset
Codex review on attention-backend selection:
- Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on
pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first
generation with no fallback. select_attention_backend now applies
_cudnn_attention_supported() to an explicit cuDNN request too.
- flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a
Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3
is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a
(min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+.
- apply_attention_backend's success path left diffusers' process-wide active
backend pinned to the kernel it set; a later component whose processors are
unconfigured (backend None) would inherit it. It now resets the global registry
to native after a successful per-transformer set (the transformer keeps its own
backend), best-effort. Also fixed _active_attention_backend: get_active_backend()
returns a (name, fn) tuple, so the prior code stringified the tuple and never
matched a name, defeating the native-restore short-circuit.
Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on
SM90; global registry reset after a successful set; _active_attention_backend
reads the tuple return.
* Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder
Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300
(Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So
_is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the
auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting
llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they
are treated as data-center class and keep the intended fp8-first behavior.
Test: extends the datacenter parametrize with 'NVIDIA B300' and
'NVIDIA GH200 480GB' (now _is_consumer_gpu False).
* Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder
Codex review: the M=1 modulation/embedder exclusion was wired only into the dense
runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called
make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint
quantised the AdaLN modulation and conditioning-embedder linears, and loading it
via transformer_prequant_path (the load path only loads already-quantised tensors,
it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes
for the runtime path.
Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 ->
the M=1 exclusion, every other scheme -> none) and use it in both the runtime
quantiser and the builder, so a prequant artifact's quantised-layer set always
matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion).
Test: int8_exclude_name_tokens returns the exclusion for int8 and () for
fp8/nvfp4/mxfp8.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 16) review round 2: native CPU arbiter, status offload, load race
Codex review on the native-engine routing:
- The /images/load route took the GPU arbiter (acquire_for(DIFFUSION) -> evict chat)
unconditionally after engine selection. A native sd.cpp load on a pure-CPU host
never touches the GPU, so that needlessly tore down the resident chat model. The
handoff is now gated: diffusers always takes it, a force-native sd.cpp load on a
CUDA/XPU/MPS box still takes it, but a native sd.cpp load on a CPU host skips it.
- sd_cpp status() hardcoded offload_policy 'none' / cpu_offload False even when
_run_load computed real offload flags (balanced/low_vram/cpu_offload off-CPU), so
the setting was unverifiable. status now derives them from state.offload_flags
(still 'none' on CPU, where the flags are empty).
- _run_load committed the new state without cancelling/waiting on a generation that
started during the (slow) asset download, so a stale sd-cli run against the OLD
model could finish afterward and persist an image from the previous model once the
new load reported ready. The commit now signals the in-flight cancel and waits on
_generate_lock before swapping _state (taken only at commit, so the download never
serialises against generation), mirroring the diffusers load path.
Tests: CPU native load skips the arbiter while a GPU native load takes it; status
reports offload active when flags are set; _run_load cancels and waits for an
in-flight generation before committing.
* Studio diffusion (Phase 14) review round 2: align helper name with the stack
Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the
identical helper already present higher in the diffusion stack (Phase 16). The
helper definition, the runtime quantiser call, and the offline builder are now
byte-identical to that version, so the two branches no longer introduce a
divergent name for the same single-source-of-truth and the stack merges without
a conflict on this fix. No behavior change.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion: eager patches + torch.compile cache speed phase
Adds the opt-in speed path for the GGUF diffusion transformer behind a
selectable speed mode (default off, so output is unchanged until a profile
is chosen):
- diffusion_eager_patches.py: shared eager fast-paths (channels_last,
attention/backend selection, fused norms and QKV) installed at load and
rolled back on unload or failed load.
- diffusion_compile_cache.py / diffusion_gguf_compile.py: a persistent
torch.compile cache and the GGUF-transformer compile wiring.
- diffusion_arch_patches.py: architecture-specific patches.
- diffusion_patch_backend.py: shared install/restore plumbing.
- diffusion_speed.py: speed-profile planning.
Tests for each module plus the benchmarking and probe scripts used to
measure speed, memory, and accuracy of the path.
* Studio diffusion: image workflows (safetensors, image-conditioned, editing) + Images UI
Backend:
- Load non-GGUF safetensors models: full bnb-4bit pipelines and single-file
fp8 transformers, gated to the unsloth org plus a curated allowlist.
- Image-conditioned workflows built with Pipeline.from_pipe so they reuse the
loaded transformer/VAE/text-encoder with no extra VRAM: img2img, inpaint,
outpaint, and a hires-fix upscale pass.
- Instruction editing as its own family kind (Qwen-Image-Edit-2511,
FLUX.1-Kontext-dev) and FLUX.2-klein reference conditioning (single and
multi-reference) plus klein inpaint.
- Auto-resize odd-sized inputs to a multiple of 16 (and resize the matched
mask) so img2img/inpaint/edit no longer reject non-/16 uploads. Bound the
decoded image size and cap upscale output to avoid OOM on large inputs.
- Fixes: from_pipe defaulting to a float32 recast that crashed torchao
quantized transformers; image-conditioned calls forcing the slider size
onto the input image. Native sd.cpp engine rejects image-conditioned and
reference requests it cannot serve.
Frontend:
- Redesigned Images page with capability-gated workflow tabs (Create,
Transform, Inpaint, Extend, Upscale, Reference, Edit), a brush mask editor,
client-side outpaint, and a multi-reference picker.
- Advanced options moved to a right-docked panel mirroring Chat: closed by
default, toggled by a single fixed top-bar button that stays in place.
sd.cpp installer: pin the release, verify each download's sha256, add a
download timeout, and make the source repo configurable for a future mirror.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio Images: correct the Advanced panel comment (closed by default, fixed toggle)
* Studio diffusion: LoRA adapters for the Images workflow
Add community LoRA support across both diffusion backends, the single
biggest step toward broad image-workflow coverage.
Backend
- New shared module core/inference/diffusion_lora.py: adapter discovery
(local scan + curated catalog + owner/name[:file] Hub refs), download
via hf_hub_download_with_xet_fallback, alias sanitization, native
managed-dir materialization with collision-broken aliases, prompt-tag
injection (deduped against user-typed tags), and a supports_lora gate.
- Native sd-cli: resolve + materialize selected LoRAs into a per-run
managed dir, inject <lora:ALIAS:w> tags, pass --lora-model-dir with
--lora-apply-mode auto. The arg builder already emitted these flags.
- Diffusers: non-fused load_lora_weights + set_adapters manager, tracked
on the pipe so an unchanged selection is a no-op and a model swap
resets; cleared on unload. Never fuses (breaks quantized transformers
and blocks live weight tweaks).
- Gated off where unsupported: torchao fp8/int8 dense, GGUF-via-diffusers,
and native Qwen-Image (no LoRA name-conversion branch upstream).
- Request contract: optional loras on DiffusionGenerateRequest; empty or
omitted is identical to today. supports_lora surfaced in status; chosen
LoRAs persisted in gallery recipe metadata.
- New GET /api/models/diffusion-loras for the picker (family-filtered).
Frontend
- Repeatable multi-LoRA picker (adapter select + weight slider 0..2 +
remove), gated by the loaded model's supports_lora and family, max 8.
Tests
- New test_diffusion_lora.py (14): helpers, request validation, native
tag/dir wiring, diffusers set_adapters manager, supports_lora matrix.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion: ControlNet for the Images workflow (diffusers)
Add ControlNet conditioning, the #2 most-used diffusion workflow after
LoRA, on the diffusers backend for the families with ControlNet pipelines
(FLUX.1 and Qwen-Image), with Union models as the default picks.
Backend
- New core/inference/diffusion_controlnet.py: family-gated discovery
(curated Union models + local dirs + bare owner/name repos), resolution
to a loadable repo/dir, control-image preprocessing (passthrough +
a dependency-free canny edge map), and a supports_controlnet gate.
- diffusion.py: a ControlNet manager parallel to the LoRA one. Loads the
(small) ControlNet model once via from_pretrained (cached by id) and
builds the family's ControlNet pipeline via Pipeline.from_pipe(base,
controlnet=model), reusing the resident base modules at their loaded
dtype (no reload, no recast). Passes the control image + conditioning
scale + guidance start/end at generate time; cleared on unload.
- Families: FLUX.1 -> FluxControlNetPipeline/Model, Qwen-Image ->
QwenImageControlNetPipeline/Model. Others declare none (gated off).
- Gated off for the native engine, GGUF-via-diffusers, and torchao
fp8/int8 dense (same rule as LoRA). v1 conditions txt2img only.
- Request contract: optional controlnet on DiffusionGenerateRequest;
supports_controlnet in status; the choice persisted in gallery meta.
- New GET /api/models/diffusion-controlnets for the picker.
Frontend
- A ControlNet control in the Images rail (model select + control-image
upload + control-type select + strength slider), gated by the loaded
model's supports_controlnet + family, shown for text-to-image.
Tests
- New test_diffusion_controlnet.py (10): discovery/resolve/preprocess/gate
helpers, request validation, family wiring, and the diffusers pipe
manager (loads once, caches, from_pipe with controlnet, rejects
unsupported families).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio ControlNet: show the picker on the Create tab (workflow id is 'create', not 'txt2img')
The ControlNet control gated on workflow === 'txt2img', but the Images workflow tab ids are create/transform/inpaint/extend/upscale/reference/edit -- there is no 'txt2img'. So the picker never rendered even with a ControlNet-capable model loaded. Gate on 'create' (the text-to-image tab) for both the picker and the request wiring. Found via a live Playwright capture of the running Studio.
* Studio: do not force diffusers pipelines cross-tagged gguf into the GGUF variant expander
Some diffusers image repos (e.g. unsloth/Qwen-Image-2512-unsloth-bnb-4bit) carry a
stray "gguf" tag on the Hub but ship no .gguf files. The model search classified
them as GGUF from the bare tag, so the picker rendered the GGUF variant expander,
which then dead-ended at "No GGUF variants found." Trust the bare gguf tag only when
the repo is not a diffusers pipeline; the -GGUF name suffix and real gguf metadata
(populated via expand=gguf) remain authoritative, so genuine GGUF repos are unaffected.
* Studio Images: load non-curated unsloth/on-device diffusers repos instead of no-op
handleModelSelect only loaded curated safetensors ids and GGUF variant picks; any other
non-GGUF pick (an on-device diffusers folder, or a future unsloth diffusers image repo
surfaced by search) silently did nothing. Treat such a pick as a full diffusers pipeline
load when the id is unsloth-hosted or on-device (the backend infers the family + base repo
and gates loads to unsloth/* or local paths), and show a clear message otherwise instead
of silently ignoring the click. Curated and GGUF paths are unchanged.
* Studio Images: keep curated safetensors models in Recommended after download
The curated bnb-4bit / fp8 diffusion rows were filtered out of the Images picker's
Recommended list once cached (curatedSafetensorsRows dropped anything in downloadedSet),
so they vanished from the picker after the first load and could only be found by typing an
exact search. The row already renders a downloaded badge, matching how GGUF Recommended
rows stay visible when cached. Drop the exclusion so the curated safetensors always list.
* Studio diffusion LoRA: sanitize dots out of adapter aliases
The LoRA alias is used as the diffusers PEFT adapter name, and PEFT rejects names
containing "." (module name can't contain "."). sanitize_alias kept dots, so a LoRA whose
filename carries a version tag (e.g. Qwen-Image-2512-Lightning-8steps-V1.0-bf16) failed to
apply with a 400. Replace dots too; the alias stays a valid native <lora:NAME:w> filename
stem. Adds regression coverage for internal dots.
* Studio Images: clarify the GGUF transformer-quant Advanced control
Renamed the confusing "Transformer quant / GGUF default" control to "GGUF speed mode"
with an "Off (run the GGUF)" default, and reworded the hint to state plainly that FP8/INT8/
FP4 load the FULL base model (larger download + more VRAM) rather than re-packing the GGUF,
falling back to the GGUF if it can't fit. Behavior unchanged; labels/hint only.
* Studio Images: list on-device unsloth diffusion models in the picker
The Images picker's On Device tab hid every non-GGUF cached repo whenever a
task filter was active, so downloaded unsloth diffusion pipelines (bnb-4bit
and FP8 safetensors) never showed up there. List cached repos that pass the
task gate, limited under a filter to unsloth-hosted ones so base repos (which
fail the diffusion load trust gate) don't appear only to dead-end on click.
Chat behavior is unchanged: the task gate still drops image repos there.
* Studio: hide single-file image checkpoints from the chat model picker
The chat picker treats a cached repo as an image model, and hides it, only
when it ships a diffusers model_index.json. Single-file, ComfyUI, and
ControlNet image checkpoints (an FP8 Qwen-Image, a z-image safetensors, a
Qwen-Image ControlNet) carry none, so they surfaced as loadable chat models.
Fall back to resolving the repo id against the known diffusion families, the
same resolver the Images backend loads from, so these checkpoints are tagged
text-to-image and stay in the Images picker only.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio Images: add the FLUX.2-dev model family
Loading unsloth/FLUX.2-dev-GGUF failed because detect_family knew only the
Qwen3-based FLUX.2-klein, so FLUX.2-dev (the full, Mistral-based Flux2Pipeline)
resolved to nothing and the load errored. Add a flux.2-dev family: Flux2Pipeline
+ Flux2Transformer2DModel over the black-forest-labs/FLUX.2-dev base repo (gated,
reachable with an HF token), with its FLUX.2 32-channel VAE and Mistral text
encoder wired for the sd-cli path from the open Comfy-Org/flux2-dev mirror.
text-to-image only: diffusers 0.38 ships no Flux2 img2img / inpaint pipeline for
dev. Frontend gets sensible dev defaults (28 steps, guidance 4), distinct from
klein's turbo defaults. Verified live: GGUF load resolves the family + gated base
repo and generates a real 1024x1024 image on GPU.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio Images: clearer error for an unsupported diffusion model
When a repo id resolves to no diffusion family the load raised 'Could not infer a
diffusion family... Pass family_override (z-image)', which points at an unrelated
family and doesn't say what is supported. Replace it with a message that lists the
supported families (from a new supported_family_names helper) and notes that video
models and image models whose diffusers transformer has no single-file loader are
not supported. Applies to both the diffusers and native sd.cpp load paths. Also
refreshes two stale family-registry comments that still called FLUX.2-dev omitted.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add SDXL diffusion family (U-Net pipeline support)
SDXL is the first U-Net family in the diffusion backend: its denoiser is
pipe.unet (UNet2DConditionModel), not a DiT pipe.transformer, and a single-file
.safetensors is the whole pipeline rather than a transformer-only file. The
backend previously assumed a DiT transformer everywhere, so add the two hooks a
U-Net family needs and register SDXL.
DiffusionFamily gains denoiser_attr ("transformer" for DiT, "unet" for SDXL) and
single_file_is_pipeline (SDXL loads a single file via pipeline_class.from_single_file
with the base repo as config, instead of transformer_class.from_single_file plus a
companion assembly). _align_vae_dtype now reads the denoiser generically so img2img
and inpaint keep the VAE and U-Net dtypes aligned.
The non-GGUF trust gate is extended with a short, exact-match, safetensors-only
allowlist of official base repos (the SDXL base/refiner and sdxl-turbo), because
SDXL ships only as a full pipeline and has no unsloth-hosted GGUF. Local paths stay
trusted as before; a random repo, even one that detects as SDXL, is still rejected.
The image-conditioned and ControlNet workflows are the standard SDXL pipelines,
built around the resident modules via from_pipe like every other family, so SDXL
gets txt2img, img2img, inpaint, outpaint, upscale, LoRA and ControlNet. There is no
native sd.cpp mapping yet, so the no-GPU route falls back to diffusers.
Frontend catalog gains SDXL Base 1.0 and SDXL Turbo entries with SDXL step/guidance
defaults (Turbo: few steps, no CFG; base: ~30 steps, real CFG).
Tests: new test_diffusion_sdxl.py (family shape, detection, trust allowlist, model
kind, U-Net VAE-dtype alignment, LoRA gate) plus loader-branch tests in
test_diffusion_backend.py (pipeline-kind from_pretrained, single-file whole-pipeline
from_single_file, allowlist accept/reject). Verified live on GPU: sdxl-turbo loads
both as a pipeline and as a single file and generates coherent txt2img + img2img.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Images: LoRA free-text Hugging Face entry + recipe round-trip
The backend has always accepted a bare Hugging Face repo id (owner/name, or
owner/name:weight-file.safetensors) as a LoRA, downloading and applying it. But the
picker only rendered when the curated catalog had entries, and the catalog is empty,
so there was no UI path to apply any LoRA. Show the LoRA section whenever the loaded
model supports LoRA, and replace the curated-only dropdown with a text input: type a
Hub repo id, or pick a discovered adapter from a datalist of suggestions when the
catalog is populated.
Also restore LoRAs when loading a recipe. restoreSettings now parses the recipe's
"id:weight" strings (splitting on the last colon, since the id itself may contain one
for a specific weight file) back into the selection, so replaying a saved image
reproduces its adapters. The generate payload trims hand-typed ids and drops empty /
zero-weight rows, and a model swap clears the selection (a LoRA is family-specific)
without discarding a free-text pick that is not in the curated list.
* Add diffusion LoRA training (SDXL text-to-image)
First diffusion training path in Studio: train a LoRA on the SDXL U-Net from an
image + caption dataset and export it as a diffusers .safetensors that the existing
diffusion LoRA loader (and any diffusers pipeline) can load.
core/training/diffusion_lora_trainer.py:
- DiffusionLoraConfig with validation/defaults (rank, alpha, targets, lr, steps, grad
accumulation, resolution, min-SNR gamma, gradient checkpointing, lr scheduler, seed,
mixed precision).
- discover_image_caption_pairs: captions from metadata.jsonl / captions.jsonl, per-image
.txt/.caption sidecars, or a dreambooth instance_prompt fallback (pure, unit-tested).
- run_diffusion_lora_training: the loop -- freeze base, PEFT-wrap the U-Net attention
projections, VAE-encode (fp32 VAE to avoid the SDXL fp16 overflow), sample noise +
timesteps, predict, MSE loss with optional min-SNR weighting (epsilon / v-prediction),
AdamW + get_scheduler + grad accumulation + grad clipping, then export via
save_lora_weights. Emits worker-protocol events (model_load_*, progress, complete) and
polls should_stop for a clean stop with a partial save.
- run_diffusion_training_process: mp.Queue subprocess adapter (event_queue / stop_queue),
so the training worker can spawn it; plus a CLI entry point.
Only SDXL (U-Net) is trained here; DiT families and the Studio UI form + route wiring are
follow-ups. The trainer is decoupled and worker-ready.
Tests: test_diffusion_lora_trainer.py covers caption discovery (metadata / sidecar /
instance prompt / skip-uncaptioned / errors), config normalisation + validation, the SDXL
add-time-ids, and the dict->config adapter. Verified live on GPU: a 60-step SDXL LoRA run
lowers the loss, exports a ~45 MB adapter, and loading it back shifts generation from
baseline (mean abs pixel diff ~55/255).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* diffusion trainer: emit learning_rate in progress events (Studio pump compatibility)
The Studio training pump reads 'learning_rate' from progress events; the diffusion
trainer emitted 'lr'. Rename the field (and the CLI reader) so the trainer's events are
directly consumable by the existing training status/SSE machinery when it is wired into
the worker, without a translation shim.
* Wire diffusion LoRA training into the Studio API
Make the SDXL LoRA trainer reachable from the app with a small, self-contained job
service and JSON routes, deliberately separate from the LLM TrainingBackend (whose
lifecycle -- LLM config build, per-run SQLite rows, matplotlib plots, transfer-to-chat-
inference -- is text-training specific and would mis-handle a diffusion run).
core/training/diffusion_training_service.py: DiffusionTrainingService runs one job at a
time -- validate the config cheaply (before any spawn), spawn the trainer subprocess
(spawn context, parent-lifetime bound), pump its events (model_load_* / progress /
complete / error) into an in-memory status snapshot, and support a clean stop. The
subprocess context and target are injectable so the full start -> pump -> status ->
complete path is unit-tested without real multiprocessing or torch.
routes/training.py: POST /api/train/diffusion/start (400 on a bad config, 409 when a job
is already running), POST /api/train/diffusion/stop, GET /api/train/diffusion/status
(JSON poll). models/training.py: DiffusionTrainingStartRequest + response schemas
mirroring DiffusionLoraConfig, so model_dump() passes straight through.
Tests: test_diffusion_training.py -- service happy path, bad-config-before-spawn,
concurrent-job rejection, clean stop, crash-without-terminal-event, event transitions;
plus route wiring via the FastAPI TestClient (start / 422 / 400 / 409 / status / stop)
with a mocked service. The diffusion trainer's progress events already use the field
names this path expects.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Images: add a Train LoRA (SDXL) dialog
Surface the diffusion training API in the Images page. A "Train LoRA" button in the top
bar opens a self-contained dialog to fine-tune an SDXL LoRA on a folder of images: pick
the base model, dataset folder, output folder, an optional instance prompt, and the core
hyperparameters (steps, rank, resolution, batch, learning rate), then Start. The dialog
polls the training status while open and shows a progress bar, step count, live loss, and
the saved adapter path, with a Stop button for a clean stop.
The dialog is independent of the loaded generation model (training runs in its own
subprocess), and prefills the base model with the loaded checkpoint when it is SDXL, else
the SDXL base. api.ts gains startDiffusionTraining / stopDiffusionTraining /
getDiffusionTrainingStatus plus their types, matching the /api/train/diffusion routes.
* Import diffusion training schemas from models.training directly
The import-hoist lint flags newly re-exported names in the models/__init__.py hub as
unused (it does not treat __all__ membership as a use). Import the three diffusion
training schemas straight from models.training in routes/training.py, where they are
used in the route annotations and calls, and drop the __init__ re-export.
* Remove stray async task scratch outputs committed by mistake
* ControlNet: reject filesystem-like ids and do not cache a model past an unload race
Two review findings on the ControlNet path:
- resolve_controlnet's bare-repo fallback accepted any id with a slash, so a
path-shaped id (/tmp/x, ../x) reached from_pretrained as a local directory.
Restrict the fallback to a strict owner/name HF repo id shape.
- _controlnet_pipe now re-checks the cancel event after the blocking
from_pretrained: an unload that raced the download had already cleared the
caches, so caching the late module would pin it past the unload.
* Pipeline prefetch: fetch only the default torch weights
A full-pipeline prefetch kept every repo file outside assets/, so an official
repo that ships multiple formats (SDXL Base: fp16 variants, ONNX, OpenVINO,
Flax, a top-level single-file twin) downloaded tens of GB from_pretrained never
loads. Skip non-torch exports and dtype-variant twins in
_pipeline_file_downloaded, and drop a component .bin when the same directory
carries a picked safetensors weight (diffusers' own preference).
* Diffusion LoRA training: fall back to fp16 when CUDA lacks bf16
The default mixed_precision=bf16 hard-fails on pre-Ampere GPUs (T4 / V100 /
RTX 20xx) which have no bf16 compute; check torch.cuda.is_bf16_supported()
and drop to fp16 there.
* Diffusion training service: join the old pump outside the lock
start() joined a finished job's pump thread while holding the service lock,
but the pump's final state writes need that same lock, so the join always
burned its full timeout and a stale pump could then overwrite the new job's
state. Join outside the lock (with a re-check after), and fence _apply_event
and the exit handler by process identity so a superseded pump can never touch
the current job's state. Adds regression tests for both.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion LoRA training: harden config handling, cancellation, SDXL conditioning, and safety
Addresses review findings on the SDXL LoRA trainer:
- Gate the base model with the same trust check as inference (unsloth/*, allowlisted
official bases, or a local path) before from_pretrained, so an untrusted remote repo
is never fetched or deserialised.
- Check the stop signal before the (slow) model load, not only between steps, so a
cancel during download is honoured; a stop may carry save=False to cancel without
leaving a partial adapter.
- Per-sample SDXL add_time_ids from the actual crop (original size + crop offset, with
the offset mirrored on horizontal flip) instead of a fixed uncropped-square tensor.
- Apply EXIF orientation before resize/crop so rotated photos train upright.
- Skip gradient clipping when max_grad_norm <= 0 (the Studio 'disable' value) instead
of scaling every gradient to zero.
- Coerce Studio config strings/blanks: learning_rate string to float, blank hf_token to
anonymous, gradient_checkpointing 'none'/'true'/'unsloth' to bool; reject a zero/negative
lora_alpha or learning_rate.
- Alias the generic Studio training payload keys (model_name/max_steps/batch_size/lora_r/
lr_scheduler_type/random_seed) onto the diffusion field names.
- Mirror the trained adapter into loras/diffusion so the Images LoRA picker discovers it.
- Report worker exceptions in both message and error keys so the failure is not lost.
Adds regression tests for the config coercion/validation and aliasing.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* ControlNet: address review findings on the diffusers path
- resolve_controlnet enforces catalog family compatibility so a direct API call
cannot load a ControlNet built for another family through the wrong pipeline.
- Unknown ControlNet ids now surface as a 400 (call site maps FileNotFoundError
to ValueError) instead of a generic 500.
- strength 0 disables ControlNet entirely, so a no-op selection never pays the
download / VRAM cost; the control image is decoded and validated BEFORE the
ControlNet is resolved or built, so a malformed image fails fast for the same reason.
- ControlNet loads use the base compute dtype (state.dtype is a display string,
not a torch.dtype, so it silently fell back to float32) and honor the base
offload policy via group offloading instead of forcing the module resident.
- Empty/malformed HF token coerced to anonymous access.
- Flux Union ControlNet control_mode mapped from the selected control type.
- resolve_controlnet drops the unused hf_token/cancel_event params.
- ControlNetSpec validates guidance_start <= guidance_end (clean 422).
- Images UI ControlNet Select shows its placeholder when nothing is selected.
Adds regression tests for family enforcement and the union control-mode map.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion training API: LLM interlock, pre-spawn VRAM free, path containment, no dropped knobs
Four review findings on the diffusion training start path:
- It spawned the SDXL trainer without checking the LLM TrainingBackend, so a
start while an LLM run was active put two trainers on the same GPU. Add a
symmetric interlock: diffusion start returns 409 when LLM training is active,
and LLM start refuses while a diffusion job is active.
- It went straight to service.start() without freeing GPU residents. Add a
pre-spawn free of the export subprocess, the resident Images pipeline (with an
arbiter release), and chat models, mirroring the LLM start path.
- data_dir / output_dir were passed through unresolved, so Studio-relative names
failed and absolute paths bypassed containment. Resolve them with
resolve_dataset_path / resolve_output_dir before spawn (400 on an uncontained
path).
- The request model dropped max_grad_norm and lora_target_modules, so runs that
set them trained with defaults. Add both fields.
The gemini pump-join deadlock was already fixed earlier (join outside the lock +
proc-identity fence). Note: honoring a stop DURING model load is a trainer-loop
change owned by the diffusion training engine PR (should_stop polled before the
first optimizer step). Adds route + model regression tests.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion LoRA: harden resolution, native tag precedence, and diffusers teardown
Address review findings on the LoRA path:
- resolve_one: normalise a blank/whitespace hf_token to None (anonymous access)
and reject a client-supplied weight file with traversal / absolute path.
- resolve_specs: convert FileNotFoundError from an unknown/stale id to ValueError
so the route returns 400 instead of a generic 500.
- _scan_local: disambiguate local adapters that share a stem (foo.safetensors vs
foo.gguf) so each is uniquely addressable.
- inject_prompt_tags: the backend-validated weight now wins over a user-typed
<lora:ALIAS:...> for a selected adapter; unselected user tags are left alone.
- diffusers _apply_loras: reject a .gguf adapter with a clear error before touching
the pipe (diffusers loads safetensors only).
- _unload_locked: drop the explicit unload_lora_weights() on teardown; the pipe is
dropped wholesale (freeing adapters), so the previous call could race an in-flight
denoise on the same pipe.
- Images page: use a stable LoRA key and clear the selection (not just the options)
when the catalog refresh fails.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion: guard trust check against OSError and validate conditioning inputs
- _is_trusted_diffusion_repo: wrap Path.exists() so a repo id with invalid
characters (or a bare owner/name id) can't raise OSError; treat any failure as
not-a-local-path and fall through to the unsloth/ allowlist. validate_load_request
still raises the clear FileNotFoundError for a genuinely missing local pick.
- generate(): reject mask_image / upscale / reference_images supplied without an
input image, and reject reference_images on a family that does not support
reference conditioning, instead of silently degrading to txt2img / img2img.
* SDXL: reject GGUF up front, skip unused base weights, drop refiner, and harden helpers
Addresses review findings on the SDXL family:
- Reject a GGUF load for single_file_is_pipeline families (SDXL) in validate_load_request,
before the route evicts the current model; SDXL has no transformer-only GGUF variant.
- Skip base-repo weight files when a whole-pipeline single file is loaded: from_single_file
(config=base) needs only the base config/tokenizer/scheduler, so a local .safetensors no
longer triggers a multi-GB base download.
- Remove the SDXL refiner from the non-GGUF trust allowlist: it is an img2img-only pipeline
but this backend loads every sdxl repo as the base txt2img pipeline.
- Normalize a blank/whitespace hf_token to None once in load_pipeline so every load branch
degrades to anonymous instead of erroring on a malformed token.
- Read the denoiser dtype from a parameter (compile-wrapped modules may lack .dtype) and
access state.family.denoiser_attr directly.
Adds/updates regression tests for the trust allowlist, GGUF rejection, and base-config filter.
* Images: preserve restored LoRAs through model load and never send hidden LoRAs
- The LoRA effect cleared the selection on every load->capable transition, which
wiped adapters restored from a gallery recipe before the model finished loading.
Track the previously-loaded family in a ref and clear only on a real family swap;
keep the selection on the initial load and on unload.
- Gate the generate payload's loras on loraCapable so a restored selection that is
hidden (loaded model does not support LoRA) is never sent to the backend.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Images Train LoRA dialog: token, validation, precision, base-repo prefill, gating, refresh
Nine review findings on the SDXL training dialog:
- Forward the saved Hub token so a gated/private SDXL base can be trained (the
image load flow already sends it).
- Re-seed the base-model field from the current default each time the dialog
opens; the keep-alive dialog otherwise kept its mount-time default after a
model loaded.
- Prefill from base_repo (the diffusers pipeline) rather than repo_id, which for
a GGUF/single-file SDXL load is the checkpoint path from_pretrained can't open.
- Add client-side validation of steps/rank/resolution/batch/learning-rate before
the request.
- Expose a precision selector (bf16/fp16/fp32) so non-bf16 GPUs can train from
the UI, not only the API.
- Gate the dialog on the active Images route (active && trainOpen) so switching
tabs closes it and stops its polling.
- Rescan the LoRA picker when a run completes, so a freshly-trained adapter
appears without a model reload.
- Cap the dialog height and scroll the body so the Start/Stop footer stays
reachable on short viewports.
- Correct the copy to not over-promise picker auto-discovery.
Freeing the resident Images pipeline before training is handled backend-side in
the diffusion training start route.
* Merge diffusion-sdxl into diffusion-lora-ux; keep options-only LoRA catch
The catalog-refresh .catch from the lower branch clears the selected adapters
too, which is right for its catalog-only picker but wrong here: this picker
holds free-text HF repo ids that are valid without being in the catalog, so a
transient refresh failure must not wipe them. Family swaps still clear the
selection and hidden LoRAs are never sent.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Train LoRA dialog: stop suggesting absolute paths the backend rejects
The dataset and output placeholders showed /path/to/... examples, but the
training routes resolve those fields inside the Studio home and reject
absolute paths outside the approved roots, so following the placeholder
produced a 400. Use folder-name placeholders and say in the labels and the
dialog description where each folder resolves.
* Align the VAE to the denoiser's first FLOATING dtype, not its first parameter
A GGUF-quantized transformer's leading parameters are packed uint8 storage,
so reading next(parameters()).dtype handed nn.Module.to() an integer dtype
and every image-conditioned generation on a GGUF model (Qwen-Image-Edit)
failed with a 500. Probe the parameters for the first floating dtype, treat
an all-integer module as a no-op, and also catch TypeError so an unexpected
dtype can never break generation. Regression test included.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Count LR scheduler warmup/decay in optimizer steps, not micro-steps
lr_sched.step() runs once per outer optimizer step (after the gradient
accumulation inner loop), for train_steps total. The scheduler was
configured with num_warmup_steps and num_training_steps multiplied by
gradient_accumulation_steps, so with accumulation > 1 a warmup or
non-constant schedule stretched past the run and never reached the
intended decay. Count both in optimizer steps.
* Address Codex review findings on the image-workflows PR
Keep diffusion.py importable without torch: the compile/arch patch modules
import torch at module level, so import them lazily at their load/unload
call sites instead of at module load. This restores the torchless contract
so get_diffusion_backend() works on a CPU/native sd.cpp install.
Match family reject keywords and aliases as whole path/name segments, not
raw substrings, so an unrelated word like edited, edition, or kontextual no
longer misroutes or hides a valid base image model, while supported edit
families (Qwen-Image-Edit, FLUX Kontext) still resolve. Mirror the same
segment matching in the picker task filter.
Route FLUX.2-dev native guidance through --guidance like the other FLUX
families rather than --cfg-scale. Reject native upscale requests that have
no input image. Read image header dimensions and reject over-limit inputs
before decoding pixels, so a crafted small-payload image cannot spike
memory. Reject an upscale that would shrink the source below its input
size. Validate the model_kind against the filename extension before the
GPU handoff. Estimate a local diffusers pipeline's size from its on-disk
weights so auto memory planning does not skip offload and OOM. Report
workflows: [txt2img] from the native backend status so the Create tab
stays enabled for a loaded native model. Clamp the outpaint canvas to the
backend's 4096px decode limit.
Adds regression tests for segment matching and kind/extension validation.
* Guard inference loads and worker lifetime against diffusion training
Teach the chat and image load guards about an active diffusion (SDXL) LoRA
job: a chat load is refused (its footprint cannot be fit-checked against the
trainer) and an image load is refused outright, mirroring the existing LLM
training guards, so a load can no longer allocate GPU memory alongside the
trainer and undo the pre-start cleanup.
Bind the diffusion trainer subprocess to the parent's lifetime and scrub the
native path lease secret from it by running the child through
run_without_native_path_secret, matching the inference/export/LLM workers, so
a Studio crash or kill no longer leaves the trainer holding the GPU.
Reset in_model_load on the complete and error terminal events: a stop or
failure during model loading otherwise leaves the status reporting a stale
loading indicator after the job has ended.
* Harden diffusion LoRA handling on the diffusers and native paths
Reject LoRA on a torch.compile'd diffusers transformer (Speed=default/max):
diffusers requires the adapter loaded before compilation, so applying one to
the already-compiled module fails with adapter-key mismatches. The status
gate now hides the picker and generate raises a clear message instead.
Convert a cancelled Hub LoRA download (RuntimeError Cancelled) to the
diffusion cancellation sentinel in resolve_specs, so an unload/superseding
load during resolution maps to a 409 instead of a generic server error.
Drop weight-0 LoRA rows before the native support gate so a request carrying
only disabled adapters stays a no-op on families where native LoRA is
unsupported, matching the diffusers path.
Reject duplicate LoRA ids in the request model: both apply paths suffix
colliding names, so a repeated id would stack the same adapter past its
per-adapter weight bound.
Strip all user-typed <lora:...> prompt tags on the native path (only the
selected adapters are materialized in the managed lora-model-dir, so an
unselected tag can never resolve), and restore saved LoRA selections from a
gallery recipe so restore reproduces a LoRA image.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Harden ControlNet resolve, gallery metadata, and the control-type picker
Check cancellation immediately after a ControlNet from_pretrained and before
any device placement, so an unload/eviction that raced the download does not
allocate several GB onto the GPU after the load was already cleared.
Require a loadable weight or shard index (not just config.json) before a local
ControlNet folder is advertised, so an interrupted copy is hidden instead of
failing deep in from_pretrained as a generic 500.
Do not record a strength-0 ControlNet in the gallery recipe: it is treated as
disabled and skipped, so the image is unconditioned and the metadata must not
claim a ControlNet was applied.
Build the control-type picker from the selected ControlNet's advertised
control_types instead of a hardcoded passthrough/canny pair, so a union model
with a precomputed depth or pose map sends the correct control_mode.
* Address further Codex findings on the image-workflows PR
- Persist the actual output image size in the gallery recipe instead of the
request sliders: Transform/Inpaint/Edit derive the size from the uploaded
image, Extend grows the canvas, and Upscale resizes it, so the sliders
recorded (and later restored) the wrong dimensions for those workflows.
- Reject a remote '*-GGUF' repo loaded as a full pipeline (no single-file
name) in validate_load_request, so the unloadable pick fails before chat is
evicted rather than deep in from_pretrained.
- Only publish an image-conditioned from_pipe wrapper to the shared aux cache
when the load is still current: from_pipe runs under the generate lock but
not the state lock, so an unload racing its construction could otherwise
cache a wrapper over torn-down modules that a later load would reuse.
- Verify the Windows CUDA runtime archive checksum before extracting it, like
the main sd-cli archive, so a corrupt or tampered runtime is rejected rather
than extracted next to the binary.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Refuse non-SDXL base models at diffusion training start
The trainer only supports the SDXL U-Net, but a FLUX / Qwen-Image / Z-Image
repo or a GGUF filename passed as base_model was accepted and then failed
minutes later inside StableDiffusionXLPipeline.from_pretrained with an
unrelated-looking error. Add a name-based guard in normalized() so known
DiT-family names and .gguf checkpoints are rejected up front, which the API
start route surfaces as an immediate 400 with a message that says exactly
which bases are trainable. Unrecognisable names still pass through so custom
local SDXL checkpoints keep working.
* Add diffusion dataset upload and training info endpoints
Training an image LoRA required knowing the Studio home layout and copying
files onto the server by hand, which is the most confusing step of the whole
flow. Two small endpoints fix that:
- GET /api/train/diffusion/info reports the datasets and outputs roots plus
every dataset folder that contains images (with image/caption counts), so
the UI can offer a picker instead of a blind free-text path.
- POST /api/train/diffusion/dataset uploads images and optional caption
.txt / metadata.jsonl files into a named folder under the datasets root,
creating it on first use and accumulating on repeat uploads so large sets
can arrive in batches. Names are validated to a single path component and
files stream to disk under the same per-upload size cap as LLM dataset
uploads. The returned name is a valid data_dir for /diffusion/start.
* Rework the Train LoRA dialog into a guided SDXL flow
The dialog assumed users knew the Studio home layout and that only SDXL is
trainable, and hid both facts behind free-text fields. Restructure it around
the three real decisions:
- Base model is a dropdown of the trainable SDXL picks (Base 1.0, Turbo, the
loaded SDXL pipeline when there is one) with a custom repo/path escape
hatch, instead of a bare text field defaulting to a repo id.
- Training images come from an in-browser upload (new dataset endpoints) or
a picker over existing dataset folders with image/caption counts. No shell
access or knowledge of the datasets root is needed any more, and the
captioning rules are explained inline.
- The output field is now Adapter name and the instance prompt is labelled
as the trigger prompt, with a no-captions warning wired to the selected
dataset's actual caption count.
Hyperparameters collapse behind a training settings toggle since the
defaults suit a first run. A completed run says where the adapter went and
offers Done / Train another, and the top-bar button gets an icon and a
plainer description. The dialog title states the SDXL-only scope and that
other families load LoRAs but cannot train them yet.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Validate diffusion training config before freeing the GPU
The start route freed resident GPU workloads (export, Images pipeline, chat)
before the service validated the config, so a start that was then refused,
now including a non-SDXL base model, tore down the user's loaded model for
nothing. Run the same cheap normalise pass first; the LLM path already
follows this rule via its before_spawn hook.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Refactor diffusion LoRA training into a family-aware platform
Split the SDXL trainer into a shared, architecture-agnostic layer so more model
families can be trained without duplicating the plumbing:
- New core/training/diffusion_train_common.py holds the config + validation, dataset
discovery, event emission, stop protocol, adapter publishing, and a lazy trainer
registry (get_trainer). diffusion_lora_trainer.py keeps the SDXL-specific loop and
re-exports the moved names so existing imports are unchanged.
- The SDXL-only base-model blocklist becomes a positive check: the family is resolved
from the base model (or an explicit model_family) via the diffusion family registry,
and a known-but-not-yet-trainable family is refused with a clear message. Unknown
custom names still default to the SDXL trainer.
- DiffusionFamily gains a trainable flag and train_base_repos; SDXL is marked trainable.
DiT families flip on when their trainers land.
- Trained adapters now write a <name>.json metadata sidecar (family, base model, rank,
trigger prompt, ...) that the LoRA scanner reads to family-gate the adapter in the
picker instead of showing it as unknown for every model.
- The training base-model trust allowlist adds the official FLUX.1-dev, Z-Image-Turbo,
and Qwen-Image repos (safetensors-only, no remote code).
* Retain diffusion training loss history and expose it in status
The training service kept only the latest loss, so a live loss chart could show a
single point. Fold each progress event into bounded (step, loss, lr) history arrays
(capped at 4000 points, decimated when full) plus the latest throughput and peak VRAM,
and record the family / base model / catalog path on completion. The status endpoint
returns these as a nested metric_history object the UI can chart directly, and the
start request accepts an optional model_family override.
* Tests for the diffusion training platform
Cover the trainer registry (get_trainer resolves SDXL, unknown family raises),
family resolution (explicit model_family validation, resolved_family on the config),
the metadata sidecar write + scan read with family gating, and the service loss-history
folding (append, bad-point skipping, decimation at cap, family/perf fields) plus the
status route nesting metric_history.
* Add diffusion dataset labeling and example-import endpoints
The Train tab needs to let users caption small datasets in the browser and
pull in a ready-made set to see training work end to end, neither of which
the upload-only endpoint supported.
Add, under /api/train/diffusion/dataset:
- GET {name}/images lists every image with its resolved caption (metadata
beats a per-image sidecar, matching the trainer's discovery order) so
uncaptioned images are visible and flaggable.
- GET {name}/image/{filename} serves an image, with ?thumb=<px> returning a
cached downscaled JPEG kept in a hidden .thumbs subdir (regenerated when
the source is newer) so the labeling grid stays light.
- PUT {name}/caption/{filename} writes, or when blank clears, the .txt
sidecar; DELETE {name}/image/{filename} removes the image plus its
sidecars and thumbnails.
- GET dataset-examples lists a curated, license-labelled registry, and
POST dataset/import-example materializes one into a dataset folder as
numbered images + .txt captions. Two loaders cover the shapes seen in the
wild: streaming rows from datasets.load_dataset (dog-example, Tuxemon) and
a snapshot + jsonl walk for imagefolder repos whose captions live in a
non-standard *.jsonl (the public-domain tarot set). Imports are idempotent
and cap the image count.
Filenames and dataset names are validated against path traversal and pinned
inside the datasets root.
* Test diffusion dataset labeling and example-import endpoints
Cover caption precedence, thumbnail generation and .thumbs exclusion,
caption write/clear, image delete cleanup, path-traversal rejection on
names and filenames, and example import with a mocked datasets.load_dataset
(files plus sidecars written, idempotent second call, cap respected, load
failure mapped to 502).
* Add flow-matching DiT LoRA trainers (FLUX.1-dev, Qwen-Image, Z-Image)
Extends diffusion LoRA training beyond SDXL to the three popular DiT families
via a single shared flow-matching loop parameterised by small per-family specs
(loading, prompt/latent encoding, transformer forward, save). Verified against
diffusers 0.38.0:
- FLUX.1-dev: 2x2 latent packing + image ids, guidance-embed forward, on-the-fly
nf4 QLoRA of the 12B transformer (the dev repo is gated, so training needs the
user's HF token).
- Qwen-Image: 5D VAE latents normalised by the per-channel latents_mean/std,
img_shapes forward, prequant nf4 base by default (on-the-fly nf4 for the bf16
base).
- Z-Image: list I/O with the reversed timestep convention and a negated
prediction, bf16 only.
The registry (get_trainer) and DiffusionFamily.trainable / train_base_repos now
route these families to the DiT trainer; the SDXL blocklist guard is replaced by
a positive family resolution that also rejects GGUF repos (inference-only) and
still-unsupported families. Per-family defaults + labels + VRAM notes are exposed
via family_train_infos for the Train UI.
Memory: caption embeddings are precomputed once and the text encoders freed
before the loop; gradient checkpointing (non-reentrant, required for bnb 4-bit)
and 8-bit AdamW are on by default.
* Speed up + shrink SDXL LoRA training (precompute text embeds, 8-bit AdamW)
SDXL re-encoded every caption with both CLIP text encoders on every step (pure
waste, since captions are constant) and kept the encoders resident. Precompute
each unique caption's embeddings once, then free the text encoders before the
loop: numerically identical (embeddings are deterministic and this consumes no
torch RNG, so the noise/timestep stream is unchanged) but faster and ~1.5 GB
lighter. Default the optimizer to 8-bit AdamW (bitsandbytes) with an fp32
fallback, halving optimizer state with no meaningful LoRA quality cost. Env
toggles (UNSLOTH_DIFFUSION_NO_PRECOMPUTE / _FP32_OPTIM) let the accuracy guard
A/B the paths.
* Expose trainable families in /diffusion/info and preflight gated bases
The training info endpoint now returns the trainable model families (name,
label, default + allowed base repos, recommended defaults, and a VRAM/access
note) so the Train UI can offer a base picker with realistic guidance. The start
route preflights a gated base repo (HEAD model_index.json with the user's token)
BEFORE freeing resident GPU workloads, so a missing FLUX.1-dev license/token
fails fast with an actionable 400 instead of evicting the loaded model and then
hitting a confusing mid-load 401.
* Tests for DiT trainers, family resolution, info families, gated preflight
Cover the DiT spec table, the QLoRA prequant heuristic, the Z-Image bf16-only
guard, the gated-repo name check, family resolution now that FLUX/Qwen/Z-Image
are trainable (and GGUF repos are rejected as inference-only), the families list
in /diffusion/info, and the gated-base 400 preflight that leaves the GPU
untouched.
* Wrap the DiT training forward in bf16 autocast
The fp32 LoRA parameters and the bnb 4-bit base matmuls need a single
compute dtype during the forward, exactly like the diffusers dreambooth
scripts run under accelerator.autocast. Without it the 4-bit backward on
FLUX.1-dev fails with an illegal-address CUBLAS error partway into the
first step. Z-Image and Qwen-Image smokes are unaffected and the SDXL
path (its own trainer) is untouched.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add Smithsonian Butterflies and Nouns example datasets
Two permissive ~100-image sets for the Train tab: huggan/smithsonian_butterflies_subset
(CC0, the classic diffusers-docs training set, imported as a subject set with a trigger
prompt since its metadata columns are species names not captions) and m1guelpf/nouns
(CC0, captioned pixel-art avatars via the text column). Both cap at 100 images.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override
Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want
fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and
prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match
on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40;
GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it.
Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s),
no overflow, quality below the quant noise floor. So the default leans to accuracy on
data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force)
lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off).
187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override
threading).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9
scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer.
Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with
use_triton_kernel=False (the default triton path needs the missing MSLK library), but
only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than
using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's
CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto
on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder
Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4
tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS,
2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the
diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both
slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166
vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is
correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add
scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 9): pre-quantized transformer loading
The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on
the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs
13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline
(scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton
on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True)
the quantized weights, so the dense bf16 never touches the GPU.
Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching
GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same
torchao config + min_features filter the runtime path uses, applied ahead of time.
New core/inference/diffusion_prequant.py (resolve_prequant_source +
load_prequantized_transformer, best-effort, lazy imports). diffusion.py
_load_dense_quant_pipeline tries the pre-quant source first and falls back to the
dense materialise+quantise path, then to GGUF, so the default is unchanged.
DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty
prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for
the resolver, the meta-init+assign loader, and the backend branch selection +
fallbacks; GPU verification via scripts/verify_prequant_backend.py.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 10): attention-backend selection
Add a selectable attention kernel via the diffusers set_attention_backend
dispatcher. Attention is memory-bandwidth bound, so a better kernel is an
end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV
matmuls torchao never touches) and composes with torch.compile.
auto picks the best exact backend for the device: cuDNN fused attention
(_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x
end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default
(below the compile/quant noise floor); native SDPA elsewhere and when speed=off
(so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/
xformers/aiter are honored, and an unavailable kernel falls back to the default
rather than failing the load.
New core/inference/diffusion_attention.py (normalize + per-device select + apply,
best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline;
attention_backend threads through begin_load / load_pipeline / status like the other
load knobs. New request field attention_backend + status field. Hermetic CPU tests
for normalize / select policy / apply fallback, plus route threading + 422. Measured
via scripts/perf_levers_probe.py.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder
Consumer / workstation GPUs halve fp8 (and fp16/bf16) FP32-accumulate tensor-core
throughput, while int8 runs at full rate (int32 accumulate is not nerfed). Public
benchmarks (SDNQ across RTX 3090/4090/5090, AMD, Intel) confirm int8 via torch._int_mm
is as fast or faster than fp8 on every consumer part, and the only path on pre-Ada
consumer cards without fp8 tensor cores. So when transformer_quant=auto, reorder the
arch tier to put int8 first on a consumer/workstation GPU (detected by the existing
_is_consumer_gpu name heuristic), while data-center HBM parts keep fp8 first.
Pure ladder reorder via _prefer_consumer_scheme; no new flags. Verified non-regression
on a B200 (still picks fp8). Hermetic tests for consumer Blackwell/Ada/workstation
(-> int8) and data-center Ada/Hopper/Blackwell (-> fp8).
* Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT
Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across
denoise steps a DiT's output settles, so once the first block's residual barely
changes the remaining blocks are skipped and their cached output reused. diffusers
ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the
standalone apply_first_block_cache hook as a fallback).
Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 ->
2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar.
OFF by default and a per-load opt-in: the win scales with step count, so it is for
many-step models (Flux / Qwen-Image) and pointless for few-step distilled models
(e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction
of the trajectory. It composes with regional compile only with fullgraph=False (the
cache's per-step decision is a torch.compiler.disable graph break), which the speed
layer now switches to automatically when a cache is engaged. Best-effort: a model
whose block signature the hook does not recognise is caught and the load proceeds
uncached.
- new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache
(enable_cache / apply_first_block_cache fallback; threshold auto-raised for a
quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import).
- diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph
when a cache is engaged.
- diffusion.py: apply_step_cache before compile; thread transformer_cache /
transformer_cache_threshold through begin_load -> load_pipeline and report the
engaged mode in status().
- models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and
transformer_cache_threshold request fields, engaged mode in the status response.
- hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold
selection, and best-effort failure handling, plus route threading + validation.
- scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM /
LPIPS vs the compiled no-cache baseline).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 14): fix int8 dense quant on Flux / Qwen (skip M=1 modulation linears)
The opt-in dense int8 transformer path crashed on Flux.1 and Qwen-Image with
'torch._int_mm: self.size(0) needs to be greater than 16, but got 1'. int8 dynamic quant
goes through torch._int_mm, which requires the activation row count M > 16. A DiT's AdaLN
modulation projections (Flux norm1.linear 3072->18432, Qwen img_mod.1 / txt_mod.1, Flux.2
*_modulation.linear) and its timestep / guidance / pooled-text conditioning embedders are
computed once from the [batch, dim] conditioning vector (M = batch = 1), not per token, so
they hit _int_mm at M=1 and crash. Their feature dims are large, so the existing
min_features filter did not exclude them.
Fix: the int8 filter now also skips any Linear whose fully-qualified name matches a
modulation / conditioning-embedder token (norm, _mod, modulation, timestep_embed,
guidance_embed, time_text_embed, pooled). These layers run at M=1 once per block and are a
negligible share of the FLOPs, so int8 keeps the full speedup on the attention / FFN layers
(M = sequence length). fp8 / nvfp4 / mxfp8 use scaled_mm, which has no M>16 limit and
quantises these layers fine, so the exclusion is int8-only. Sequence embedders
(context_embedder / x_embedder / txt_in, M = seq) are deliberately not excluded -- note
'context_embedder' contains the substring 'text_embed', which is why the token is the
specific 'time_text_embed', not 'text_embed'.
Measured on a B200 (1024px, transformer_quant=int8 + speed=default), int8 now runs on every
supported model and is the fastest dense path on Flux/Qwen (int8 runs full-rate vs fp8's
FP32-accumulate): FLUX.1-dev 9.62s eager -> 1.98s (4.86x, vs fp8 2.15s), Qwen-Image -> 1.87s
(5.57x, vs fp8 2.09s), FLUX.1-schnell -> 0.41s (3.59x). Z-Image and Flux.2-klein (already
working) are unchanged.
- diffusion_transformer_quant.py: add _INT8_EXCLUDE_NAME_TOKENS; make_filter_fn takes
exclude_name_tokens; quantize_transformer passes it for int8 only.
- hermetic test that the int8 filter excludes the modulation / embedder linears (and keeps
attention / FFN / sequence-embedder linears), while fp8 keeps them.
- scripts/int8_linear_probe.py: the meta-device probe used to enumerate each transformer's
Linear layers and derive the exclusion list.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 15): build int8 pre-quantized checkpoints (skip M=1 modulation linears)
The prequant-checkpoint builder applied the dense quant filter without the int8-only
M=1 modulation / conditioning-embedder exclusion the runtime path uses, so a built int8
checkpoint baked those projections as int8 and crashed (torch._int_mm needs M>16) at the
first denoise step on Flux / Qwen. Factor the scheme->exclusion decision into a shared
exclude_tokens_for_scheme() used by both the runtime quantise path and the offline builder
so they can never drift, and apply it in build_prequant_checkpoint.py. int8 prequant now
produces a working checkpoint on every supported model, giving int8 (the consumer-preferred
scheme) the same ~2x load-VRAM and download reduction fp8 already had.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 16): route no-GPU loads to the native sd.cpp engine
When no CUDA/ROCm/XPU GPU is available, route diffusion load/generate to the
native stable-diffusion.cpp engine instead of diffusers, with diffusers as the
guaranteed fallback. On CPU sd.cpp is 1.4-2.8x faster and uses 1.5-2.2x less RAM.
- diffusion_engine_router: centralised engine selection (built on the existing
select_diffusion_engine), env opt-outs, MPS gating, recorded fallback reason.
- sd_cpp_backend (SdCppDiffusionBackend): the diffusers backend method surface
backed by sd-cli, with lazy binary install, registry-driven asset fetch,
step-progress parsing, and cancellation.
- diffusion_families: per-family single-file VAE + text-encoder asset mapping.
- sd_cpp_engine: cancellation support (process-group kill + SdCppCancelled).
- routes/inference + gpu_arbiter: drive the active engine via the router; the
API now reports the active engine and any fallback reason.
- tests for the backend, router, route selection, and cancellation.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Phase 16 review fixes: engine-switch unload, sd.cpp error mapping, per-image seeds, Qwen sampler
Address review feedback on #6724:
- engine router: unload the engine being deactivated on a switch, so the old
model is not left resident-but-unreachable (the evictor only targets the active
engine).
- generate route: sd.cpp execution errors (nonzero exit / timeout / missing
output) now map to 500, not 409 (which only means not-loaded / cancelled).
- native batch: return per-image seeds and persist the actual seed for each image
so every batch image is reproducible.
- Qwen-Image native path: apply --sampling-method euler --flow-shift 3 per the
stable-diffusion.cpp docs; other families keep sd-cli defaults.
- honor speed_mode (native --diffusion-fa) and, off-CPU, memory_mode/cpu_offload
offload flags on the native load instead of hardcoding them off.
- fail the load when the sd-cli binary is present but not runnable (version()
now returns None on exec error / nonzero exit).
- size estimate: only treat the transformer asset as a possible local path.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in
load_prequantized_transformer ends in torch.load(weights_only=False), which executes
arbitrary code from the pickle. The transformer_prequant_path load-request field reached
that unpickle for any local file an authenticated caller named, so a request could trigger
remote code execution. Refuse the source.kind=='path' branch unless the operator sets
UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted
and unaffected. Document the requirement on the API field and add gate tests.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa
- apply_attention_backend now restores the native default when no backend is requested or a
kernel fails. diffusers keeps a process-wide active attention backend that
set_attention_backend updates, and a fresh transformer's processors follow it, so a load
that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile
load pinned, breaking the bit-identical/off guarantee.
- select_attention_backend drops flash3/flash4 up front when the CUDA capability is below
Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit
request on the wrong card set fine then crashed mid-generation; it now falls back to native.
- Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a
valid alias of native) is accepted instead of 422-rejected by Pydantic.
- Drop the dead replace('-','_') normalization (no alias uses dashes/underscores).
- perf_levers_probe.py output dir is now relative to the script, not a hardcoded path.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF
- apply_step_cache now engages only via the transformer's native enable_cache (the diffusers
CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a
cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin
transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported
transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model
now runs uncached per the best-effort contract.
- GGUF transformers are quantized (the default Studio load path), so they now use the higher
quantized FBCache threshold when the caller leaves it unset, instead of the dense default
that could keep the cache from triggering.
- fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so
fullgraph=True failed warmup and silently measured an eager cached run); output dir is now
relative to the script, not a hardcoded path.
* Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder
_is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as
consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead
of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already
classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the
same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate.
Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0).
* Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config
Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just
TypeError on older torchao), which would break the mxfp8 config helper instead of falling
back to the default. Catch both so the fallback is robust.
quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds
CUDA memory during the per-row VRAM probe; output dir relative to the script.
* Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load
- snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform
missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the
whole snapshot. restore_backend_flags restores each flag independently so one failure can't
leave the others leaked process-wide.
- load_pipeline restores the flags (and clears the GPU cache) when the build fails after
apply_speed_optims mutated the process-wide flags but before _state captured them for unload
to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and
contaminated later off generations.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output
Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs
without producing output (or without closing stdout) would never reach proc.wait and the
wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the
PROCESS, so the main thread always enforces the timeout and kills a hung process (which
closes the pipe and ends the reader). Add a test that times out even when stdout blocks,
and make the no-binary test hermetic so a host-installed sd-cli can't leak in.
* Studio diffusion (Phase 14): guard the int8 exclusion filter against a None fqn
The filter callback can be invoked without a module name, so fqn.lower() would raise
AttributeError on None. Fall back to an empty name (nothing matches the exclusion tokens,
so the linear is kept) instead of crashing the quantise pass.
* Studio diffusion (Phase 16) review fixes: native engine robustness
- sd_cpp_backend: stop truncating explicit seeds to 53 bits (mask to int64);
a large requested seed was silently collapsed (2**53 -> 0) and distinct seeds
aliased to the same image. Random seeds stay 53-bit (JS-safe).
- sd_cpp_backend: sanitize empty/whitespace hf_token to None so HfApi/hf_hub
fall back to anonymous instead of failing auth on a blank token.
- sd_cpp_backend: a superseding load now cancels the in-flight generation, so the
old sd-cli can no longer return/persist an image from the previous model.
- diffusion_engine_router: run the previous engine's unload() OUTSIDE the lock so a
slow 10+ GB free / CUDA sync does not block engine selection.
- diffusion_engine_router: probe sd-cli runnability (version()) before committing to
native, so a present-but-unrunnable binary falls back to diffusers at selection.
- diffusion_device: resolve a torch-free CPU target when torch is unavailable, so a
CPU-only install can still reach the native sd.cpp engine instead of failing load.
- tests updated for the runnability probe + a not-runnable fallback case.
* Studio diffusion (Phase 9) review fixes: prequant safety + validation
- SECURITY: a request-supplied local pre-quant path is now unpickled only when it
resolves inside an operator-configured ALLOWLIST of directories
(UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in,
once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on
any path a load request named (arbitrary code execution). realpath() blocks symlink
escapes; a bare on/off toggle is no longer a wildcard.
- Validate the checkpoint's min_features against the runtime Linear filter, so a
checkpoint that quantised a different layer set is rejected instead of silently
loading a model that mismatches the dense path while reporting the same scheme.
- Tolerant base_model_id compare (exact or same final path/repo segment), so a local
path or fork of the canonical base is accepted instead of falling back to dense.
- _has_meta_tensors uses any(chain(...)) (no intermediate lists).
- prequant verify/probe scripts use repo-relative paths (+ env overrides), not the
author's absolute /mnt paths.
- tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts
- diffusion_memory: when group offload is unavailable and the plan falls back to
whole-module offload, enable VAE tiling (the group plan left it off, but the fallback
is the low-VRAM path where the decode spike can OOM). Covers both the group and
sequential fallback branches.
- perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a
balanced bit-identity regression actually fails the check instead of exiting 0.
- compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and
the load-progress poll has a 30 min deadline instead of looping forever on a hang.
- test for the group->model fallback enabling VAE tiling.
* Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path
- diffusion: a torchao-quantized transformer is committed only compiled. A dense model
resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF
it replaced), so when transformer_quant engaged and speed resolved to off, promote to
default (regional compile); warn loudly if compile still does not engage.
- diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the
CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK);
otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF.
- nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path).
- test asserts the eager-quant -> default-compile promotion.
* Studio diffusion (Phase 10) review fixes: attention gating + probe isolation
- diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere
NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation,
so auto now stays on native SDPA there.
- diffusion_attention: _active_attention_backend handles get_active_backend() returning an
enum/None (not a tuple); the old unpack always raised and was swallowed, so
the native-restore short-circuit never fired.
- perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS
on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul
so the inductor_flags variant doesn't leak into later compiled rows.
- tests for the SM80 cuDNN gate.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening
- install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit
timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket);
extract through a per-member containment check (Zip-Slip guard); expanduser the
--install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch
the separately-published cudart runtime DLL archive so sd-cli.exe can start.
- sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the
installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start
sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend
crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie.
- tests: Zip-Slip rejection, normal extraction, studio-home discovery.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs
Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes
the numbered files <stem>_<idx><suffix> (base_0.png, base_1.png, ...) instead of
the literal --output path. SdCppEngine.generate checked only the literal path, so
a batch generation would exit 0 and then raise 'no image' (or return a stale
file). generate now returns the literal path when present and otherwise falls
back to the numbered siblings; single-image behavior is unchanged.
Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected
without error.
* Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats
Codex review on the native engine arg builder:
- build_sd_cpp_command emitted --width/--height unconditionally, so an
img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of
the input. width/height are now Optional (None = unset): an image-conditioned
run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives
the size from the input image (set_width_and_height_if_unset); a plain txt2img
run with unset dims keeps the prior 1024x1024 default; explicit dims are always
honored. width/height are read only by the builder, so the type change is local.
- build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...)
that silently swallowed repeats=0 into sd-cli's default of one pass, turning an
explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError
and emits the flag for any explicit value != 1.
Tests: img2img unset dims omit width/height (init_img and ref_images), explicit
dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at
the default. (Two pre-existing binary-discovery tests fail only because a real
sd-cli is installed in this dev environment; unrelated to this change.)
* Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc
Codex review: the transformer_prequant_path field description still told operators
to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the
prior security fix made that variable a directory allowlist -- _allowed_prequant_roots
deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator
following the documented =1 would have every transformer_prequant_path request
silently refused. The description now states it must name one or more allowlisted
directories and that a bare on/off value is not accepted.
Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does
not say =1, and describes an allowlist/directory (guards against doc drift).
* Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset
Codex review on attention-backend selection:
- Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on
pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first
generation with no fallback. select_attention_backend now applies
_cudnn_attention_supported() to an explicit cuDNN request too.
- flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a
Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3
is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a
(min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+.
- apply_attention_backend's success path left diffusers' process-wide active
backend pinned to the kernel it set; a later component whose processors are
unconfigured (backend None) would inherit it. It now resets the global registry
to native after a successful per-transformer set (the transformer keeps its own
backend), best-effort. Also fixed _active_attention_backend: get_active_backend()
returns a (name, fn) tuple, so the prior code stringified the tuple and never
matched a name, defeating the native-restore short-circuit.
Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on
SM90; global registry reset after a successful set; _active_attention_backend
reads the tuple return.
* Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder
Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300
(Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So
_is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the
auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting
llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they
are treated as data-center class and keep the intended fp8-first behavior.
Test: extends the datacenter parametrize with 'NVIDIA B300' and
'NVIDIA GH200 480GB' (now _is_consumer_gpu False).
* Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder
Codex review: the M=1 modulation/embedder exclusion was wired only into the dense
runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called
make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint
quantised the AdaLN modulation and conditioning-embedder linears, and loading it
via transformer_prequant_path (the load path only loads already-quantised tensors,
it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes
for the runtime path.
Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 ->
the M=1 exclusion, every other scheme -> none) and use it in both the runtime
quantiser and the builder, so a prequant artifact's quantised-layer set always
matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion).
Test: int8_exclude_name_tokens returns the exclusion for int8 and () for
fp8/nvfp4/mxfp8.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 16) review round 2: native CPU arbiter, status offload, load race
Codex review on the native-engine routing:
- The /images/load route took the GPU arbiter (acquire_for(DIFFUSION) -> evict chat)
unconditionally after engine selection. A native sd.cpp load on a pure-CPU host
never touches the GPU, so that needlessly tore down the resident chat model. The
handoff is now gated: diffusers always takes it, a force-native sd.cpp load on a
CUDA/XPU/MPS box still takes it, but a native sd.cpp load on a CPU host skips it.
- sd_cpp status() hardcoded offload_policy 'none' / cpu_offload False even when
_run_load computed real offload flags (balanced/low_vram/cpu_offload off-CPU), so
the setting was unverifiable. status now derives them from state.offload_flags
(still 'none' on CPU, where the flags are empty).
- _run_load committed the new state without cancelling/waiting on a generation that
started during the (slow) asset download, so a stale sd-cli run against the OLD
model could finish afterward and persist an image from the previous model once the
new load reported ready. The commit now signals the in-flight cancel and waits on
_generate_lock before swapping _state (taken only at commit, so the download never
serialises against generation), mirroring the diffusers load path.
Tests: CPU native load skips the arbiter while a GPU native load takes it; status
reports offload active when flags are set; _run_load cancels and waits for an
in-flight generation before committing.
* Studio diffusion (Phase 14) review round 2: align helper name with the stack
Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the
identical helper already present higher in the diffusion stack (Phase 16). The
helper definition, the runtime quantiser call, and the offline builder are now
byte-identical to that version, so the two branches no longer introduce a
divergent name for the same single-source-of-truth and the stack merges without
a conflict on this fix. No behavior change.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion: eager patches + torch.compile cache speed phase
Adds the opt-in speed path for the GGUF diffusion transformer behind a
selectable speed mode (default off, so output is unchanged until a profile
is chosen):
- diffusion_eager_patches.py: shared eager fast-paths (channels_last,
attention/backend selection, fused norms and QKV) installed at load and
rolled back on unload or failed load.
- diffusion_compile_cache.py / diffusion_gguf_compile.py: a persistent
torch.compile cache and the GGUF-transformer compile wiring.
- diffusion_arch_patches.py: architecture-specific patches.
- diffusion_patch_backend.py: shared install/restore plumbing.
- diffusion_speed.py: speed-profile planning.
Tests for each module plus the benchmarking and probe scripts used to
measure speed, memory, and accuracy of the path.
* Studio diffusion: image workflows (safetensors, image-conditioned, editing) + Images UI
Backend:
- Load non-GGUF safetensors models: full bnb-4bit pipelines and single-file
fp8 transformers, gated to the unsloth org plus a curated allowlist.
- Image-conditioned workflows built with Pipeline.from_pipe so they reuse the
loaded transformer/VAE/text-encoder with no extra VRAM: img2img, inpaint,
outpaint, and a hires-fix upscale pass.
- Instruction editing as its own family kind (Qwen-Image-Edit-2511,
FLUX.1-Kontext-dev) and FLUX.2-klein reference conditioning (single and
multi-reference) plus klein inpaint.
- Auto-resize odd-sized inputs to a multiple of 16 (and resize the matched
mask) so img2img/inpaint/edit no longer reject non-/16 uploads. Bound the
decoded image size and cap upscale output to avoid OOM on large inputs.
- Fixes: from_pipe defaulting to a float32 recast that crashed torchao
quantized transformers; image-conditioned calls forcing the slider size
onto the input image. Native sd.cpp engine rejects image-conditioned and
reference requests it cannot serve.
Frontend:
- Redesigned Images page with capability-gated workflow tabs (Create,
Transform, Inpaint, Extend, Upscale, Reference, Edit), a brush mask editor,
client-side outpaint, and a multi-reference picker.
- Advanced options moved to a right-docked panel mirroring Chat: closed by
default, toggled by a single fixed top-bar button that stays in place.
sd.cpp installer: pin the release, verify each download's sha256, add a
download timeout, and make the source repo configurable for a future mirror.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio Images: correct the Advanced panel comment (closed by default, fixed toggle)
* Studio diffusion: LoRA adapters for the Images workflow
Add community LoRA support across both diffusion backends, the single
biggest step toward broad image-workflow coverage.
Backend
- New shared module core/inference/diffusion_lora.py: adapter discovery
(local scan + curated catalog + owner/name[:file] Hub refs), download
via hf_hub_download_with_xet_fallback, alias sanitization, native
managed-dir materialization with collision-broken aliases, prompt-tag
injection (deduped against user-typed tags), and a supports_lora gate.
- Native sd-cli: resolve + materialize selected LoRAs into a per-run
managed dir, inject <lora:ALIAS:w> tags, pass --lora-model-dir with
--lora-apply-mode auto. The arg builder already emitted these flags.
- Diffusers: non-fused load_lora_weights + set_adapters manager, tracked
on the pipe so an unchanged selection is a no-op and a model swap
resets; cleared on unload. Never fuses (breaks quantized transformers
and blocks live weight tweaks).
- Gated off where unsupported: torchao fp8/int8 dense, GGUF-via-diffusers,
and native Qwen-Image (no LoRA name-conversion branch upstream).
- Request contract: optional loras on DiffusionGenerateRequest; empty or
omitted is identical to today. supports_lora surfaced in status; chosen
LoRAs persisted in gallery recipe metadata.
- New GET /api/models/diffusion-loras for the picker (family-filtered).
Frontend
- Repeatable multi-LoRA picker (adapter select + weight slider 0..2 +
remove), gated by the loaded model's supports_lora and family, max 8.
Tests
- New test_diffusion_lora.py (14): helpers, request validation, native
tag/dir wiring, diffusers set_adapters manager, supports_lora matrix.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion: ControlNet for the Images workflow (diffusers)
Add ControlNet conditioning, the #2 most-used diffusion workflow after
LoRA, on the diffusers backend for the families with ControlNet pipelines
(FLUX.1 and Qwen-Image), with Union models as the default picks.
Backend
- New core/inference/diffusion_controlnet.py: family-gated discovery
(curated Union models + local dirs + bare owner/name repos), resolution
to a loadable repo/dir, control-image preprocessing (passthrough +
a dependency-free canny edge map), and a supports_controlnet gate.
- diffusion.py: a ControlNet manager parallel to the LoRA one. Loads the
(small) ControlNet model once via from_pretrained (cached by id) and
builds the family's ControlNet pipeline via Pipeline.from_pipe(base,
controlnet=model), reusing the resident base modules at their loaded
dtype (no reload, no recast). Passes the control image + conditioning
scale + guidance start/end at generate time; cleared on unload.
- Families: FLUX.1 -> FluxControlNetPipeline/Model, Qwen-Image ->
QwenImageControlNetPipeline/Model. Others declare none (gated off).
- Gated off for the native engine, GGUF-via-diffusers, and torchao
fp8/int8 dense (same rule as LoRA). v1 conditions txt2img only.
- Request contract: optional controlnet on DiffusionGenerateRequest;
supports_controlnet in status; the choice persisted in gallery meta.
- New GET /api/models/diffusion-controlnets for the picker.
Frontend
- A ControlNet control in the Images rail (model select + control-image
upload + control-type select + strength slider), gated by the loaded
model's supports_controlnet + family, shown for text-to-image.
Tests
- New test_diffusion_controlnet.py (10): discovery/resolve/preprocess/gate
helpers, request validation, family wiring, and the diffusers pipe
manager (loads once, caches, from_pipe with controlnet, rejects
unsupported families).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio ControlNet: show the picker on the Create tab (workflow id is 'create', not 'txt2img')
The ControlNet control gated on workflow === 'txt2img', but the Images workflow tab ids are create/transform/inpaint/extend/upscale/reference/edit -- there is no 'txt2img'. So the picker never rendered even with a ControlNet-capable model loaded. Gate on 'create' (the text-to-image tab) for both the picker and the request wiring. Found via a live Playwright capture of the running Studio.
* Studio: do not force diffusers pipelines cross-tagged gguf into the GGUF variant expander
Some diffusers image repos (e.g. unsloth/Qwen-Image-2512-unsloth-bnb-4bit) carry a
stray "gguf" tag on the Hub but ship no .gguf files. The model search classified
them as GGUF from the bare tag, so the picker rendered the GGUF variant expander,
which then dead-ended at "No GGUF variants found." Trust the bare gguf tag only when
the repo is not a diffusers pipeline; the -GGUF name suffix and real gguf metadata
(populated via expand=gguf) remain authoritative, so genuine GGUF repos are unaffected.
* Studio Images: load non-curated unsloth/on-device diffusers repos instead of no-op
handleModelSelect only loaded curated safetensors ids and GGUF variant picks; any other
non-GGUF pick (an on-device diffusers folder, or a future unsloth diffusers image repo
surfaced by search) silently did nothing. Treat such a pick as a full diffusers pipeline
load when the id is unsloth-hosted or on-device (the backend infers the family + base repo
and gates loads to unsloth/* or local paths), and show a clear message otherwise instead
of silently ignoring the click. Curated and GGUF paths are unchanged.
* Studio Images: keep curated safetensors models in Recommended after download
The curated bnb-4bit / fp8 diffusion rows were filtered out of the Images picker's
Recommended list once cached (curatedSafetensorsRows dropped anything in downloadedSet),
so they vanished from the picker after the first load and could only be found by typing an
exact search. The row already renders a downloaded badge, matching how GGUF Recommended
rows stay visible when cached. Drop the exclusion so the curated safetensors always list.
* Studio diffusion LoRA: sanitize dots out of adapter aliases
The LoRA alias is used as the diffusers PEFT adapter name, and PEFT rejects names
containing "." (module name can't contain "."). sanitize_alias kept dots, so a LoRA whose
filename carries a version tag (e.g. Qwen-Image-2512-Lightning-8steps-V1.0-bf16) failed to
apply with a 400. Replace dots too; the alias stays a valid native <lora:NAME:w> filename
stem. Adds regression coverage for internal dots.
* Studio Images: clarify the GGUF transformer-quant Advanced control
Renamed the confusing "Transformer quant / GGUF default" control to "GGUF speed mode"
with an "Off (run the GGUF)" default, and reworded the hint to state plainly that FP8/INT8/
FP4 load the FULL base model (larger download + more VRAM) rather than re-packing the GGUF,
falling back to the GGUF if it can't fit. Behavior unchanged; labels/hint only.
* Studio Images: list on-device unsloth diffusion models in the picker
The Images picker's On Device tab hid every non-GGUF cached repo whenever a
task filter was active, so downloaded unsloth diffusion pipelines (bnb-4bit
and FP8 safetensors) never showed up there. List cached repos that pass the
task gate, limited under a filter to unsloth-hosted ones so base repos (which
fail the diffusion load trust gate) don't appear only to dead-end on click.
Chat behavior is unchanged: the task gate still drops image repos there.
* Studio: hide single-file image checkpoints from the chat model picker
The chat picker treats a cached repo as an image model, and hides it, only
when it ships a diffusers model_index.json. Single-file, ComfyUI, and
ControlNet image checkpoints (an FP8 Qwen-Image, a z-image safetensors, a
Qwen-Image ControlNet) carry none, so they surfaced as loadable chat models.
Fall back to resolving the repo id against the known diffusion families, the
same resolver the Images backend loads from, so these checkpoints are tagged
text-to-image and stay in the Images picker only.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio Images: add the FLUX.2-dev model family
Loading unsloth/FLUX.2-dev-GGUF failed because detect_family knew only the
Qwen3-based FLUX.2-klein, so FLUX.2-dev (the full, Mistral-based Flux2Pipeline)
resolved to nothing and the load errored. Add a flux.2-dev family: Flux2Pipeline
+ Flux2Transformer2DModel over the black-forest-labs/FLUX.2-dev base repo (gated,
reachable with an HF token), with its FLUX.2 32-channel VAE and Mistral text
encoder wired for the sd-cli path from the open Comfy-Org/flux2-dev mirror.
text-to-image only: diffusers 0.38 ships no Flux2 img2img / inpaint pipeline for
dev. Frontend gets sensible dev defaults (28 steps, guidance 4), distinct from
klein's turbo defaults. Verified live: GGUF load resolves the family + gated base
repo and generates a real 1024x1024 image on GPU.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio Images: clearer error for an unsupported diffusion model
When a repo id resolves to no diffusion family the load raised 'Could not infer a
diffusion family... Pass family_override (z-image)', which points at an unrelated
family and doesn't say what is supported. Replace it with a message that lists the
supported families (from a new supported_family_names helper) and notes that video
models and image models whose diffusers transformer has no single-file loader are
not supported. Applies to both the diffusers and native sd.cpp load paths. Also
refreshes two stale family-registry comments that still called FLUX.2-dev omitted.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add SDXL diffusion family (U-Net pipeline support)
SDXL is the first U-Net family in the diffusion backend: its denoiser is
pipe.unet (UNet2DConditionModel), not a DiT pipe.transformer, and a single-file
.safetensors is the whole pipeline rather than a transformer-only file. The
backend previously assumed a DiT transformer everywhere, so add the two hooks a
U-Net family needs and register SDXL.
DiffusionFamily gains denoiser_attr ("transformer" for DiT, "unet" for SDXL) and
single_file_is_pipeline (SDXL loads a single file via pipeline_class.from_single_file
with the base repo as config, instead of transformer_class.from_single_file plus a
companion assembly). _align_vae_dtype now reads the denoiser generically so img2img
and inpaint keep the VAE and U-Net dtypes aligned.
The non-GGUF trust gate is extended with a short, exact-match, safetensors-only
allowlist of official base repos (the SDXL base/refiner and sdxl-turbo), because
SDXL ships only as a full pipeline and has no unsloth-hosted GGUF. Local paths stay
trusted as before; a random repo, even one that detects as SDXL, is still rejected.
The image-conditioned and ControlNet workflows are the standard SDXL pipelines,
built around the resident modules via from_pipe like every other family, so SDXL
gets txt2img, img2img, inpaint, outpaint, upscale, LoRA and ControlNet. There is no
native sd.cpp mapping yet, so the no-GPU route falls back to diffusers.
Frontend catalog gains SDXL Base 1.0 and SDXL Turbo entries with SDXL step/guidance
defaults (Turbo: few steps, no CFG; base: ~30 steps, real CFG).
Tests: new test_diffusion_sdxl.py (family shape, detection, trust allowlist, model
kind, U-Net VAE-dtype alignment, LoRA gate) plus loader-branch tests in
test_diffusion_backend.py (pipeline-kind from_pretrained, single-file whole-pipeline
from_single_file, allowlist accept/reject). Verified live on GPU: sdxl-turbo loads
both as a pipeline and as a single file and generates coherent txt2img + img2img.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Images: LoRA free-text Hugging Face entry + recipe round-trip
The backend has always accepted a bare Hugging Face repo id (owner/name, or
owner/name:weight-file.safetensors) as a LoRA, downloading and applying it. But the
picker only rendered when the curated catalog had entries, and the catalog is empty,
so there was no UI path to apply any LoRA. Show the LoRA section whenever the loaded
model supports LoRA, and replace the curated-only dropdown with a text input: type a
Hub repo id, or pick a discovered adapter from a datalist of suggestions when the
catalog is populated.
Also restore LoRAs when loading a recipe. restoreSettings now parses the recipe's
"id:weight" strings (splitting on the last colon, since the id itself may contain one
for a specific weight file) back into the selection, so replaying a saved image
reproduces its adapters. The generate payload trims hand-typed ids and drops empty /
zero-weight rows, and a model swap clears the selection (a LoRA is family-specific)
without discarding a free-text pick that is not in the curated list.
* Add diffusion LoRA training (SDXL text-to-image)
First diffusion training path in Studio: train a LoRA on the SDXL U-Net from an
image + caption dataset and export it as a diffusers .safetensors that the existing
diffusion LoRA loader (and any diffusers pipeline) can load.
core/training/diffusion_lora_trainer.py:
- DiffusionLoraConfig with validation/defaults (rank, alpha, targets, lr, steps, grad
accumulation, resolution, min-SNR gamma, gradient checkpointing, lr scheduler, seed,
mixed precision).
- discover_image_caption_pairs: captions from metadata.jsonl / captions.jsonl, per-image
.txt/.caption sidecars, or a dreambooth instance_prompt fallback (pure, unit-tested).
- run_diffusion_lora_training: the loop -- freeze base, PEFT-wrap the U-Net attention
projections, VAE-encode (fp32 VAE to avoid the SDXL fp16 overflow), sample noise +
timesteps, predict, MSE loss with optional min-SNR weighting (epsilon / v-prediction),
AdamW + get_scheduler + grad accumulation + grad clipping, then export via
save_lora_weights. Emits worker-protocol events (model_load_*, progress, complete) and
polls should_stop for a clean stop with a partial save.
- run_diffusion_training_process: mp.Queue subprocess adapter (event_queue / stop_queue),
so the training worker can spawn it; plus a CLI entry point.
Only SDXL (U-Net) is trained here; DiT families and the Studio UI form + route wiring are
follow-ups. The trainer is decoupled and worker-ready.
Tests: test_diffusion_lora_trainer.py covers caption discovery (metadata / sidecar /
instance prompt / skip-uncaptioned / errors), config normalisation + validation, the SDXL
add-time-ids, and the dict->config adapter. Verified live on GPU: a 60-step SDXL LoRA run
lowers the loss, exports a ~45 MB adapter, and loading it back shifts generation from
baseline (mean abs pixel diff ~55/255).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* diffusion trainer: emit learning_rate in progress events (Studio pump compatibility)
The Studio training pump reads 'learning_rate' from progress events; the diffusion
trainer emitted 'lr'. Rename the field (and the CLI reader) so the trainer's events are
directly consumable by the existing training status/SSE machinery when it is wired into
the worker, without a translation shim.
* Wire diffusion LoRA training into the Studio API
Make the SDXL LoRA trainer reachable from the app with a small, self-contained job
service and JSON routes, deliberately separate from the LLM TrainingBackend (whose
lifecycle -- LLM config build, per-run SQLite rows, matplotlib plots, transfer-to-chat-
inference -- is text-training specific and would mis-handle a diffusion run).
core/training/diffusion_training_service.py: DiffusionTrainingService runs one job at a
time -- validate the config cheaply (before any spawn), spawn the trainer subprocess
(spawn context, parent-lifetime bound), pump its events (model_load_* / progress /
complete / error) into an in-memory status snapshot, and support a clean stop. The
subprocess context and target are injectable so the full start -> pump -> status ->
complete path is unit-tested without real multiprocessing or torch.
routes/training.py: POST /api/train/diffusion/start (400 on a bad config, 409 when a job
is already running), POST /api/train/diffusion/stop, GET /api/train/diffusion/status
(JSON poll). models/training.py: DiffusionTrainingStartRequest + response schemas
mirroring DiffusionLoraConfig, so model_dump() passes straight through.
Tests: test_diffusion_training.py -- service happy path, bad-config-before-spawn,
concurrent-job rejection, clean stop, crash-without-terminal-event, event transitions;
plus route wiring via the FastAPI TestClient (start / 422 / 400 / 409 / status / stop)
with a mocked service. The diffusion trainer's progress events already use the field
names this path expects.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Images: add a Train LoRA (SDXL) dialog
Surface the diffusion training API in the Images page. A "Train LoRA" button in the top
bar opens a self-contained dialog to fine-tune an SDXL LoRA on a folder of images: pick
the base model, dataset folder, output folder, an optional instance prompt, and the core
hyperparameters (steps, rank, resolution, batch, learning rate), then Start. The dialog
polls the training status while open and shows a progress bar, step count, live loss, and
the saved adapter path, with a Stop button for a clean stop.
The dialog is independent of the loaded generation model (training runs in its own
subprocess), and prefills the base model with the loaded checkpoint when it is SDXL, else
the SDXL base. api.ts gains startDiffusionTraining / stopDiffusionTraining /
getDiffusionTrainingStatus plus their types, matching the /api/train/diffusion routes.
* Import diffusion training schemas from models.training directly
The import-hoist lint flags newly re-exported names in the models/__init__.py hub as
unused (it does not treat __all__ membership as a use). Import the three diffusion
training schemas straight from models.training in routes/training.py, where they are
used in the route annotations and calls, and drop the __init__ re-export.
* Remove stray async task scratch outputs committed by mistake
* ControlNet: reject filesystem-like ids and do not cache a model past an unload race
Two review findings on the ControlNet path:
- resolve_controlnet's bare-repo fallback accepted any id with a slash, so a
path-shaped id (/tmp/x, ../x) reached from_pretrained as a local directory.
Restrict the fallback to a strict owner/name HF repo id shape.
- _controlnet_pipe now re-checks the cancel event after the blocking
from_pretrained: an unload that raced the download had already cleared the
caches, so caching the late module would pin it past the unload.
* Pipeline prefetch: fetch only the default torch weights
A full-pipeline prefetch kept every repo file outside assets/, so an official
repo that ships multiple formats (SDXL Base: fp16 variants, ONNX, OpenVINO,
Flax, a top-level single-file twin) downloaded tens of GB from_pretrained never
loads. Skip non-torch exports and dtype-variant twins in
_pipeline_file_downloaded, and drop a component .bin when the same directory
carries a picked safetensors weight (diffusers' own preference).
* Diffusion LoRA training: fall back to fp16 when CUDA lacks bf16
The default mixed_precision=bf16 hard-fails on pre-Ampere GPUs (T4 / V100 /
RTX 20xx) which have no bf16 compute; check torch.cuda.is_bf16_supported()
and drop to fp16 there.
* Diffusion training service: join the old pump outside the lock
start() joined a finished job's pump thread while holding the service lock,
but the pump's final state writes need that same lock, so the join always
burned its full timeout and a stale pump could then overwrite the new job's
state. Join outside the lock (with a re-check after), and fence _apply_event
and the exit handler by process identity so a superseded pump can never touch
the current job's state. Adds regression tests for both.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion LoRA training: harden config handling, cancellation, SDXL conditioning, and safety
Addresses review findings on the SDXL LoRA trainer:
- Gate the base model with the same trust check as inference (unsloth/*, allowlisted
official bases, or a local path) before from_pretrained, so an untrusted remote repo
is never fetched or deserialised.
- Check the stop signal before the (slow) model load, not only between steps, so a
cancel during download is honoured; a stop may carry save=False to cancel without
leaving a partial adapter.
- Per-sample SDXL add_time_ids from the actual crop (original size + crop offset, with
the offset mirrored on horizontal flip) instead of a fixed uncropped-square tensor.
- Apply EXIF orientation before resize/crop so rotated photos train upright.
- Skip gradient clipping when max_grad_norm <= 0 (the Studio 'disable' value) instead
of scaling every gradient to zero.
- Coerce Studio config strings/blanks: learning_rate string to float, blank hf_token to
anonymous, gradient_checkpointing 'none'/'true'/'unsloth' to bool; reject a zero/negative
lora_alpha or learning_rate.
- Alias the generic Studio training payload keys (model_name/max_steps/batch_size/lora_r/
lr_scheduler_type/random_seed) onto the diffusion field names.
- Mirror the trained adapter into loras/diffusion so the Images LoRA picker discovers it.
- Report worker exceptions in both message and error keys so the failure is not lost.
Adds regression tests for the config coercion/validation and aliasing.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* ControlNet: address review findings on the diffusers path
- resolve_controlnet enforces catalog family compatibility so a direct API call
cannot load a ControlNet built for another family through the wrong pipeline.
- Unknown ControlNet ids now surface as a 400 (call site maps FileNotFoundError
to ValueError) instead of a generic 500.
- strength 0 disables ControlNet entirely, so a no-op selection never pays the
download / VRAM cost; the control image is decoded and validated BEFORE the
ControlNet is resolved or built, so a malformed image fails fast for the same reason.
- ControlNet loads use the base compute dtype (state.dtype is a display string,
not a torch.dtype, so it silently fell back to float32) and honor the base
offload policy via group offloading instead of forcing the module resident.
- Empty/malformed HF token coerced to anonymous access.
- Flux Union ControlNet control_mode mapped from the selected control type.
- resolve_controlnet drops the unused hf_token/cancel_event params.
- ControlNetSpec validates guidance_start <= guidance_end (clean 422).
- Images UI ControlNet Select shows its placeholder when nothing is selected.
Adds regression tests for family enforcement and the union control-mode map.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion training API: LLM interlock, pre-spawn VRAM free, path containment, no dropped knobs
Four review findings on the diffusion training start path:
- It spawned the SDXL trainer without checking the LLM TrainingBackend, so a
start while an LLM run was active put two trainers on the same GPU. Add a
symmetric interlock: diffusion start returns 409 when LLM training is active,
and LLM start refuses while a diffusion job is active.
- It went straight to service.start() without freeing GPU residents. Add a
pre-spawn free of the export subprocess, the resident Images pipeline (with an
arbiter release), and chat models, mirroring the LLM start path.
- data_dir / output_dir were passed through unresolved, so Studio-relative names
failed and absolute paths bypassed containment. Resolve them with
resolve_dataset_path / resolve_output_dir before spawn (400 on an uncontained
path).
- The request model dropped max_grad_norm and lora_target_modules, so runs that
set them trained with defaults. Add both fields.
The gemini pump-join deadlock was already fixed earlier (join outside the lock +
proc-identity fence). Note: honoring a stop DURING model load is a trainer-loop
change owned by the diffusion training engine PR (should_stop polled before the
first optimizer step). Adds route + model regression tests.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion LoRA: harden resolution, native tag precedence, and diffusers teardown
Address review findings on the LoRA path:
- resolve_one: normalise a blank/whitespace hf_token to None (anonymous access)
and reject a client-supplied weight file with traversal / absolute path.
- resolve_specs: convert FileNotFoundError from an unknown/stale id to ValueError
so the route returns 400 instead of a generic 500.
- _scan_local: disambiguate local adapters that share a stem (foo.safetensors vs
foo.gguf) so each is uniquely addressable.
- inject_prompt_tags: the backend-validated weight now wins over a user-typed
<lora:ALIAS:...> for a selected adapter; unselected user tags are left alone.
- diffusers _apply_loras: reject a .gguf adapter with a clear error before touching
the pipe (diffusers loads safetensors only).
- _unload_locked: drop the explicit unload_lora_weights() on teardown; the pipe is
dropped wholesale (freeing adapters), so the previous call could race an in-flight
denoise on the same pipe.
- Images page: use a stable LoRA key and clear the selection (not just the options)
when the catalog refresh fails.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion: guard trust check against OSError and validate conditioning inputs
- _is_trusted_diffusion_repo: wrap Path.exists() so a repo id with invalid
characters (or a bare owner/name id) can't raise OSError; treat any failure as
not-a-local-path and fall through to the unsloth/ allowlist. validate_load_request
still raises the clear FileNotFoundError for a genuinely missing local pick.
- generate(): reject mask_image / upscale / reference_images supplied without an
input image, and reject reference_images on a family that does not support
reference conditioning, instead of silently degrading to txt2img / img2img.
* SDXL: reject GGUF up front, skip unused base weights, drop refiner, and harden helpers
Addresses review findings on the SDXL family:
- Reject a GGUF load for single_file_is_pipeline families (SDXL) in validate_load_request,
before the route evicts the current model; SDXL has no transformer-only GGUF variant.
- Skip base-repo weight files when a whole-pipeline single file is loaded: from_single_file
(config=base) needs only the base config/tokenizer/scheduler, so a local .safetensors no
longer triggers a multi-GB base download.
- Remove the SDXL refiner from the non-GGUF trust allowlist: it is an img2img-only pipeline
but this backend loads every sdxl repo as the base txt2img pipeline.
- Normalize a blank/whitespace hf_token to None once in load_pipeline so every load branch
degrades to anonymous instead of erroring on a malformed token.
- Read the denoiser dtype from a parameter (compile-wrapped modules may lack .dtype) and
access state.family.denoiser_attr directly.
Adds/updates regression tests for the trust allowlist, GGUF rejection, and base-config filter.
* Images: preserve restored LoRAs through model load and never send hidden LoRAs
- The LoRA effect cleared the selection on every load->capable transition, which
wiped adapters restored from a gallery recipe before the model finished loading.
Track the previously-loaded family in a ref and clear only on a real family swap;
keep the selection on the initial load and on unload.
- Gate the generate payload's loras on loraCapable so a restored selection that is
hidden (loaded model does not support LoRA) is never sent to the backend.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Images Train LoRA dialog: token, validation, precision, base-repo prefill, gating, refresh
Nine review findings on the SDXL training dialog:
- Forward the saved Hub token so a gated/private SDXL base can be trained (the
image load flow already sends it).
- Re-seed the base-model field from the current default each time the dialog
opens; the keep-alive dialog otherwise kept its mount-time default after a
model loaded.
- Prefill from base_repo (the diffusers pipeline) rather than repo_id, which for
a GGUF/single-file SDXL load is the checkpoint path from_pretrained can't open.
- Add client-side validation of steps/rank/resolution/batch/learning-rate before
the request.
- Expose a precision selector (bf16/fp16/fp32) so non-bf16 GPUs can train from
the UI, not only the API.
- Gate the dialog on the active Images route (active && trainOpen) so switching
tabs closes it and stops its polling.
- Rescan the LoRA picker when a run completes, so a freshly-trained adapter
appears without a model reload.
- Cap the dialog height and scroll the body so the Start/Stop footer stays
reachable on short viewports.
- Correct the copy to not over-promise picker auto-discovery.
Freeing the resident Images pipeline before training is handled backend-side in
the diffusion training start route.
* Merge diffusion-sdxl into diffusion-lora-ux; keep options-only LoRA catch
The catalog-refresh .catch from the lower branch clears the selected adapters
too, which is right for its catalog-only picker but wrong here: this picker
holds free-text HF repo ids that are valid without being in the catalog, so a
transient refresh failure must not wipe them. Family swaps still clear the
selection and hidden LoRAs are never sent.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Train LoRA dialog: stop suggesting absolute paths the backend rejects
The dataset and output placeholders showed /path/to/... examples, but the
training routes resolve those fields inside the Studio home and reject
absolute paths outside the approved roots, so following the placeholder
produced a 400. Use folder-name placeholders and say in the labels and the
dialog description where each folder resolves.
* Align the VAE to the denoiser's first FLOATING dtype, not its first parameter
A GGUF-quantized transformer's leading parameters are packed uint8 storage,
so reading next(parameters()).dtype handed nn.Module.to() an integer dtype
and every image-conditioned generation on a GGUF model (Qwen-Image-Edit)
failed with a 500. Probe the parameters for the first floating dtype, treat
an all-integer module as a no-op, and also catch TypeError so an unexpected
dtype can never break generation. Regression test included.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Count LR scheduler warmup/decay in optimizer steps, not micro-steps
lr_sched.step() runs once per outer optimizer step (after the gradient
accumulation inner loop), for train_steps total. The scheduler was
configured with num_warmup_steps and num_training_steps multiplied by
gradient_accumulation_steps, so with accumulation > 1 a warmup or
non-constant schedule stretched past the run and never reached the
intended decay. Count both in optimizer steps.
* Address Codex review findings on the image-workflows PR
Keep diffusion.py importable without torch: the compile/arch patch modules
import torch at module level, so import them lazily at their load/unload
call sites instead of at module load. This restores the torchless contract
so get_diffusion_backend() works on a CPU/native sd.cpp install.
Match family reject keywords and aliases as whole path/name segments, not
raw substrings, so an unrelated word like edited, edition, or kontextual no
longer misroutes or hides a valid base image model, while supported edit
families (Qwen-Image-Edit, FLUX Kontext) still resolve. Mirror the same
segment matching in the picker task filter.
Route FLUX.2-dev native guidance through --guidance like the other FLUX
families rather than --cfg-scale. Reject native upscale requests that have
no input image. Read image header dimensions and reject over-limit inputs
before decoding pixels, so a crafted small-payload image cannot spike
memory. Reject an upscale that would shrink the source below its input
size. Validate the model_kind against the filename extension before the
GPU handoff. Estimate a local diffusers pipeline's size from its on-disk
weights so auto memory planning does not skip offload and OOM. Report
workflows: [txt2img] from the native backend status so the Create tab
stays enabled for a loaded native model. Clamp the outpaint canvas to the
backend's 4096px decode limit.
Adds regression tests for segment matching and kind/extension validation.
* Guard inference loads and worker lifetime against diffusion training
Teach the chat and image load guards about an active diffusion (SDXL) LoRA
job: a chat load is refused (its footprint cannot be fit-checked against the
trainer) and an image load is refused outright, mirroring the existing LLM
training guards, so a load can no longer allocate GPU memory alongside the
trainer and undo the pre-start cleanup.
Bind the diffusion trainer subprocess to the parent's lifetime and scrub the
native path lease secret from it by running the child through
run_without_native_path_secret, matching the inference/export/LLM workers, so
a Studio crash or kill no longer leaves the trainer holding the GPU.
Reset in_model_load on the complete and error terminal events: a stop or
failure during model loading otherwise leaves the status reporting a stale
loading indicator after the job has ended.
* Harden diffusion LoRA handling on the diffusers and native paths
Reject LoRA on a torch.compile'd diffusers transformer (Speed=default/max):
diffusers requires the adapter loaded before compilation, so applying one to
the already-compiled module fails with adapter-key mismatches. The status
gate now hides the picker and generate raises a clear message instead.
Convert a cancelled Hub LoRA download (RuntimeError Cancelled) to the
diffusion cancellation sentinel in resolve_specs, so an unload/superseding
load during resolution maps to a 409 instead of a generic server error.
Drop weight-0 LoRA rows before the native support gate so a request carrying
only disabled adapters stays a no-op on families where native LoRA is
unsupported, matching the diffusers path.
Reject duplicate LoRA ids in the request model: both apply paths suffix
colliding names, so a repeated id would stack the same adapter past its
per-adapter weight bound.
Strip all user-typed <lora:...> prompt tags on the native path (only the
selected adapters are materialized in the managed lora-model-dir, so an
unselected tag can never resolve), and restore saved LoRA selections from a
gallery recipe so restore reproduces a LoRA image.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Harden ControlNet resolve, gallery metadata, and the control-type picker
Check cancellation immediately after a ControlNet from_pretrained and before
any device placement, so an unload/eviction that raced the download does not
allocate several GB onto the GPU after the load was already cleared.
Require a loadable weight or shard index (not just config.json) before a local
ControlNet folder is advertised, so an interrupted copy is hidden instead of
failing deep in from_pretrained as a generic 500.
Do not record a strength-0 ControlNet in the gallery recipe: it is treated as
disabled and skipped, so the image is unconditioned and the metadata must not
claim a ControlNet was applied.
Build the control-type picker from the selected ControlNet's advertised
control_types instead of a hardcoded passthrough/canny pair, so a union model
with a precomputed depth or pose map sends the correct control_mode.
* Address further Codex findings on the image-workflows PR
- Persist the actual output image size in the gallery recipe instead of the
request sliders: Transform/Inpaint/Edit derive the size from the uploaded
image, Extend grows the canvas, and Upscale resizes it, so the sliders
recorded (and later restored) the wrong dimensions for those workflows.
- Reject a remote '*-GGUF' repo loaded as a full pipeline (no single-file
name) in validate_load_request, so the unloadable pick fails before chat is
evicted rather than deep in from_pretrained.
- Only publish an image-conditioned from_pipe wrapper to the shared aux cache
when the load is still current: from_pipe runs under the generate lock but
not the state lock, so an unload racing its construction could otherwise
cache a wrapper over torn-down modules that a later load would reuse.
- Verify the Windows CUDA runtime archive checksum before extracting it, like
the main sd-cli archive, so a corrupt or tampered runtime is rejected rather
than extracted next to the binary.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Refuse non-SDXL base models at diffusion training start
The trainer only supports the SDXL U-Net, but a FLUX / Qwen-Image / Z-Image
repo or a GGUF filename passed as base_model was accepted and then failed
minutes later inside StableDiffusionXLPipeline.from_pretrained with an
unrelated-looking error. Add a name-based guard in normalized() so known
DiT-family names and .gguf checkpoints are rejected up front, which the API
start route surfaces as an immediate 400 with a message that says exactly
which bases are trainable. Unrecognisable names still pass through so custom
local SDXL checkpoints keep working.
* Add diffusion dataset upload and training info endpoints
Training an image LoRA required knowing the Studio home layout and copying
files onto the server by hand, which is the most confusing step of the whole
flow. Two small endpoints fix that:
- GET /api/train/diffusion/info reports the datasets and outputs roots plus
every dataset folder that contains images (with image/caption counts), so
the UI can offer a picker instead of a blind free-text path.
- POST /api/train/diffusion/dataset uploads images and optional caption
.txt / metadata.jsonl files into a named folder under the datasets root,
creating it on first use and accumulating on repeat uploads so large sets
can arrive in batches. Names are validated to a single path component and
files stream to disk under the same per-upload size cap as LLM dataset
uploads. The returned name is a valid data_dir for /diffusion/start.
* Rework the Train LoRA dialog into a guided SDXL flow
The dialog assumed users knew the Studio home layout and that only SDXL is
trainable, and hid both facts behind free-text fields. Restructure it around
the three real decisions:
- Base model is a dropdown of the trainable SDXL picks (Base 1.0, Turbo, the
loaded SDXL pipeline when there is one) with a custom repo/path escape
hatch, instead of a bare text field defaulting to a repo id.
- Training images come from an in-browser upload (new dataset endpoints) or
a picker over existing dataset folders with image/caption counts. No shell
access or knowledge of the datasets root is needed any more, and the
captioning rules are explained inline.
- The output field is now Adapter name and the instance prompt is labelled
as the trigger prompt, with a no-captions warning wired to the selected
dataset's actual caption count.
Hyperparameters collapse behind a training settings toggle since the
defaults suit a first run. A completed run says where the adapter went and
offers Done / Train another, and the top-bar button gets an icon and a
plainer description. The dialog title states the SDXL-only scope and that
other families load LoRAs but cannot train them yet.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Validate diffusion training config before freeing the GPU
The start route freed resident GPU workloads (export, Images pipeline, chat)
before the service validated the config, so a start that was then refused,
now including a non-SDXL base model, tore down the user's loaded model for
nothing. Run the same cheap normalise pass first; the LLM path already
follows this rule via its before_spawn hook.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Refactor diffusion LoRA training into a family-aware platform
Split the SDXL trainer into a shared, architecture-agnostic layer so more model
families can be trained without duplicating the plumbing:
- New core/training/diffusion_train_common.py holds the config + validation, dataset
discovery, event emission, stop protocol, adapter publishing, and a lazy trainer
registry (get_trainer). diffusion_lora_trainer.py keeps the SDXL-specific loop and
re-exports the moved names so existing imports are unchanged.
- The SDXL-only base-model blocklist becomes a positive check: the family is resolved
from the base model (or an explicit model_family) via the diffusion family registry,
and a known-but-not-yet-trainable family is refused with a clear message. Unknown
custom names still default to the SDXL trainer.
- DiffusionFamily gains a trainable flag and train_base_repos; SDXL is marked trainable.
DiT families flip on when their trainers land.
- Trained adapters now write a <name>.json metadata sidecar (family, base model, rank,
trigger prompt, ...) that the LoRA scanner reads to family-gate the adapter in the
picker instead of showing it as unknown for every model.
- The training base-model trust allowlist adds the official FLUX.1-dev, Z-Image-Turbo,
and Qwen-Image repos (safetensors-only, no remote code).
* Retain diffusion training loss history and expose it in status
The training service kept only the latest loss, so a live loss chart could show a
single point. Fold each progress event into bounded (step, loss, lr) history arrays
(capped at 4000 points, decimated when full) plus the latest throughput and peak VRAM,
and record the family / base model / catalog path on completion. The status endpoint
returns these as a nested metric_history object the UI can chart directly, and the
start request accepts an optional model_family override.
* Tests for the diffusion training platform
Cover the trainer registry (get_trainer resolves SDXL, unknown family raises),
family resolution (explicit model_family validation, resolved_family on the config),
the metadata sidecar write + scan read with family gating, and the service loss-history
folding (append, bad-point skipping, decimation at cap, family/perf fields) plus the
status route nesting metric_history.
* Add flow-matching DiT LoRA trainers (FLUX.1-dev, Qwen-Image, Z-Image)
Extends diffusion LoRA training beyond SDXL to the three popular DiT families
via a single shared flow-matching loop parameterised by small per-family specs
(loading, prompt/latent encoding, transformer forward, save). Verified against
diffusers 0.38.0:
- FLUX.1-dev: 2x2 latent packing + image ids, guidance-embed forward, on-the-fly
nf4 QLoRA of the 12B transformer (the dev repo is gated, so training needs the
user's HF token).
- Qwen-Image: 5D VAE latents normalised by the per-channel latents_mean/std,
img_shapes forward, prequant nf4 base by default (on-the-fly nf4 for the bf16
base).
- Z-Image: list I/O with the reversed timestep convention and a negated
prediction, bf16 only.
The registry (get_trainer) and DiffusionFamily.trainable / train_base_repos now
route these families to the DiT trainer; the SDXL blocklist guard is replaced by
a positive family resolution that also rejects GGUF repos (inference-only) and
still-unsupported families. Per-family defaults + labels + VRAM notes are exposed
via family_train_infos for the Train UI.
Memory: caption embeddings are precomputed once and the text encoders freed
before the loop; gradient checkpointing (non-reentrant, required for bnb 4-bit)
and 8-bit AdamW are on by default.
* Speed up + shrink SDXL LoRA training (precompute text embeds, 8-bit AdamW)
SDXL re-encoded every caption with both CLIP text encoders on every step (pure
waste, since captions are constant) and kept the encoders resident. Precompute
each unique caption's embeddings once, then free the text encoders before the
loop: numerically identical (embeddings are deterministic and this consumes no
torch RNG, so the noise/timestep stream is unchanged) but faster and ~1.5 GB
lighter. Default the optimizer to 8-bit AdamW (bitsandbytes) with an fp32
fallback, halving optimizer state with no meaningful LoRA quality cost. Env
toggles (UNSLOTH_DIFFUSION_NO_PRECOMPUTE / _FP32_OPTIM) let the accuracy guard
A/B the paths.
* Expose trainable families in /diffusion/info and preflight gated bases
The training info endpoint now returns the trainable model families (name,
label, default + allowed base repos, recommended defaults, and a VRAM/access
note) so the Train UI can offer a base picker with realistic guidance. The start
route preflights a gated base repo (HEAD model_index.json with the user's token)
BEFORE freeing resident GPU workloads, so a missing FLUX.1-dev license/token
fails fast with an actionable 400 instead of evicting the loaded model and then
hitting a confusing mid-load 401.
* Tests for DiT trainers, family resolution, info families, gated preflight
Cover the DiT spec table, the QLoRA prequant heuristic, the Z-Image bf16-only
guard, the gated-repo name check, family resolution now that FLUX/Qwen/Z-Image
are trainable (and GGUF repos are rejected as inference-only), the families list
in /diffusion/info, and the gated-base 400 preflight that leaves the GPU
untouched.
* Wrap the DiT training forward in bf16 autocast
The fp32 LoRA parameters and the bnb 4-bit base matmuls need a single
compute dtype during the forward, exactly like the diffusers dreambooth
scripts run under accelerator.autocast. Without it the 4-bit backward on
FLUX.1-dev fails with an illegal-address CUBLAS error partway into the
first step. Z-Image and Qwen-Image smokes are unaffected and the SDXL
path (its own trainer) is untouched.
* [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>
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
* Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source)
Add an opt-in transformer_quant mode that loads the dense bf16 transformer and
torchao-quantises it onto the low-precision tensor cores, instead of the GGUF
transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a
B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's
0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor.
GGUF+compile stays the low-memory default and the fallback. The mode is gated on
CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB);
any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged
reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe
(Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips
the tiny projections that crash int8's torch._int_mm. New module mirrors
diffusion_precision.py; quant runs before compile before placement.
184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route
coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py
is the standalone torchao lever probe.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity
Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838
TFLOPS with FP16 accumulate; bf16 209), so:
- fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on;
pinning it guards consumer cards against a default change. On B200 it is identical
speed and slightly better quality (LPIPS 0.050 vs 0.091).
- the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate).
2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py):
2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the
cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile
(our main ~2x). Documented as a dead end, not shipped.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe
scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image
generation and reports max-abs + non-finite counts for use_fast_accum True vs False.
Confirms fast accumulation is an accumulation-precision knob, not an overflow one:
across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16),
0 non-finite elements and identical max-abs for both modes.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override
Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want
fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and
prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match
on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40;
GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it.
Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s),
no overflow, quality below the quant noise floor. So the default leans to accuracy on
data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force)
lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off).
187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override
threading).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9
scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer.
Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with
use_triton_kernel=False (the default triton path needs the missing MSLK library), but
only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than
using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's
CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto
on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder
Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4
tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS,
2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the
diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both
slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166
vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is
correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add
scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 9): pre-quantized transformer loading
The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on
the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs
13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline
(scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton
on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True)
the quantized weights, so the dense bf16 never touches the GPU.
Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching
GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same
torchao config + min_features filter the runtime path uses, applied ahead of time.
New core/inference/diffusion_prequant.py (resolve_prequant_source +
load_prequantized_transformer, best-effort, lazy imports). diffusion.py
_load_dense_quant_pipeline tries the pre-quant source first and falls back to the
dense materialise+quantise path, then to GGUF, so the default is unchanged.
DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty
prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for
the resolver, the meta-init+assign loader, and the backend branch selection +
fallbacks; GPU verification via scripts/verify_prequant_backend.py.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 10): attention-backend selection
Add a selectable attention kernel via the diffusers set_attention_backend
dispatcher. Attention is memory-bandwidth bound, so a better kernel is an
end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV
matmuls torchao never touches) and composes with torch.compile.
auto picks the best exact backend for the device: cuDNN fused attention
(_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x
end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default
(below the compile/quant noise floor); native SDPA elsewhere and when speed=off
(so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/
xformers/aiter are honored, and an unavailable kernel falls back to the default
rather than failing the load.
New core/inference/diffusion_attention.py (normalize + per-device select + apply,
best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline;
attention_backend threads through begin_load / load_pipeline / status like the other
load knobs. New request field attention_backend + status field. Hermetic CPU tests
for normalize / select policy / apply fallback, plus route threading + 422. Measured
via scripts/perf_levers_probe.py.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder
Consumer / workstation GPUs halve fp8 (and fp16/bf16) FP32-accumulate tensor-core
throughput, while int8 runs at full rate (int32 accumulate is not nerfed). Public
benchmarks (SDNQ across RTX 3090/4090/5090, AMD, Intel) confirm int8 via torch._int_mm
is as fast or faster than fp8 on every consumer part, and the only path on pre-Ada
consumer cards without fp8 tensor cores. So when transformer_quant=auto, reorder the
arch tier to put int8 first on a consumer/workstation GPU (detected by the existing
_is_consumer_gpu name heuristic), while data-center HBM parts keep fp8 first.
Pure ladder reorder via _prefer_consumer_scheme; no new flags. Verified non-regression
on a B200 (still picks fp8). Hermetic tests for consumer Blackwell/Ada/workstation
(-> int8) and data-center Ada/Hopper/Blackwell (-> fp8).
* Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT
Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across
denoise steps a DiT's output settles, so once the first block's residual barely
changes the remaining blocks are skipped and their cached output reused. diffusers
ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the
standalone apply_first_block_cache hook as a fallback).
Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 ->
2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar.
OFF by default and a per-load opt-in: the win scales with step count, so it is for
many-step models (Flux / Qwen-Image) and pointless for few-step distilled models
(e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction
of the trajectory. It composes with regional compile only with fullgraph=False (the
cache's per-step decision is a torch.compiler.disable graph break), which the speed
layer now switches to automatically when a cache is engaged. Best-effort: a model
whose block signature the hook does not recognise is caught and the load proceeds
uncached.
- new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache
(enable_cache / apply_first_block_cache fallback; threshold auto-raised for a
quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import).
- diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph
when a cache is engaged.
- diffusion.py: apply_step_cache before compile; thread transformer_cache /
transformer_cache_threshold through begin_load -> load_pipeline and report the
engaged mode in status().
- models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and
transformer_cache_threshold request fields, engaged mode in the status response.
- hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold
selection, and best-effort failure handling, plus route threading + validation.
- scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM /
LPIPS vs the compiled no-cache baseline).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 14): fix int8 dense quant on Flux / Qwen (skip M=1 modulation linears)
The opt-in dense int8 transformer path crashed on Flux.1 and Qwen-Image with
'torch._int_mm: self.size(0) needs to be greater than 16, but got 1'. int8 dynamic quant
goes through torch._int_mm, which requires the activation row count M > 16. A DiT's AdaLN
modulation projections (Flux norm1.linear 3072->18432, Qwen img_mod.1 / txt_mod.1, Flux.2
*_modulation.linear) and its timestep / guidance / pooled-text conditioning embedders are
computed once from the [batch, dim] conditioning vector (M = batch = 1), not per token, so
they hit _int_mm at M=1 and crash. Their feature dims are large, so the existing
min_features filter did not exclude them.
Fix: the int8 filter now also skips any Linear whose fully-qualified name matches a
modulation / conditioning-embedder token (norm, _mod, modulation, timestep_embed,
guidance_embed, time_text_embed, pooled). These layers run at M=1 once per block and are a
negligible share of the FLOPs, so int8 keeps the full speedup on the attention / FFN layers
(M = sequence length). fp8 / nvfp4 / mxfp8 use scaled_mm, which has no M>16 limit and
quantises these layers fine, so the exclusion is int8-only. Sequence embedders
(context_embedder / x_embedder / txt_in, M = seq) are deliberately not excluded -- note
'context_embedder' contains the substring 'text_embed', which is why the token is the
specific 'time_text_embed', not 'text_embed'.
Measured on a B200 (1024px, transformer_quant=int8 + speed=default), int8 now runs on every
supported model and is the fastest dense path on Flux/Qwen (int8 runs full-rate vs fp8's
FP32-accumulate): FLUX.1-dev 9.62s eager -> 1.98s (4.86x, vs fp8 2.15s), Qwen-Image -> 1.87s
(5.57x, vs fp8 2.09s), FLUX.1-schnell -> 0.41s (3.59x). Z-Image and Flux.2-klein (already
working) are unchanged.
- diffusion_transformer_quant.py: add _INT8_EXCLUDE_NAME_TOKENS; make_filter_fn takes
exclude_name_tokens; quantize_transformer passes it for int8 only.
- hermetic test that the int8 filter excludes the modulation / embedder linears (and keeps
attention / FFN / sequence-embedder linears), while fp8 keeps them.
- scripts/int8_linear_probe.py: the meta-device probe used to enumerate each transformer's
Linear layers and derive the exclusion list.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 15): build int8 pre-quantized checkpoints (skip M=1 modulation linears)
The prequant-checkpoint builder applied the dense quant filter without the int8-only
M=1 modulation / conditioning-embedder exclusion the runtime path uses, so a built int8
checkpoint baked those projections as int8 and crashed (torch._int_mm needs M>16) at the
first denoise step on Flux / Qwen. Factor the scheme->exclusion decision into a shared
exclude_tokens_for_scheme() used by both the runtime quantise path and the offline builder
so they can never drift, and apply it in build_prequant_checkpoint.py. int8 prequant now
produces a working checkpoint on every supported model, giving int8 (the consumer-preferred
scheme) the same ~2x load-VRAM and download reduction fp8 already had.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 16): route no-GPU loads to the native sd.cpp engine
When no CUDA/ROCm/XPU GPU is available, route diffusion load/generate to the
native stable-diffusion.cpp engine instead of diffusers, with diffusers as the
guaranteed fallback. On CPU sd.cpp is 1.4-2.8x faster and uses 1.5-2.2x less RAM.
- diffusion_engine_router: centralised engine selection (built on the existing
select_diffusion_engine), env opt-outs, MPS gating, recorded fallback reason.
- sd_cpp_backend (SdCppDiffusionBackend): the diffusers backend method surface
backed by sd-cli, with lazy binary install, registry-driven asset fetch,
step-progress parsing, and cancellation.
- diffusion_families: per-family single-file VAE + text-encoder asset mapping.
- sd_cpp_engine: cancellation support (process-group kill + SdCppCancelled).
- routes/inference + gpu_arbiter: drive the active engine via the router; the
API now reports the active engine and any fallback reason.
- tests for the backend, router, route selection, and cancellation.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Phase 16 review fixes: engine-switch unload, sd.cpp error mapping, per-image seeds, Qwen sampler
Address review feedback on #6724:
- engine router: unload the engine being deactivated on a switch, so the old
model is not left resident-but-unreachable (the evictor only targets the active
engine).
- generate route: sd.cpp execution errors (nonzero exit / timeout / missing
output) now map to 500, not 409 (which only means not-loaded / cancelled).
- native batch: return per-image seeds and persist the actual seed for each image
so every batch image is reproducible.
- Qwen-Image native path: apply --sampling-method euler --flow-shift 3 per the
stable-diffusion.cpp docs; other families keep sd-cli defaults.
- honor speed_mode (native --diffusion-fa) and, off-CPU, memory_mode/cpu_offload
offload flags on the native load instead of hardcoding them off.
- fail the load when the sd-cli binary is present but not runnable (version()
now returns None on exec error / nonzero exit).
- size estimate: only treat the transformer asset as a possible local path.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in
load_prequantized_transformer ends in torch.load(weights_only=False), which executes
arbitrary code from the pickle. The transformer_prequant_path load-request field reached
that unpickle for any local file an authenticated caller named, so a request could trigger
remote code execution. Refuse the source.kind=='path' branch unless the operator sets
UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted
and unaffected. Document the requirement on the API field and add gate tests.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa
- apply_attention_backend now restores the native default when no backend is requested or a
kernel fails. diffusers keeps a process-wide active attention backend that
set_attention_backend updates, and a fresh transformer's processors follow it, so a load
that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile
load pinned, breaking the bit-identical/off guarantee.
- select_attention_backend drops flash3/flash4 up front when the CUDA capability is below
Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit
request on the wrong card set fine then crashed mid-generation; it now falls back to native.
- Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a
valid alias of native) is accepted instead of 422-rejected by Pydantic.
- Drop the dead replace('-','_') normalization (no alias uses dashes/underscores).
- perf_levers_probe.py output dir is now relative to the script, not a hardcoded path.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF
- apply_step_cache now engages only via the transformer's native enable_cache (the diffusers
CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a
cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin
transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported
transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model
now runs uncached per the best-effort contract.
- GGUF transformers are quantized (the default Studio load path), so they now use the higher
quantized FBCache threshold when the caller leaves it unset, instead of the dense default
that could keep the cache from triggering.
- fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so
fullgraph=True failed warmup and silently measured an eager cached run); output dir is now
relative to the script, not a hardcoded path.
* Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder
_is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as
consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead
of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already
classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the
same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate.
Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0).
* Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config
Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just
TypeError on older torchao), which would break the mxfp8 config helper instead of falling
back to the default. Catch both so the fallback is robust.
quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds
CUDA memory during the per-row VRAM probe; output dir relative to the script.
* Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load
- snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform
missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the
whole snapshot. restore_backend_flags restores each flag independently so one failure can't
leave the others leaked process-wide.
- load_pipeline restores the flags (and clears the GPU cache) when the build fails after
apply_speed_optims mutated the process-wide flags but before _state captured them for unload
to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and
contaminated later off generations.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output
Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs
without producing output (or without closing stdout) would never reach proc.wait and the
wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the
PROCESS, so the main thread always enforces the timeout and kills a hung process (which
closes the pipe and ends the reader). Add a test that times out even when stdout blocks,
and make the no-binary test hermetic so a host-installed sd-cli can't leak in.
* Studio diffusion (Phase 14): guard the int8 exclusion filter against a None fqn
The filter callback can be invoked without a module name, so fqn.lower() would raise
AttributeError on None. Fall back to an empty name (nothing matches the exclusion tokens,
so the linear is kept) instead of crashing the quantise pass.
* Studio diffusion (Phase 16) review fixes: native engine robustness
- sd_cpp_backend: stop truncating explicit seeds to 53 bits (mask to int64);
a large requested seed was silently collapsed (2**53 -> 0) and distinct seeds
aliased to the same image. Random seeds stay 53-bit (JS-safe).
- sd_cpp_backend: sanitize empty/whitespace hf_token to None so HfApi/hf_hub
fall back to anonymous instead of failing auth on a blank token.
- sd_cpp_backend: a superseding load now cancels the in-flight generation, so the
old sd-cli can no longer return/persist an image from the previous model.
- diffusion_engine_router: run the previous engine's unload() OUTSIDE the lock so a
slow 10+ GB free / CUDA sync does not block engine selection.
- diffusion_engine_router: probe sd-cli runnability (version()) before committing to
native, so a present-but-unrunnable binary falls back to diffusers at selection.
- diffusion_device: resolve a torch-free CPU target when torch is unavailable, so a
CPU-only install can still reach the native sd.cpp engine instead of failing load.
- tests updated for the runnability probe + a not-runnable fallback case.
* Studio diffusion (Phase 9) review fixes: prequant safety + validation
- SECURITY: a request-supplied local pre-quant path is now unpickled only when it
resolves inside an operator-configured ALLOWLIST of directories
(UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in,
once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on
any path a load request named (arbitrary code execution). realpath() blocks symlink
escapes; a bare on/off toggle is no longer a wildcard.
- Validate the checkpoint's min_features against the runtime Linear filter, so a
checkpoint that quantised a different layer set is rejected instead of silently
loading a model that mismatches the dense path while reporting the same scheme.
- Tolerant base_model_id compare (exact or same final path/repo segment), so a local
path or fork of the canonical base is accepted instead of falling back to dense.
- _has_meta_tensors uses any(chain(...)) (no intermediate lists).
- prequant verify/probe scripts use repo-relative paths (+ env overrides), not the
author's absolute /mnt paths.
- tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts
- diffusion_memory: when group offload is unavailable and the plan falls back to
whole-module offload, enable VAE tiling (the group plan left it off, but the fallback
is the low-VRAM path where the decode spike can OOM). Covers both the group and
sequential fallback branches.
- perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a
balanced bit-identity regression actually fails the check instead of exiting 0.
- compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and
the load-progress poll has a 30 min deadline instead of looping forever on a hang.
- test for the group->model fallback enabling VAE tiling.
* Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path
- diffusion: a torchao-quantized transformer is committed only compiled. A dense model
resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF
it replaced), so when transformer_quant engaged and speed resolved to off, promote to
default (regional compile); warn loudly if compile still does not engage.
- diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the
CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK);
otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF.
- nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path).
- test asserts the eager-quant -> default-compile promotion.
* Studio diffusion (Phase 10) review fixes: attention gating + probe isolation
- diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere
NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation,
so auto now stays on native SDPA there.
- diffusion_attention: _active_attention_backend handles get_active_backend() returning an
enum/None (not a tuple); the old unpack always raised and was swallowed, so
the native-restore short-circuit never fired.
- perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS
on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul
so the inductor_flags variant doesn't leak into later compiled rows.
- tests for the SM80 cuDNN gate.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening
- install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit
timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket);
extract through a per-member containment check (Zip-Slip guard); expanduser the
--install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch
the separately-published cudart runtime DLL archive so sd-cli.exe can start.
- sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the
installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start
sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend
crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie.
- tests: Zip-Slip rejection, normal extraction, studio-home discovery.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs
Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes
the numbered files <stem>_<idx><suffix> (base_0.png, base_1.png, ...) instead of
the literal --output path. SdCppEngine.generate checked only the literal path, so
a batch generation would exit 0 and then raise 'no image' (or return a stale
file). generate now returns the literal path when present and otherwise falls
back to the numbered siblings; single-image behavior is unchanged.
Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected
without error.
* Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats
Codex review on the native engine arg builder:
- build_sd_cpp_command emitted --width/--height unconditionally, so an
img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of
the input. width/height are now Optional (None = unset): an image-conditioned
run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives
the size from the input image (set_width_and_height_if_unset); a plain txt2img
run with unset dims keeps the prior 1024x1024 default; explicit dims are always
honored. width/height are read only by the builder, so the type change is local.
- build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...)
that silently swallowed repeats=0 into sd-cli's default of one pass, turning an
explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError
and emits the flag for any explicit value != 1.
Tests: img2img unset dims omit width/height (init_img and ref_images), explicit
dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at
the default. (Two pre-existing binary-discovery tests fail only because a real
sd-cli is installed in this dev environment; unrelated to this change.)
* Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc
Codex review: the transformer_prequant_path field description still told operators
to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the
prior security fix made that variable a directory allowlist -- _allowed_prequant_roots
deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator
following the documented =1 would have every transformer_prequant_path request
silently refused. The description now states it must name one or more allowlisted
directories and that a bare on/off value is not accepted.
Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does
not say =1, and describes an allowlist/directory (guards against doc drift).
* Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset
Codex review on attention-backend selection:
- Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on
pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first
generation with no fallback. select_attention_backend now applies
_cudnn_attention_supported() to an explicit cuDNN request too.
- flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a
Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3
is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a
(min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+.
- apply_attention_backend's success path left diffusers' process-wide active
backend pinned to the kernel it set; a later component whose processors are
unconfigured (backend None) would inherit it. It now resets the global registry
to native after a successful per-transformer set (the transformer keeps its own
backend), best-effort. Also fixed _active_attention_backend: get_active_backend()
returns a (name, fn) tuple, so the prior code stringified the tuple and never
matched a name, defeating the native-restore short-circuit.
Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on
SM90; global registry reset after a successful set; _active_attention_backend
reads the tuple return.
* Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder
Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300
(Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So
_is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the
auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting
llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they
are treated as data-center class and keep the intended fp8-first behavior.
Test: extends the datacenter parametrize with 'NVIDIA B300' and
'NVIDIA GH200 480GB' (now _is_consumer_gpu False).
* Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder
Codex review: the M=1 modulation/embedder exclusion was wired only into the dense
runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called
make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint
quantised the AdaLN modulation and conditioning-embedder linears, and loading it
via transformer_prequant_path (the load path only loads already-quantised tensors,
it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes
for the runtime path.
Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 ->
the M=1 exclusion, every other scheme -> none) and use it in both the runtime
quantiser and the builder, so a prequant artifact's quantised-layer set always
matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion).
Test: int8_exclude_name_tokens returns the exclusion for int8 and () for
fp8/nvfp4/mxfp8.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 16) review round 2: native CPU arbiter, status offload, load race
Codex review on the native-engine routing:
- The /images/load route took the GPU arbiter (acquire_for(DIFFUSION) -> evict chat)
unconditionally after engine selection. A native sd.cpp load on a pure-CPU host
never touches the GPU, so that needlessly tore down the resident chat model. The
handoff is now gated: diffusers always takes it, a force-native sd.cpp load on a
CUDA/XPU/MPS box still takes it, but a native sd.cpp load on a CPU host skips it.
- sd_cpp status() hardcoded offload_policy 'none' / cpu_offload False even when
_run_load computed real offload flags (balanced/low_vram/cpu_offload off-CPU), so
the setting was unverifiable. status now derives them from state.offload_flags
(still 'none' on CPU, where the flags are empty).
- _run_load committed the new state without cancelling/waiting on a generation that
started during the (slow) asset download, so a stale sd-cli run against the OLD
model could finish afterward and persist an image from the previous model once the
new load reported ready. The commit now signals the in-flight cancel and waits on
_generate_lock before swapping _state (taken only at commit, so the download never
serialises against generation), mirroring the diffusers load path.
Tests: CPU native load skips the arbiter while a GPU native load takes it; status
reports offload active when flags are set; _run_load cancels and waits for an
in-flight generation before committing.
* Studio diffusion (Phase 14) review round 2: align helper name with the stack
Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the
identical helper already present higher in the diffusion stack (Phase 16). The
helper definition, the runtime quantiser call, and the offline builder are now
byte-identical to that version, so the two branches no longer introduce a
divergent name for the same single-source-of-truth and the stack merges without
a conflict on this fix. No behavior change.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion: eager patches + torch.compile cache speed phase
Adds the opt-in speed path for the GGUF diffusion transformer behind a
selectable speed mode (default off, so output is unchanged until a profile
is chosen):
- diffusion_eager_patches.py: shared eager fast-paths (channels_last,
attention/backend selection, fused norms and QKV) installed at load and
rolled back on unload or failed load.
- diffusion_compile_cache.py / diffusion_gguf_compile.py: a persistent
torch.compile cache and the GGUF-transformer compile wiring.
- diffusion_arch_patches.py: architecture-specific patches.
- diffusion_patch_backend.py: shared install/restore plumbing.
- diffusion_speed.py: speed-profile planning.
Tests for each module plus the benchmarking and probe scripts used to
measure speed, memory, and accuracy of the path.
* Studio diffusion: image workflows (safetensors, image-conditioned, editing) + Images UI
Backend:
- Load non-GGUF safetensors models: full bnb-4bit pipelines and single-file
fp8 transformers, gated to the unsloth org plus a curated allowlist.
- Image-conditioned workflows built with Pipeline.from_pipe so they reuse the
loaded transformer/VAE/text-encoder with no extra VRAM: img2img, inpaint,
outpaint, and a hires-fix upscale pass.
- Instruction editing as its own family kind (Qwen-Image-Edit-2511,
FLUX.1-Kontext-dev) and FLUX.2-klein reference conditioning (single and
multi-reference) plus klein inpaint.
- Auto-resize odd-sized inputs to a multiple of 16 (and resize the matched
mask) so img2img/inpaint/edit no longer reject non-/16 uploads. Bound the
decoded image size and cap upscale output to avoid OOM on large inputs.
- Fixes: from_pipe defaulting to a float32 recast that crashed torchao
quantized transformers; image-conditioned calls forcing the slider size
onto the input image. Native sd.cpp engine rejects image-conditioned and
reference requests it cannot serve.
Frontend:
- Redesigned Images page with capability-gated workflow tabs (Create,
Transform, Inpaint, Extend, Upscale, Reference, Edit), a brush mask editor,
client-side outpaint, and a multi-reference picker.
- Advanced options moved to a right-docked panel mirroring Chat: closed by
default, toggled by a single fixed top-bar button that stays in place.
sd.cpp installer: pin the release, verify each download's sha256, add a
download timeout, and make the source repo configurable for a future mirror.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio Images: correct the Advanced panel comment (closed by default, fixed toggle)
* Studio diffusion: LoRA adapters for the Images workflow
Add community LoRA support across both diffusion backends, the single
biggest step toward broad image-workflow coverage.
Backend
- New shared module core/inference/diffusion_lora.py: adapter discovery
(local scan + curated catalog + owner/name[:file] Hub refs), download
via hf_hub_download_with_xet_fallback, alias sanitization, native
managed-dir materialization with collision-broken aliases, prompt-tag
injection (deduped against user-typed tags), and a supports_lora gate.
- Native sd-cli: resolve + materialize selected LoRAs into a per-run
managed dir, inject <lora:ALIAS:w> tags, pass --lora-model-dir with
--lora-apply-mode auto. The arg builder already emitted these flags.
- Diffusers: non-fused load_lora_weights + set_adapters manager, tracked
on the pipe so an unchanged selection is a no-op and a model swap
resets; cleared on unload. Never fuses (breaks quantized transformers
and blocks live weight tweaks).
- Gated off where unsupported: torchao fp8/int8 dense, GGUF-via-diffusers,
and native Qwen-Image (no LoRA name-conversion branch upstream).
- Request contract: optional loras on DiffusionGenerateRequest; empty or
omitted is identical to today. supports_lora surfaced in status; chosen
LoRAs persisted in gallery recipe metadata.
- New GET /api/models/diffusion-loras for the picker (family-filtered).
Frontend
- Repeatable multi-LoRA picker (adapter select + weight slider 0..2 +
remove), gated by the loaded model's supports_lora and family, max 8.
Tests
- New test_diffusion_lora.py (14): helpers, request validation, native
tag/dir wiring, diffusers set_adapters manager, supports_lora matrix.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion: ControlNet for the Images workflow (diffusers)
Add ControlNet conditioning, the #2 most-used diffusion workflow after
LoRA, on the diffusers backend for the families with ControlNet pipelines
(FLUX.1 and Qwen-Image), with Union models as the default picks.
Backend
- New core/inference/diffusion_controlnet.py: family-gated discovery
(curated Union models + local dirs + bare owner/name repos), resolution
to a loadable repo/dir, control-image preprocessing (passthrough +
a dependency-free canny edge map), and a supports_controlnet gate.
- diffusion.py: a ControlNet manager parallel to the LoRA one. Loads the
(small) ControlNet model once via from_pretrained (cached by id) and
builds the family's ControlNet pipeline via Pipeline.from_pipe(base,
controlnet=model), reusing the resident base modules at their loaded
dtype (no reload, no recast). Passes the control image + conditioning
scale + guidance start/end at generate time; cleared on unload.
- Families: FLUX.1 -> FluxControlNetPipeline/Model, Qwen-Image ->
QwenImageControlNetPipeline/Model. Others declare none (gated off).
- Gated off for the native engine, GGUF-via-diffusers, and torchao
fp8/int8 dense (same rule as LoRA). v1 conditions txt2img only.
- Request contract: optional controlnet on DiffusionGenerateRequest;
supports_controlnet in status; the choice persisted in gallery meta.
- New GET /api/models/diffusion-controlnets for the picker.
Frontend
- A ControlNet control in the Images rail (model select + control-image
upload + control-type select + strength slider), gated by the loaded
model's supports_controlnet + family, shown for text-to-image.
Tests
- New test_diffusion_controlnet.py (10): discovery/resolve/preprocess/gate
helpers, request validation, family wiring, and the diffusers pipe
manager (loads once, caches, from_pipe with controlnet, rejects
unsupported families).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio ControlNet: show the picker on the Create tab (workflow id is 'create', not 'txt2img')
The ControlNet control gated on workflow === 'txt2img', but the Images workflow tab ids are create/transform/inpaint/extend/upscale/reference/edit -- there is no 'txt2img'. So the picker never rendered even with a ControlNet-capable model loaded. Gate on 'create' (the text-to-image tab) for both the picker and the request wiring. Found via a live Playwright capture of the running Studio.
* Studio: do not force diffusers pipelines cross-tagged gguf into the GGUF variant expander
Some diffusers image repos (e.g. unsloth/Qwen-Image-2512-unsloth-bnb-4bit) carry a
stray "gguf" tag on the Hub but ship no .gguf files. The model search classified
them as GGUF from the bare tag, so the picker rendered the GGUF variant expander,
which then dead-ended at "No GGUF variants found." Trust the bare gguf tag only when
the repo is not a diffusers pipeline; the -GGUF name suffix and real gguf metadata
(populated via expand=gguf) remain authoritative, so genuine GGUF repos are unaffected.
* Studio Images: load non-curated unsloth/on-device diffusers repos instead of no-op
handleModelSelect only loaded curated safetensors ids and GGUF variant picks; any other
non-GGUF pick (an on-device diffusers folder, or a future unsloth diffusers image repo
surfaced by search) silently did nothing. Treat such a pick as a full diffusers pipeline
load when the id is unsloth-hosted or on-device (the backend infers the family + base repo
and gates loads to unsloth/* or local paths), and show a clear message otherwise instead
of silently ignoring the click. Curated and GGUF paths are unchanged.
* Studio Images: keep curated safetensors models in Recommended after download
The curated bnb-4bit / fp8 diffusion rows were filtered out of the Images picker's
Recommended list once cached (curatedSafetensorsRows dropped anything in downloadedSet),
so they vanished from the picker after the first load and could only be found by typing an
exact search. The row already renders a downloaded badge, matching how GGUF Recommended
rows stay visible when cached. Drop the exclusion so the curated safetensors always list.
* Studio diffusion LoRA: sanitize dots out of adapter aliases
The LoRA alias is used as the diffusers PEFT adapter name, and PEFT rejects names
containing "." (module name can't contain "."). sanitize_alias kept dots, so a LoRA whose
filename carries a version tag (e.g. Qwen-Image-2512-Lightning-8steps-V1.0-bf16) failed to
apply with a 400. Replace dots too; the alias stays a valid native <lora:NAME:w> filename
stem. Adds regression coverage for internal dots.
* Studio Images: clarify the GGUF transformer-quant Advanced control
Renamed the confusing "Transformer quant / GGUF default" control to "GGUF speed mode"
with an "Off (run the GGUF)" default, and reworded the hint to state plainly that FP8/INT8/
FP4 load the FULL base model (larger download + more VRAM) rather than re-packing the GGUF,
falling back to the GGUF if it can't fit. Behavior unchanged; labels/hint only.
* Studio Images: list on-device unsloth diffusion models in the picker
The Images picker's On Device tab hid every non-GGUF cached repo whenever a
task filter was active, so downloaded unsloth diffusion pipelines (bnb-4bit
and FP8 safetensors) never showed up there. List cached repos that pass the
task gate, limited under a filter to unsloth-hosted ones so base repos (which
fail the diffusion load trust gate) don't appear only to dead-end on click.
Chat behavior is unchanged: the task gate still drops image repos there.
* Studio: hide single-file image checkpoints from the chat model picker
The chat picker treats a cached repo as an image model, and hides it, only
when it ships a diffusers model_index.json. Single-file, ComfyUI, and
ControlNet image checkpoints (an FP8 Qwen-Image, a z-image safetensors, a
Qwen-Image ControlNet) carry none, so they surfaced as loadable chat models.
Fall back to resolving the repo id against the known diffusion families, the
same resolver the Images backend loads from, so these checkpoints are tagged
text-to-image and stay in the Images picker only.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio Images: add the FLUX.2-dev model family
Loading unsloth/FLUX.2-dev-GGUF failed because detect_family knew only the
Qwen3-based FLUX.2-klein, so FLUX.2-dev (the full, Mistral-based Flux2Pipeline)
resolved to nothing and the load errored. Add a flux.2-dev family: Flux2Pipeline
+ Flux2Transformer2DModel over the black-forest-labs/FLUX.2-dev base repo (gated,
reachable with an HF token), with its FLUX.2 32-channel VAE and Mistral text
encoder wired for the sd-cli path from the open Comfy-Org/flux2-dev mirror.
text-to-image only: diffusers 0.38 ships no Flux2 img2img / inpaint pipeline for
dev. Frontend gets sensible dev defaults (28 steps, guidance 4), distinct from
klein's turbo defaults. Verified live: GGUF load resolves the family + gated base
repo and generates a real 1024x1024 image on GPU.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio Images: clearer error for an unsupported diffusion model
When a repo id resolves to no diffusion family the load raised 'Could not infer a
diffusion family... Pass family_override (z-image)', which points at an unrelated
family and doesn't say what is supported. Replace it with a message that lists the
supported families (from a new supported_family_names helper) and notes that video
models and image models whose diffusers transformer has no single-file loader are
not supported. Applies to both the diffusers and native sd.cpp load paths. Also
refreshes two stale family-registry comments that still called FLUX.2-dev omitted.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add SDXL diffusion family (U-Net pipeline support)
SDXL is the first U-Net family in the diffusion backend: its denoiser is
pipe.unet (UNet2DConditionModel), not a DiT pipe.transformer, and a single-file
.safetensors is the whole pipeline rather than a transformer-only file. The
backend previously assumed a DiT transformer everywhere, so add the two hooks a
U-Net family needs and register SDXL.
DiffusionFamily gains denoiser_attr ("transformer" for DiT, "unet" for SDXL) and
single_file_is_pipeline (SDXL loads a single file via pipeline_class.from_single_file
with the base repo as config, instead of transformer_class.from_single_file plus a
companion assembly). _align_vae_dtype now reads the denoiser generically so img2img
and inpaint keep the VAE and U-Net dtypes aligned.
The non-GGUF trust gate is extended with a short, exact-match, safetensors-only
allowlist of official base repos (the SDXL base/refiner and sdxl-turbo), because
SDXL ships only as a full pipeline and has no unsloth-hosted GGUF. Local paths stay
trusted as before; a random repo, even one that detects as SDXL, is still rejected.
The image-conditioned and ControlNet workflows are the standard SDXL pipelines,
built around the resident modules via from_pipe like every other family, so SDXL
gets txt2img, img2img, inpaint, outpaint, upscale, LoRA and ControlNet. There is no
native sd.cpp mapping yet, so the no-GPU route falls back to diffusers.
Frontend catalog gains SDXL Base 1.0 and SDXL Turbo entries with SDXL step/guidance
defaults (Turbo: few steps, no CFG; base: ~30 steps, real CFG).
Tests: new test_diffusion_sdxl.py (family shape, detection, trust allowlist, model
kind, U-Net VAE-dtype alignment, LoRA gate) plus loader-branch tests in
test_diffusion_backend.py (pipeline-kind from_pretrained, single-file whole-pipeline
from_single_file, allowlist accept/reject). Verified live on GPU: sdxl-turbo loads
both as a pipeline and as a single file and generates coherent txt2img + img2img.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Images: LoRA free-text Hugging Face entry + recipe round-trip
The backend has always accepted a bare Hugging Face repo id (owner/name, or
owner/name:weight-file.safetensors) as a LoRA, downloading and applying it. But the
picker only rendered when the curated catalog had entries, and the catalog is empty,
so there was no UI path to apply any LoRA. Show the LoRA section whenever the loaded
model supports LoRA, and replace the curated-only dropdown with a text input: type a
Hub repo id, or pick a discovered adapter from a datalist of suggestions when the
catalog is populated.
Also restore LoRAs when loading a recipe. restoreSettings now parses the recipe's
"id:weight" strings (splitting on the last colon, since the id itself may contain one
for a specific weight file) back into the selection, so replaying a saved image
reproduces its adapters. The generate payload trims hand-typed ids and drops empty /
zero-weight rows, and a model swap clears the selection (a LoRA is family-specific)
without discarding a free-text pick that is not in the curated list.
* Add diffusion LoRA training (SDXL text-to-image)
First diffusion training path in Studio: train a LoRA on the SDXL U-Net from an
image + caption dataset and export it as a diffusers .safetensors that the existing
diffusion LoRA loader (and any diffusers pipeline) can load.
core/training/diffusion_lora_trainer.py:
- DiffusionLoraConfig with validation/defaults (rank, alpha, targets, lr, steps, grad
accumulation, resolution, min-SNR gamma, gradient checkpointing, lr scheduler, seed,
mixed precision).
- discover_image_caption_pairs: captions from metadata.jsonl / captions.jsonl, per-image
.txt/.caption sidecars, or a dreambooth instance_prompt fallback (pure, unit-tested).
- run_diffusion_lora_training: the loop -- freeze base, PEFT-wrap the U-Net attention
projections, VAE-encode (fp32 VAE to avoid the SDXL fp16 overflow), sample noise +
timesteps, predict, MSE loss with optional min-SNR weighting (epsilon / v-prediction),
AdamW + get_scheduler + grad accumulation + grad clipping, then export via
save_lora_weights. Emits worker-protocol events (model_load_*, progress, complete) and
polls should_stop for a clean stop with a partial save.
- run_diffusion_training_process: mp.Queue subprocess adapter (event_queue / stop_queue),
so the training worker can spawn it; plus a CLI entry point.
Only SDXL (U-Net) is trained here; DiT families and the Studio UI form + route wiring are
follow-ups. The trainer is decoupled and worker-ready.
Tests: test_diffusion_lora_trainer.py covers caption discovery (metadata / sidecar /
instance prompt / skip-uncaptioned / errors), config normalisation + validation, the SDXL
add-time-ids, and the dict->config adapter. Verified live on GPU: a 60-step SDXL LoRA run
lowers the loss, exports a ~45 MB adapter, and loading it back shifts generation from
baseline (mean abs pixel diff ~55/255).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* diffusion trainer: emit learning_rate in progress events (Studio pump compatibility)
The Studio training pump reads 'learning_rate' from progress events; the diffusion
trainer emitted 'lr'. Rename the field (and the CLI reader) so the trainer's events are
directly consumable by the existing training status/SSE machinery when it is wired into
the worker, without a translation shim.
* Wire diffusion LoRA training into the Studio API
Make the SDXL LoRA trainer reachable from the app with a small, self-contained job
service and JSON routes, deliberately separate from the LLM TrainingBackend (whose
lifecycle -- LLM config build, per-run SQLite rows, matplotlib plots, transfer-to-chat-
inference -- is text-training specific and would mis-handle a diffusion run).
core/training/diffusion_training_service.py: DiffusionTrainingService runs one job at a
time -- validate the config cheaply (before any spawn), spawn the trainer subprocess
(spawn context, parent-lifetime bound), pump its events (model_load_* / progress /
complete / error) into an in-memory status snapshot, and support a clean stop. The
subprocess context and target are injectable so the full start -> pump -> status ->
complete path is unit-tested without real multiprocessing or torch.
routes/training.py: POST /api/train/diffusion/start (400 on a bad config, 409 when a job
is already running), POST /api/train/diffusion/stop, GET /api/train/diffusion/status
(JSON poll). models/training.py: DiffusionTrainingStartRequest + response schemas
mirroring DiffusionLoraConfig, so model_dump() passes straight through.
Tests: test_diffusion_training.py -- service happy path, bad-config-before-spawn,
concurrent-job rejection, clean stop, crash-without-terminal-event, event transitions;
plus route wiring via the FastAPI TestClient (start / 422 / 400 / 409 / status / stop)
with a mocked service. The diffusion trainer's progress events already use the field
names this path expects.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Images: add a Train LoRA (SDXL) dialog
Surface the diffusion training API in the Images page. A "Train LoRA" button in the top
bar opens a self-contained dialog to fine-tune an SDXL LoRA on a folder of images: pick
the base model, dataset folder, output folder, an optional instance prompt, and the core
hyperparameters (steps, rank, resolution, batch, learning rate), then Start. The dialog
polls the training status while open and shows a progress bar, step count, live loss, and
the saved adapter path, with a Stop button for a clean stop.
The dialog is independent of the loaded generation model (training runs in its own
subprocess), and prefills the base model with the loaded checkpoint when it is SDXL, else
the SDXL base. api.ts gains startDiffusionTraining / stopDiffusionTraining /
getDiffusionTrainingStatus plus their types, matching the /api/train/diffusion routes.
* Import diffusion training schemas from models.training directly
The import-hoist lint flags newly re-exported names in the models/__init__.py hub as
unused (it does not treat __all__ membership as a use). Import the three diffusion
training schemas straight from models.training in routes/training.py, where they are
used in the route annotations and calls, and drop the __init__ re-export.
* Remove stray async task scratch outputs committed by mistake
* ControlNet: reject filesystem-like ids and do not cache a model past an unload race
Two review findings on the ControlNet path:
- resolve_controlnet's bare-repo fallback accepted any id with a slash, so a
path-shaped id (/tmp/x, ../x) reached from_pretrained as a local directory.
Restrict the fallback to a strict owner/name HF repo id shape.
- _controlnet_pipe now re-checks the cancel event after the blocking
from_pretrained: an unload that raced the download had already cleared the
caches, so caching the late module would pin it past the unload.
* Pipeline prefetch: fetch only the default torch weights
A full-pipeline prefetch kept every repo file outside assets/, so an official
repo that ships multiple formats (SDXL Base: fp16 variants, ONNX, OpenVINO,
Flax, a top-level single-file twin) downloaded tens of GB from_pretrained never
loads. Skip non-torch exports and dtype-variant twins in
_pipeline_file_downloaded, and drop a component .bin when the same directory
carries a picked safetensors weight (diffusers' own preference).
* Diffusion LoRA training: fall back to fp16 when CUDA lacks bf16
The default mixed_precision=bf16 hard-fails on pre-Ampere GPUs (T4 / V100 /
RTX 20xx) which have no bf16 compute; check torch.cuda.is_bf16_supported()
and drop to fp16 there.
* Diffusion training service: join the old pump outside the lock
start() joined a finished job's pump thread while holding the service lock,
but the pump's final state writes need that same lock, so the join always
burned its full timeout and a stale pump could then overwrite the new job's
state. Join outside the lock (with a re-check after), and fence _apply_event
and the exit handler by process identity so a superseded pump can never touch
the current job's state. Adds regression tests for both.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion LoRA training: harden config handling, cancellation, SDXL conditioning, and safety
Addresses review findings on the SDXL LoRA trainer:
- Gate the base model with the same trust check as inference (unsloth/*, allowlisted
official bases, or a local path) before from_pretrained, so an untrusted remote repo
is never fetched or deserialised.
- Check the stop signal before the (slow) model load, not only between steps, so a
cancel during download is honoured; a stop may carry save=False to cancel without
leaving a partial adapter.
- Per-sample SDXL add_time_ids from the actual crop (original size + crop offset, with
the offset mirrored on horizontal flip) instead of a fixed uncropped-square tensor.
- Apply EXIF orientation before resize/crop so rotated photos train upright.
- Skip gradient clipping when max_grad_norm <= 0 (the Studio 'disable' value) instead
of scaling every gradient to zero.
- Coerce Studio config strings/blanks: learning_rate string to float, blank hf_token to
anonymous, gradient_checkpointing 'none'/'true'/'unsloth' to bool; reject a zero/negative
lora_alpha or learning_rate.
- Alias the generic Studio training payload keys (model_name/max_steps/batch_size/lora_r/
lr_scheduler_type/random_seed) onto the diffusion field names.
- Mirror the trained adapter into loras/diffusion so the Images LoRA picker discovers it.
- Report worker exceptions in both message and error keys so the failure is not lost.
Adds regression tests for the config coercion/validation and aliasing.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* ControlNet: address review findings on the diffusers path
- resolve_controlnet enforces catalog family compatibility so a direct API call
cannot load a ControlNet built for another family through the wrong pipeline.
- Unknown ControlNet ids now surface as a 400 (call site maps FileNotFoundError
to ValueError) instead of a generic 500.
- strength 0 disables ControlNet entirely, so a no-op selection never pays the
download / VRAM cost; the control image is decoded and validated BEFORE the
ControlNet is resolved or built, so a malformed image fails fast for the same reason.
- ControlNet loads use the base compute dtype (state.dtype is a display string,
not a torch.dtype, so it silently fell back to float32) and honor the base
offload policy via group offloading instead of forcing the module resident.
- Empty/malformed HF token coerced to anonymous access.
- Flux Union ControlNet control_mode mapped from the selected control type.
- resolve_controlnet drops the unused hf_token/cancel_event params.
- ControlNetSpec validates guidance_start <= guidance_end (clean 422).
- Images UI ControlNet Select shows its placeholder when nothing is selected.
Adds regression tests for family enforcement and the union control-mode map.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion training API: LLM interlock, pre-spawn VRAM free, path containment, no dropped knobs
Four review findings on the diffusion training start path:
- It spawned the SDXL trainer without checking the LLM TrainingBackend, so a
start while an LLM run was active put two trainers on the same GPU. Add a
symmetric interlock: diffusion start returns 409 when LLM training is active,
and LLM start refuses while a diffusion job is active.
- It went straight to service.start() without freeing GPU residents. Add a
pre-spawn free of the export subprocess, the resident Images pipeline (with an
arbiter release), and chat models, mirroring the LLM start path.
- data_dir / output_dir were passed through unresolved, so Studio-relative names
failed and absolute paths bypassed containment. Resolve them with
resolve_dataset_path / resolve_output_dir before spawn (400 on an uncontained
path).
- The request model dropped max_grad_norm and lora_target_modules, so runs that
set them trained with defaults. Add both fields.
The gemini pump-join deadlock was already fixed earlier (join outside the lock +
proc-identity fence). Note: honoring a stop DURING model load is a trainer-loop
change owned by the diffusion training engine PR (should_stop polled before the
first optimizer step). Adds route + model regression tests.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion LoRA: harden resolution, native tag precedence, and diffusers teardown
Address review findings on the LoRA path:
- resolve_one: normalise a blank/whitespace hf_token to None (anonymous access)
and reject a client-supplied weight file with traversal / absolute path.
- resolve_specs: convert FileNotFoundError from an unknown/stale id to ValueError
so the route returns 400 instead of a generic 500.
- _scan_local: disambiguate local adapters that share a stem (foo.safetensors vs
foo.gguf) so each is uniquely addressable.
- inject_prompt_tags: the backend-validated weight now wins over a user-typed
<lora:ALIAS:...> for a selected adapter; unselected user tags are left alone.
- diffusers _apply_loras: reject a .gguf adapter with a clear error before touching
the pipe (diffusers loads safetensors only).
- _unload_locked: drop the explicit unload_lora_weights() on teardown; the pipe is
dropped wholesale (freeing adapters), so the previous call could race an in-flight
denoise on the same pipe.
- Images page: use a stable LoRA key and clear the selection (not just the options)
when the catalog refresh fails.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion: guard trust check against OSError and validate conditioning inputs
- _is_trusted_diffusion_repo: wrap Path.exists() so a repo id with invalid
characters (or a bare owner/name id) can't raise OSError; treat any failure as
not-a-local-path and fall through to the unsloth/ allowlist. validate_load_request
still raises the clear FileNotFoundError for a genuinely missing local pick.
- generate(): reject mask_image / upscale / reference_images supplied without an
input image, and reject reference_images on a family that does not support
reference conditioning, instead of silently degrading to txt2img / img2img.
* SDXL: reject GGUF up front, skip unused base weights, drop refiner, and harden helpers
Addresses review findings on the SDXL family:
- Reject a GGUF load for single_file_is_pipeline families (SDXL) in validate_load_request,
before the route evicts the current model; SDXL has no transformer-only GGUF variant.
- Skip base-repo weight files when a whole-pipeline single file is loaded: from_single_file
(config=base) needs only the base config/tokenizer/scheduler, so a local .safetensors no
longer triggers a multi-GB base download.
- Remove the SDXL refiner from the non-GGUF trust allowlist: it is an img2img-only pipeline
but this backend loads every sdxl repo as the base txt2img pipeline.
- Normalize a blank/whitespace hf_token to None once in load_pipeline so every load branch
degrades to anonymous instead of erroring on a malformed token.
- Read the denoiser dtype from a parameter (compile-wrapped modules may lack .dtype) and
access state.family.denoiser_attr directly.
Adds/updates regression tests for the trust allowlist, GGUF rejection, and base-config filter.
* Images: preserve restored LoRAs through model load and never send hidden LoRAs
- The LoRA effect cleared the selection on every load->capable transition, which
wiped adapters restored from a gallery recipe before the model finished loading.
Track the previously-loaded family in a ref and clear only on a real family swap;
keep the selection on the initial load and on unload.
- Gate the generate payload's loras on loraCapable so a restored selection that is
hidden (loaded model does not support LoRA) is never sent to the backend.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Images Train LoRA dialog: token, validation, precision, base-repo prefill, gating, refresh
Nine review findings on the SDXL training dialog:
- Forward the saved Hub token so a gated/private SDXL base can be trained (the
image load flow already sends it).
- Re-seed the base-model field from the current default each time the dialog
opens; the keep-alive dialog otherwise kept its mount-time default after a
model loaded.
- Prefill from base_repo (the diffusers pipeline) rather than repo_id, which for
a GGUF/single-file SDXL load is the checkpoint path from_pretrained can't open.
- Add client-side validation of steps/rank/resolution/batch/learning-rate before
the request.
- Expose a precision selector (bf16/fp16/fp32) so non-bf16 GPUs can train from
the UI, not only the API.
- Gate the dialog on the active Images route (active && trainOpen) so switching
tabs closes it and stops its polling.
- Rescan the LoRA picker when a run completes, so a freshly-trained adapter
appears without a model reload.
- Cap the dialog height and scroll the body so the Start/Stop footer stays
reachable on short viewports.
- Correct the copy to not over-promise picker auto-discovery.
Freeing the resident Images pipeline before training is handled backend-side in
the diffusion training start route.
* Merge diffusion-sdxl into diffusion-lora-ux; keep options-only LoRA catch
The catalog-refresh .catch from the lower branch clears the selected adapters
too, which is right for its catalog-only picker but wrong here: this picker
holds free-text HF repo ids that are valid without being in the catalog, so a
transient refresh failure must not wipe them. Family swaps still clear the
selection and hidden LoRAs are never sent.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Train LoRA dialog: stop suggesting absolute paths the backend rejects
The dataset and output placeholders showed /path/to/... examples, but the
training routes resolve those fields inside the Studio home and reject
absolute paths outside the approved roots, so following the placeholder
produced a 400. Use folder-name placeholders and say in the labels and the
dialog description where each folder resolves.
* Align the VAE to the denoiser's first FLOATING dtype, not its first parameter
A GGUF-quantized transformer's leading parameters are packed uint8 storage,
so reading next(parameters()).dtype handed nn.Module.to() an integer dtype
and every image-conditioned generation on a GGUF model (Qwen-Image-Edit)
failed with a 500. Probe the parameters for the first floating dtype, treat
an all-integer module as a no-op, and also catch TypeError so an unexpected
dtype can never break generation. Regression test included.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Count LR scheduler warmup/decay in optimizer steps, not micro-steps
lr_sched.step() runs once per outer optimizer step (after the gradient
accumulation inner loop), for train_steps total. The scheduler was
configured with num_warmup_steps and num_training_steps multiplied by
gradient_accumulation_steps, so with accumulation > 1 a warmup or
non-constant schedule stretched past the run and never reached the
intended decay. Count both in optimizer steps.
* Address Codex review findings on the image-workflows PR
Keep diffusion.py importable without torch: the compile/arch patch modules
import torch at module level, so import them lazily at their load/unload
call sites instead of at module load. This restores the torchless contract
so get_diffusion_backend() works on a CPU/native sd.cpp install.
Match family reject keywords and aliases as whole path/name segments, not
raw substrings, so an unrelated word like edited, edition, or kontextual no
longer misroutes or hides a valid base image model, while supported edit
families (Qwen-Image-Edit, FLUX Kontext) still resolve. Mirror the same
segment matching in the picker task filter.
Route FLUX.2-dev native guidance through --guidance like the other FLUX
families rather than --cfg-scale. Reject native upscale requests that have
no input image. Read image header dimensions and reject over-limit inputs
before decoding pixels, so a crafted small-payload image cannot spike
memory. Reject an upscale that would shrink the source below its input
size. Validate the model_kind against the filename extension before the
GPU handoff. Estimate a local diffusers pipeline's size from its on-disk
weights so auto memory planning does not skip offload and OOM. Report
workflows: [txt2img] from the native backend status so the Create tab
stays enabled for a loaded native model. Clamp the outpaint canvas to the
backend's 4096px decode limit.
Adds regression tests for segment matching and kind/extension validation.
* Guard inference loads and worker lifetime against diffusion training
Teach the chat and image load guards about an active diffusion (SDXL) LoRA
job: a chat load is refused (its footprint cannot be fit-checked against the
trainer) and an image load is refused outright, mirroring the existing LLM
training guards, so a load can no longer allocate GPU memory alongside the
trainer and undo the pre-start cleanup.
Bind the diffusion trainer subprocess to the parent's lifetime and scrub the
native path lease secret from it by running the child through
run_without_native_path_secret, matching the inference/export/LLM workers, so
a Studio crash or kill no longer leaves the trainer holding the GPU.
Reset in_model_load on the complete and error terminal events: a stop or
failure during model loading otherwise leaves the status reporting a stale
loading indicator after the job has ended.
* Harden diffusion LoRA handling on the diffusers and native paths
Reject LoRA on a torch.compile'd diffusers transformer (Speed=default/max):
diffusers requires the adapter loaded before compilation, so applying one to
the already-compiled module fails with adapter-key mismatches. The status
gate now hides the picker and generate raises a clear message instead.
Convert a cancelled Hub LoRA download (RuntimeError Cancelled) to the
diffusion cancellation sentinel in resolve_specs, so an unload/superseding
load during resolution maps to a 409 instead of a generic server error.
Drop weight-0 LoRA rows before the native support gate so a request carrying
only disabled adapters stays a no-op on families where native LoRA is
unsupported, matching the diffusers path.
Reject duplicate LoRA ids in the request model: both apply paths suffix
colliding names, so a repeated id would stack the same adapter past its
per-adapter weight bound.
Strip all user-typed <lora:...> prompt tags on the native path (only the
selected adapters are materialized in the managed lora-model-dir, so an
unselected tag can never resolve), and restore saved LoRA selections from a
gallery recipe so restore reproduces a LoRA image.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Harden ControlNet resolve, gallery metadata, and the control-type picker
Check cancellation immediately after a ControlNet from_pretrained and before
any device placement, so an unload/eviction that raced the download does not
allocate several GB onto the GPU after the load was already cleared.
Require a loadable weight or shard index (not just config.json) before a local
ControlNet folder is advertised, so an interrupted copy is hidden instead of
failing deep in from_pretrained as a generic 500.
Do not record a strength-0 ControlNet in the gallery recipe: it is treated as
disabled and skipped, so the image is unconditioned and the metadata must not
claim a ControlNet was applied.
Build the control-type picker from the selected ControlNet's advertised
control_types instead of a hardcoded passthrough/canny pair, so a union model
with a precomputed depth or pose map sends the correct control_mode.
* Address further Codex findings on the image-workflows PR
- Persist the actual output image size in the gallery recipe instead of the
request sliders: Transform/Inpaint/Edit derive the size from the uploaded
image, Extend grows the canvas, and Upscale resizes it, so the sliders
recorded (and later restored) the wrong dimensions for those workflows.
- Reject a remote '*-GGUF' repo loaded as a full pipeline (no single-file
name) in validate_load_request, so the unloadable pick fails before chat is
evicted rather than deep in from_pretrained.
- Only publish an image-conditioned from_pipe wrapper to the shared aux cache
when the load is still current: from_pipe runs under the generate lock but
not the state lock, so an unload racing its construction could otherwise
cache a wrapper over torn-down modules that a later load would reuse.
- Verify the Windows CUDA runtime archive checksum before extracting it, like
the main sd-cli archive, so a corrupt or tampered runtime is rejected rather
than extracted next to the binary.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Refuse non-SDXL base models at diffusion training start
The trainer only supports the SDXL U-Net, but a FLUX / Qwen-Image / Z-Image
repo or a GGUF filename passed as base_model was accepted and then failed
minutes later inside StableDiffusionXLPipeline.from_pretrained with an
unrelated-looking error. Add a name-based guard in normalized() so known
DiT-family names and .gguf checkpoints are rejected up front, which the API
start route surfaces as an immediate 400 with a message that says exactly
which bases are trainable. Unrecognisable names still pass through so custom
local SDXL checkpoints keep working.
* Add diffusion dataset upload and training info endpoints
Training an image LoRA required knowing the Studio home layout and copying
files onto the server by hand, which is the most confusing step of the whole
flow. Two small endpoints fix that:
- GET /api/train/diffusion/info reports the datasets and outputs roots plus
every dataset folder that contains images (with image/caption counts), so
the UI can offer a picker instead of a blind free-text path.
- POST /api/train/diffusion/dataset uploads images and optional caption
.txt / metadata.jsonl files into a named folder under the datasets root,
creating it on first use and accumulating on repeat uploads so large sets
can arrive in batches. Names are validated to a single path component and
files stream to disk under the same per-upload size cap as LLM dataset
uploads. The returned name is a valid data_dir for /diffusion/start.
* Rework the Train LoRA dialog into a guided SDXL flow
The dialog assumed users knew the Studio home layout and that only SDXL is
trainable, and hid both facts behind free-text fields. Restructure it around
the three real decisions:
- Base model is a dropdown of the trainable SDXL picks (Base 1.0, Turbo, the
loaded SDXL pipeline when there is one) with a custom repo/path escape
hatch, instead of a bare text field defaulting to a repo id.
- Training images come from an in-browser upload (new dataset endpoints) or
a picker over existing dataset folders with image/caption counts. No shell
access or knowledge of the datasets root is needed any more, and the
captioning rules are explained inline.
- The output field is now Adapter name and the instance prompt is labelled
as the trigger prompt, with a no-captions warning wired to the selected
dataset's actual caption count.
Hyperparameters collapse behind a training settings toggle since the
defaults suit a first run. A completed run says where the adapter went and
offers Done / Train another, and the top-bar button gets an icon and a
plainer description. The dialog title states the SDXL-only scope and that
other families load LoRAs but cannot train them yet.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Validate diffusion training config before freeing the GPU
The start route freed resident GPU workloads (export, Images pipeline, chat)
before the service validated the config, so a start that was then refused,
now including a non-SDXL base model, tore down the user's loaded model for
nothing. Run the same cheap normalise pass first; the LLM path already
follows this rule via its before_spawn hook.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Refactor diffusion LoRA training into a family-aware platform
Split the SDXL trainer into a shared, architecture-agnostic layer so more model
families can be trained without duplicating the plumbing:
- New core/training/diffusion_train_common.py holds the config + validation, dataset
discovery, event emission, stop protocol, adapter publishing, and a lazy trainer
registry (get_trainer). diffusion_lora_trainer.py keeps the SDXL-specific loop and
re-exports the moved names so existing imports are unchanged.
- The SDXL-only base-model blocklist becomes a positive check: the family is resolved
from the base model (or an explicit model_family) via the diffusion family registry,
and a known-but-not-yet-trainable family is refused with a clear message. Unknown
custom names still default to the SDXL trainer.
- DiffusionFamily gains a trainable flag and train_base_repos; SDXL is marked trainable.
DiT families flip on when their trainers land.
- Trained adapters now write a <name>.json metadata sidecar (family, base model, rank,
trigger prompt, ...) that the LoRA scanner reads to family-gate the adapter in the
picker instead of showing it as unknown for every model.
- The training base-model trust allowlist adds the official FLUX.1-dev, Z-Image-Turbo,
and Qwen-Image repos (safetensors-only, no remote code).
* Retain diffusion training loss history and expose it in status
The training service kept only the latest loss, so a live loss chart could show a
single point. Fold each progress event into bounded (step, loss, lr) history arrays
(capped at 4000 points, decimated when full) plus the latest throughput and peak VRAM,
and record the family / base model / catalog path on completion. The status endpoint
returns these as a nested metric_history object the UI can chart directly, and the
start request accepts an optional model_family override.
* Tests for the diffusion training platform
Cover the trainer registry (get_trainer resolves SDXL, unknown family raises),
family resolution (explicit model_family validation, resolved_family on the config),
the metadata sidecar write + scan read with family gating, and the service loss-history
folding (append, bad-point skipping, decimation at cap, family/perf fields) plus the
status route nesting metric_history.
* [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>
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
* Studio diffusion: cross-platform device policy, fp16 guard, lock split, validate-before-evict
Phase 1 of porting the richer diffusion stack onto the image-generation backend.
- Add a compartmentalized device/dtype policy module (diffusion_device.py)
resolving CUDA/ROCm/XPU/MPS/CPU with capability flags. Keeps the NVIDIA
capability-based bf16 choice; ROCm and XPU are isolated; MPS uses bf16 or
fp32, never a silent fp16 that renders a black image.
- Add a per-family fp16_incompatible flag (Z-Image) and promote a resolved
float16 to float32 for those families so they do not produce black images.
- Split the backend locks: a generation holds only _generate_lock, so status,
unload, and a new load are never blocked by a long denoise. Add per-generation
cancellation via callback_on_step_end so an eviction or a superseding load
preempts a running generation; a replacement load waits for it to stop before
allocating, so two pipelines never sit in VRAM at once.
- Validate a load request before the GPU handoff so an unloadable pick never
evicts a working chat model, and reject missing local paths up front.
- Add CPU-only tests for the device policy, dtype guard, lock split and
cancellation, and validate-before-evict, plus a GPU benchmark/regression
script (scripts/diffusion_bench.py) measuring latency, peak VRAM, and PSNR
against a saved reference.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 2A): measured-budget memory planner + offload/VAE policy
Add a lean, backend-agnostic memory policy that picks a CPU-offload policy and
VAE tiling/slicing from measured free device memory vs the model's estimated
resident footprint, then applies it to the built pipeline. auto stays resident
when the model fits (byte-identical to the prior resident path), and falls to
whole-module offload when tight; fast/balanced/low_vram are explicit overrides.
Sequential submodule offload is unreliable for GGUF transformers on diffusers
0.38, so it falls back to whole-module offload and status reports the policy
actually engaged.
Verified on Z-Image-Turbo Q4_K_M (B200): auto reproduces the resident image with
no VRAM/latency regression (PSNR inf); balanced/low_vram cut generation peak VRAM
47.9% (15951 -> 8318 MB) with byte-identical output, at the expected latency cost.
73 prior + 35 new CPU tests pass.
* Studio diffusion (Phase 2D): streamed block-level offload + functional VAE tiling
Add a streamed 'group' offload tier (diffusers apply_group_offloading, block_level,
use_stream) that keeps the transformer flowing through the GPU a few blocks at a
time while the text encoder / VAE stay resident, and fix VAE tiling to drive the
VAE submodule (pipelines like Z-Image expose enable_tiling on pipe.vae, not the
pipeline). apply_memory_plan now returns the (policy, tiling) actually engaged so
status never overstates either, and group falls back to whole-module offload when
the transformer can't be streamed.
Measured on Z-Image (B200), all lossless (PSNR inf vs resident): balanced/group
cuts generation peak VRAM 32% (15951 -> 10840 MB) at near-resident speed (2.07 ->
2.99s); low_vram/model cuts it 48% (-> 8318 MB) but is slower (7.99s). Mode names
now match that tradeoff: balanced = stream the transformer, low_vram = offload
every component. auto picks group when the companions fit resident, else model.
112 CPU tests pass.
* Studio diffusion (Phase 5): image quality-vs-quant accuracy harness
Add scripts/diffusion_quality.py, the accuracy analogue of the KLD workflow: hold
prompt + seed fixed, render a grid with a reference quant (default BF16), then render
each candidate quant and measure drift from the reference. Records mean PSNR + SSIM
(pure-numpy, no skimage/scipy) and optional CLIP text-alignment + image-similarity
(transformers, --clip), plus file size, latency, and peak VRAM, then prints a
quality-vs-cost table and recommends the smallest quant within a quality budget.
--selftest validates the metrics on synthetic images with no GPU or model.
Verified on Z-Image (B200): the table degrades monotonically with quant size
(Q8 -> Q4 -> Q2: PSNR 21.7 -> 15.5, SSIM 0.82 -> 0.61), while CLIP-text stays flat
(~0.34) -- quantization erodes fine detail far more than prompt adherence.
* Studio diffusion (Phase 3): opt-in speed layer (channels_last / compile / TF32)
Add a speed_mode knob (off by default, so the render path stays bit-identical):
default applies channels_last VAE + regional torch.compile of the denoiser's
repeated block where eligible; max also enables TF32 matmul and fused QKV. Regional
compile is gated off for the GGUF transformer (dequantises per-op) and for families
flagged not compile-friendly (a new supports_torch_compile flag, False for Z-Image),
so it activates automatically only once a non-GGUF bf16 transformer is loaded. Speed
optims run before placement/offload, per the diffusers composition order. status now
reports speed_mode + the optims actually engaged.
Verified on Z-Image (B200): default -> ['channels_last'], max -> ['channels_last',
'tf32'], compile correctly skipped for GGUF; generation works in every mode.
121 CPU tests pass.
* Studio diffusion (Phase 2B): opt-in fp8 text-encoder layerwise casting
Add a text_encoder_fp8 knob that casts the companion text encoder(s) to fp8 (e4m3)
storage via diffusers apply_layerwise_casting, upcasting per layer to the bf16
compute dtype while normalisations and embeddings stay full precision. Applied
before placement, gated to CUDA + bf16, best-effort (a failure leaves the encoder
dense). status reports which encoders were cast.
Verified on Z-Image (B200, balanced/group mode where the encoder stays resident):
generation peak VRAM dropped 37% (10840 -> 6791 MB, below the lowest-VRAM offload)
at near-resident speed. It is a memory-vs-quality tradeoff, not free -- ~20 dB PSNR
vs the bf16 encoder, a larger shift than one transformer quant step -- so it is off
by default and documented as such, with the Phase 5 harness to size the cost.
127 CPU tests pass.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 2C): NVFP4 text-encoder quant (+ generalise fp8 knob)
Generalise the text-encoder precision knob from a fp8 bool to text_encoder_quant
(fp8 | nvfp4). nvfp4 quantises the companion text encoder to 4-bit via torchao
NVFP4 weight-only (two-level microscaling) on Blackwell's FP4 tensor cores; fp8
stays the broader-hardware path (cc>=8.9). Both are gated, best-effort, and run
before placement; status reports the mode actually engaged. This is the lean
realisation of GGUF-native text-encoder quant: 4-bit on the encoder without the
3045-line port.
Verified on Z-Image (B200, balanced/group where the encoder stays resident), vs the
bf16 encoder: nvfp4 cut generation peak VRAM 48% (10840 -> 5593 MB, the lowest TE
option, below whole-model offload) at near-fp8 quality (16.4 vs 17.1 dB PSNR), and
both quants ran faster than bf16. A memory-vs-quality tradeoff (off by default);
size it per model with the Phase 5 quality harness. diffusion_bench gains
--text-encoder-quant.
129 CPU tests pass.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 4): native stable-diffusion.cpp engine for CPU/Mac
Adds the CPU / Apple-Silicon tier of the two-engine strategy, mirroring the
chat backend's llama.cpp shell-out. Diffusers stays the default on CUDA / ROCm
/ XPU; this covers the hardware diffusers serves poorly, consuming the same
split GGUF assets Studio already curates.
- sd_cpp_args.py: pure sd-cli command builder. Maps the family to its
text-encoder flag (Z-Image Qwen3 to --llm, Qwen-Image to --qwen2vl, FLUX.1
CLIP-L + T5), and the diffusers memory policy (none/group/model/sequential)
to sd.cpp's offload flags (--offload-to-cpu / --clip-on-cpu / --vae-on-cpu /
--vae-tiling / --diffusion-fa), so one user knob drives both engines.
- sd_cpp_engine.py: SdCppEngine over a located sd-cli. find_sd_cpp_binary()
with the same precedence as the llama finder (env override, then the Studio
install root, then in-tree, then PATH), an is_available/version probe, and a
one-shot subprocess generate that streams progress and returns the PNG.
runtime_env() prepends the binary's directory to the platform library path
so a prebuilt's bundled libstable-diffusion.so resolves.
select_diffusion_engine() is the pure routing decision (GPU backends to
diffusers, CPU/MPS to native when present).
- install_sd_cpp_prebuilt.py: resolve + download the per-host prebuilt
(macOS-arm64/Metal, Linux x86_64 CPU, Vulkan/ROCm/Windows variants) into the
Studio install root. resolve_release_asset() is a pure, unit-tested
host-to-asset matrix.
- scripts/sd_cpp_smoke.py: end-to-end native generation harness.
Tests (CPU-only, subprocess/filesystem stubbed): 49 new across args, engine,
routing, runtime env, and the installer resolver. Full diffusion suite 166
passing.
Verified on a B200 box: built sd-cli (CUDA) and the prebuilt (CPU) both
generate Z-Image-Turbo Q4_K end to end through SdCppEngine: balanced (group
offload, 5.0s gen), low_vram (full CPU offload + VAE tiling, 13.4s), and the
dynamically-linked CPU prebuilt (50.4s on CPU), all producing coherent images.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 6): img2img / inpaint / edit / LoRA / upscale on the native engine
Builds on Phase 4's native stable-diffusion.cpp engine, extending it from
text-to-image to the wider feature surface, since sd.cpp supports all of these
through the binary already. Pure command-builder additions plus one engine
method, so the txt2img path is unchanged.
- sd_cpp_args.py: SdCppGenParams gains image-conditioning fields. init_img +
strength make a run img2img, adding mask makes it inpaint, ref_images drives
FLUX-Kontext / Qwen-Image-Edit style editing (repeated --ref-image), and
lora_dir + the <lora:name:weight> prompt syntax select LoRAs. New
SdCppUpscaleParams + build_sd_cpp_upscale_command for the ESRGAN upscale run
mode (input image + esrgan model, no prompt / text encoders).
- sd_cpp_engine.py: the subprocess runner is factored into a shared _run() so
generate() (now carrying the conditioning flags) and a new upscale() reuse
the same streaming / error / output-check path.
- scripts/sd_cpp_smoke.py: --task {txt2img,img2img,upscale} with --init-img /
--strength / --upscale-model / --upscale-repeats.
Tests: 10 new across the img2img / inpaint / edit / LoRA flag construction, the
upscale builder and its validation, and the engine's img2img + upscale paths.
Full diffusion suite 176 passing.
Verified on a B200 box through SdCppEngine: img2img (Z-Image-Turbo Q4_K, the
init image conditioned at strength 0.6, 4.8s) and ESRGAN upscale
(512x512 -> 2048x2048 via RealESRGAN_x4plus_anime_6B, 2.7s), both producing
coherent images. Video and the diffusers-path feature wiring are deferred.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 7): accuracy-preserving speed pass
Re-review of the diffusion stack (#6675/#6679/#6680) surfaced one real accuracy
bug and a dead-on-arrival speed path; this fixes both and adds the lossless /
near-lossless wins, all measured on a B200.
Correctness:
- TF32 global-state leak (fix). speed_mode=max flipped torch.backends.*.allow_tf32
process-wide and never restored them, so a later `off` load silently inherited
TF32 and was no longer bit-identical. Added snapshot_backend_flags /
restore_backend_flags (TF32 + cudnn.benchmark), captured before the speed layer
runs and restored on unload. Verified: load max -> unload -> load off is now
byte-identical (PSNR inf) to a fresh off.
- sd-cli timeout could hang forever. _run() blocked in `for line in stdout` and
only checked the timeout after EOF, so a child stuck in model load / GPU init
with no output ignored the timeout. Drained stdout on a reader thread with a
wall-clock deadline. Added a silent-hang regression test.
Speed (diffusers path), near-lossless, opt-in tiers:
- Regional torch.compile now runs on the GGUF transformer. The is_gguf gate (and
Z-Image's supports_torch_compile=False) were stale: compile_repeated_blocks
compiles and runs ~2.2x faster on the GGUF Z-Image transformer on
torch 2.9.1 / diffusers 0.38 (the per-op dequant stays eager, the rest of the
block compiles). Measured: off 1.80s -> default 0.82s/gen (+54.7%), PSNR 37.7 dB
vs eager -- far above the Q4 quant noise floor (~21 dB), so it does not move
output quality. Gate relaxed; default tier delivers it.
- cudnn.benchmark added to the default tier (autotunes the fixed-shape VAE convs).
- torch.inference_mode() around the pipeline call (lossless, strictly faster than
the no_grad diffusers uses internally).
Memory path:
- VAE tiling (not bit-identical >1MP) restricted to the model/sequential/CPU tiers;
the balanced (group) tier keeps exact slicing only, so it is now bit-identical to
the resident image (verified PSNR inf) and slightly faster.
- Group offload adds non_blocking + record_stream on the CUDA stream path to
overlap each block's H2D copy with compute (lossless; gated on the installed
diffusers signature so older versions still work).
Native (sd.cpp) path:
- native_speed_flags: a first-class speed knob (default -> --diffusion-fa, a
near-lossless CUDA win that was previously only added on offload tiers; max also
-> --diffusion-conv-direct). conv-direct stays opt-in: measured +45% on CUDA, so
it is never auto-on. Engine generate() merges it, de-duped against offload flags.
Default profile: a GGUF model with no explicit speed_mode now resolves to the
`default` profile (resolve_speed_mode), since compile's perturbation sits below the
quantisation noise floor and so does not reduce quality versus the dense reference;
out of the box a GGUF Z-Image generation drops from 1.80s to 0.81s. Dense models
stay `off` / bit-identical, and an explicit speed_mode -- including "off" -- is
always honored, so the byte-identical path remains one flag away and is the
regression reference.
Tooling: scripts/compile_probe.py (eager vs compiled GGUF probe), scripts/
perf_verify.py (the B200 verification above), and diffusion_bench.py gains
--speed-mode so the speed tiers are benchmarkable.
Tests: 183 passing (was 166); new coverage for the backend-flag snapshot/restore,
GGUF compile eligibility, the balanced tiling/slicing split, native_speed_flags +
the engine de-dup, and the sd-cli silent-hang timeout.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 7): max tier uses max-autotune-no-cudagraphs + engine/lever benchmarks
The opt-in `max` speed tier now compiles the repeated block with
mode=max-autotune-no-cudagraphs (dynamic=False) instead of the default mode:
Triton autotuning for GEMM/conv-heavier models, gated to the tier where a longer
cold compile is acceptable. CUDA-graph modes (reduce-overhead / max-autotune) are
deliberately avoided -- both crash on the regionally-compiled block (its static
output buffer is overwritten across denoise steps), measured.
Adds two reproducible benchmarks used to validate the optimization research:
- scripts/compare_engines.py: PyTorch (diffusers GGUF) vs native sd.cpp head-to-head.
- scripts/leverage_probe.py: coordinate_descent_tuning + FirstBlockCache probes.
Measured on B200 (Z-Image Q4_K_M, 1024px, 8 steps): default compile 0.80s/gen;
coordinate_descent_tuning 0.79s (within noise, already covered by max-autotune);
FirstBlockCache does not run on Z-Image (diffusers 0.38 block-detection / Dynamo).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 8): opt-in fast transformer (torchao int8/fp8/fp4 on a dense source)
Add an opt-in transformer_quant mode that loads the dense bf16 transformer and
torchao-quantises it onto the low-precision tensor cores, instead of the GGUF
transformer (which dequantises to bf16 per matmul and so runs at bf16 rate). On a
B200 (Z-Image-Turbo, 1024px/8 steps): auto picks fp8 at 0.614s vs GGUF+compile's
0.823s (1.34x), int8 0.626s (1.32x), both at lower LPIPS than GGUF's own 4-bit floor.
GGUF+compile stays the low-memory default and the fallback. The mode is gated on
CUDA + bf16 + resident VRAM headroom (the dense load peaks ~21GB vs GGUF's 13GB);
any unsupported arch/scheme, OOM, or quant failure falls back to GGUF with a logged
reason. auto picks the best scheme per GPU via a real quantise+matmul smoke probe
(Blackwell nvfp4/fp8/mxfp8, Ada/Hopper fp8, Ampere int8); a min-features filter skips
the tiny projections that crash int8's torch._int_mm. New module mirrors
diffusion_precision.py; quant runs before compile before placement.
184 -> tests pass; new test_diffusion_transformer_quant.py plus backend/route
coverage. scripts/diffusion_bench.py gains --transformer-quant; scripts/quant_probe.py
is the standalone torchao lever probe.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 8): consumer-GPU tuning - lock fp8 fast accumulate, prefer fp8 over mxfp8, reject 2:4 sparsity
Consumer Blackwell halves tensor-core throughput on FP32 accumulate (fp8 419 vs 838
TFLOPS with FP16 accumulate; bf16 209), so:
- fp8 config locks use_fast_accum=True (Float8MMConfig). torchao already defaults it on;
pinning it guards consumer cards against a default change. On B200 it is identical
speed and slightly better quality (LPIPS 0.050 vs 0.091).
- the Blackwell auto ladder prefers fp8 over mxfp8 (measured faster + more accurate).
2:4 semi-structured sparsity evaluated and rejected (scripts/sparse_accum_probe.py):
2:4 magnitude-prune + fp8 gives LPIPS 0.858 (broken image) with no fine-tune, the
cuSPARSELt kernel errors on torch 2.9, and it does not compose with torch.compile
(our main ~2x). Documented as a dead end, not shipped.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 8): add fp8 fast-accum overflow verification probe
scripts/fp8_overflow_check.py hooks every quantised linear during a real Z-Image
generation and reports max-abs + non-finite counts for use_fast_accum True vs False.
Confirms fast accumulation is an accumulation-precision knob, not an overflow one:
across 276 linears, including Z-Image's ~1.0e6 activation peaks (which overflow FP16),
0 non-finite elements and identical max-abs for both modes.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 8): detect consumer vs data-center GPU for fp8 accumulate, with user override
Consumer/workstation GPUs (GDDR) halve fp8 FP32-accumulate throughput, so they want
fast (FP16) accumulate; data-center HBM parts (B200/H100/A100/L40) are not nerfed and
prefer the higher-precision FP32 accumulate. Add _is_consumer_gpu() (token-exact match
on the device name per NVIDIA's GPU list, so workstation A4000 != data-center A40;
GeForce/TITAN and unknown default to consumer) and gate the fp8 use_fast_accum on it.
Measured: fast accumulate is ~2x on consumer Blackwell and ~8% on B200 (0.608 vs 0.665s),
no overflow, quality below the quant noise floor. So the default leans to accuracy on
data-center; a new request field transformer_quant_fast_accum (null=auto, true/false=force)
lets the operator override per load (scripts/diffusion_bench.py --fp8-fast-accum auto|on|off).
187 diffusion tests pass (+ consumer detection, _resolve_fast_accum, and the override
threading).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 8): add NVFP4 probe documenting it is not yet a win on torch 2.9
scripts/nvfp4_probe.py measures NVFP4 via torchao on the real Z-Image transformer.
Finding (B200, 1024px/8 steps): NVFP4 is a torchao feature and DOES run with
use_triton_kernel=False (the default triton path needs the missing MSLK library), but
only at bf16-compile rate (0.667s vs fp8 0.592s) -- it dequantises FP4->bf16 rather than
using the FP4 tensor cores. The real FP4 speedup needs MSLK or torch>=2.11 + torchao's
CUTLASS FP4 GEMM. The smoke probe (default triton=True) already keeps NVFP4 out of auto
on this env, so auto correctly stays on fp8; NVFP4 activates automatically once fast.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 8): prefer fp8 over nvfp4 in Blackwell auto ladder
Validated NVFP4 on torch 2.11 + torchao CUTLASS FP4 in an isolated env. The FP4
tensor-core GEMM is genuinely active there (a 16384^3 GEMM hits ~3826 TFLOPS,
2.52x bf16 and 1.37x fp8), but it only beats fp8 on very large GEMMs. At the
diffusion transformer's shapes (hidden ~3072, MLP ~12288, M~4096) NVFP4 is both
slower (0.81x fp8 end to end on Z-Image 1024px) and less accurate (LPIPS 0.166
vs fp8's 0.044). Reorder the Blackwell auto ladder to fp8 before nvfp4 so auto is
correct even on a future MSLK-equipped box; nvfp4 stays an explicit opt-in. Add
scripts/nvfp4_t211_probe.py (extension diagnostics + GEMM micro + end-to-end).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 9): pre-quantized transformer loading
The Phase 8 fast transformer_quant path materialises the dense bf16 transformer on
the GPU and torchao-quantises it in place, so its load peak is ~2x GGUF's (~21 vs
13.4 GB) plus a ~12 GB download. Add a pre-quantized branch: quantise once offline
(scripts/build_prequant_checkpoint.py) and at runtime build the transformer skeleton
on the meta device (accelerate.init_empty_weights) and load_state_dict(assign=True)
the quantized weights, so the dense bf16 never touches the GPU.
Measured (B200, Z-Image fp8): full-pipeline GPU load peak 21.2 -> 14.6 GB (matching
GGUF's 13.4), on-disk 12 -> 6.28 GB, output bit-identical (LPIPS 0.0). It is the same
torchao config + min_features filter the runtime path uses, applied ahead of time.
New core/inference/diffusion_prequant.py (resolve_prequant_source +
load_prequantized_transformer, best-effort, lazy imports). diffusion.py
_load_dense_quant_pipeline tries the pre-quant source first and falls back to the
dense materialise+quantise path, then to GGUF, so the default is unchanged.
DiffusionLoadRequest gains transformer_prequant_path; DiffusionFamily gains an empty
prequant_repos map for hosted checkpoints (hosting deferred). Hermetic CPU tests for
the resolver, the meta-init+assign loader, and the backend branch selection +
fallbacks; GPU verification via scripts/verify_prequant_backend.py.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 10): attention-backend selection
Add a selectable attention kernel via the diffusers set_attention_backend
dispatcher. Attention is memory-bandwidth bound, so a better kernel is an
end-to-end win orthogonal to the linear-weight quantisation (it speeds the QK/PV
matmuls torchao never touches) and composes with torch.compile.
auto picks the best exact backend for the device: cuDNN fused attention
(_native_cudnn) on NVIDIA when a speed profile is active, measured ~1.18x
end-to-end on a B200 (Z-Image 1024px/8 steps) with LPIPS ~0.004 vs the default
(below the compile/quant noise floor); native SDPA elsewhere and when speed=off
(so off stays bit-identical). Explicit native/cudnn/flash/flash3/flash4/sage/
xformers/aiter are honored, and an unavailable kernel falls back to the default
rather than failing the load.
New core/inference/diffusion_attention.py (normalize + per-device select + apply,
best-effort, lazy imports). Set on pipe.transformer BEFORE compile in load_pipeline;
attention_backend threads through begin_load / load_pipeline / status like the other
load knobs. New request field attention_backend + status field. Hermetic CPU tests
for normalize / select policy / apply fallback, plus route threading + 422. Measured
via scripts/perf_levers_probe.py.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 11): prefer int8 on consumer GPUs in the auto ladder
Consumer / workstation GPUs halve fp8 (and fp16/bf16) FP32-accumulate tensor-core
throughput, while int8 runs at full rate (int32 accumulate is not nerfed). Public
benchmarks (SDNQ across RTX 3090/4090/5090, AMD, Intel) confirm int8 via torch._int_mm
is as fast or faster than fp8 on every consumer part, and the only path on pre-Ada
consumer cards without fp8 tensor cores. So when transformer_quant=auto, reorder the
arch tier to put int8 first on a consumer/workstation GPU (detected by the existing
_is_consumer_gpu name heuristic), while data-center HBM parts keep fp8 first.
Pure ladder reorder via _prefer_consumer_scheme; no new flags. Verified non-regression
on a B200 (still picks fp8). Hermetic tests for consumer Blackwell/Ada/workstation
(-> int8) and data-center Ada/Hopper/Blackwell (-> fp8).
* Studio diffusion (Phase 12): First-Block-Cache step caching for many-step DiT
Add opt-in step caching (First-Block-Cache) for the diffusion transformer. Across
denoise steps a DiT's output settles, so once the first block's residual barely
changes the remaining blocks are skipped and their cached output reused. diffusers
ships it natively (FirstBlockCacheConfig + transformer.enable_cache, with the
standalone apply_first_block_cache hook as a fallback).
Measured on Flux.1-dev (28 steps, 1024px): ~1.4x on top of torch.compile (2.83 ->
2.03s) at LPIPS ~0.08 vs the no-cache output, well inside the quality bar.
OFF by default and a per-load opt-in: the win scales with step count, so it is for
many-step models (Flux / Qwen-Image) and pointless for few-step distilled models
(e.g. Z-Image-Turbo at ~8 steps), where a single skipped step is a large fraction
of the trajectory. It composes with regional compile only with fullgraph=False (the
cache's per-step decision is a torch.compiler.disable graph break), which the speed
layer now switches to automatically when a cache is engaged. Best-effort: a model
whose block signature the hook does not recognise is caught and the load proceeds
uncached.
- new core/inference/diffusion_cache.py: normalize_transformer_cache + apply_step_cache
(enable_cache / apply_first_block_cache fallback; threshold auto-raised for a
quantised transformer per ParaAttention's fp8 guidance; lazy diffusers import).
- diffusion_speed.py: apply_speed_optims takes cache_active; compile drops fullgraph
when a cache is engaged.
- diffusion.py: apply_step_cache before compile; thread transformer_cache /
transformer_cache_threshold through begin_load -> load_pipeline and report the
engaged mode in status().
- models/inference.py + routes/inference.py: transformer_cache (off | fbcache) and
transformer_cache_threshold request fields, engaged mode in the status response.
- hermetic tests for normalisation, the enable_cache / hook-fallback paths, threshold
selection, and best-effort failure handling, plus route threading + validation.
- scripts/fbcache_flux_probe.py: the Flux validation probe (latency / speedup / VRAM /
LPIPS vs the compiled no-cache baseline).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 14): fix int8 dense quant on Flux / Qwen (skip M=1 modulation linears)
The opt-in dense int8 transformer path crashed on Flux.1 and Qwen-Image with
'torch._int_mm: self.size(0) needs to be greater than 16, but got 1'. int8 dynamic quant
goes through torch._int_mm, which requires the activation row count M > 16. A DiT's AdaLN
modulation projections (Flux norm1.linear 3072->18432, Qwen img_mod.1 / txt_mod.1, Flux.2
*_modulation.linear) and its timestep / guidance / pooled-text conditioning embedders are
computed once from the [batch, dim] conditioning vector (M = batch = 1), not per token, so
they hit _int_mm at M=1 and crash. Their feature dims are large, so the existing
min_features filter did not exclude them.
Fix: the int8 filter now also skips any Linear whose fully-qualified name matches a
modulation / conditioning-embedder token (norm, _mod, modulation, timestep_embed,
guidance_embed, time_text_embed, pooled). These layers run at M=1 once per block and are a
negligible share of the FLOPs, so int8 keeps the full speedup on the attention / FFN layers
(M = sequence length). fp8 / nvfp4 / mxfp8 use scaled_mm, which has no M>16 limit and
quantises these layers fine, so the exclusion is int8-only. Sequence embedders
(context_embedder / x_embedder / txt_in, M = seq) are deliberately not excluded -- note
'context_embedder' contains the substring 'text_embed', which is why the token is the
specific 'time_text_embed', not 'text_embed'.
Measured on a B200 (1024px, transformer_quant=int8 + speed=default), int8 now runs on every
supported model and is the fastest dense path on Flux/Qwen (int8 runs full-rate vs fp8's
FP32-accumulate): FLUX.1-dev 9.62s eager -> 1.98s (4.86x, vs fp8 2.15s), Qwen-Image -> 1.87s
(5.57x, vs fp8 2.09s), FLUX.1-schnell -> 0.41s (3.59x). Z-Image and Flux.2-klein (already
working) are unchanged.
- diffusion_transformer_quant.py: add _INT8_EXCLUDE_NAME_TOKENS; make_filter_fn takes
exclude_name_tokens; quantize_transformer passes it for int8 only.
- hermetic test that the int8 filter excludes the modulation / embedder linears (and keeps
attention / FFN / sequence-embedder linears), while fp8 keeps them.
- scripts/int8_linear_probe.py: the meta-device probe used to enumerate each transformer's
Linear layers and derive the exclusion list.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 15): build int8 pre-quantized checkpoints (skip M=1 modulation linears)
The prequant-checkpoint builder applied the dense quant filter without the int8-only
M=1 modulation / conditioning-embedder exclusion the runtime path uses, so a built int8
checkpoint baked those projections as int8 and crashed (torch._int_mm needs M>16) at the
first denoise step on Flux / Qwen. Factor the scheme->exclusion decision into a shared
exclude_tokens_for_scheme() used by both the runtime quantise path and the offline builder
so they can never drift, and apply it in build_prequant_checkpoint.py. int8 prequant now
produces a working checkpoint on every supported model, giving int8 (the consumer-preferred
scheme) the same ~2x load-VRAM and download reduction fp8 already had.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 16): route no-GPU loads to the native sd.cpp engine
When no CUDA/ROCm/XPU GPU is available, route diffusion load/generate to the
native stable-diffusion.cpp engine instead of diffusers, with diffusers as the
guaranteed fallback. On CPU sd.cpp is 1.4-2.8x faster and uses 1.5-2.2x less RAM.
- diffusion_engine_router: centralised engine selection (built on the existing
select_diffusion_engine), env opt-outs, MPS gating, recorded fallback reason.
- sd_cpp_backend (SdCppDiffusionBackend): the diffusers backend method surface
backed by sd-cli, with lazy binary install, registry-driven asset fetch,
step-progress parsing, and cancellation.
- diffusion_families: per-family single-file VAE + text-encoder asset mapping.
- sd_cpp_engine: cancellation support (process-group kill + SdCppCancelled).
- routes/inference + gpu_arbiter: drive the active engine via the router; the
API now reports the active engine and any fallback reason.
- tests for the backend, router, route selection, and cancellation.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Phase 16 review fixes: engine-switch unload, sd.cpp error mapping, per-image seeds, Qwen sampler
Address review feedback on #6724:
- engine router: unload the engine being deactivated on a switch, so the old
model is not left resident-but-unreachable (the evictor only targets the active
engine).
- generate route: sd.cpp execution errors (nonzero exit / timeout / missing
output) now map to 500, not 409 (which only means not-loaded / cancelled).
- native batch: return per-image seeds and persist the actual seed for each image
so every batch image is reproducible.
- Qwen-Image native path: apply --sampling-method euler --flow-shift 3 per the
stable-diffusion.cpp docs; other families keep sd-cli defaults.
- honor speed_mode (native --diffusion-fa) and, off-CPU, memory_mode/cpu_offload
offload flags on the native load instead of hardcoding them off.
- fail the load when the sd-cli binary is present but not runnable (version()
now returns None on exec error / nonzero exit).
- size estimate: only treat the transformer asset as a possible local path.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 9): gate request-supplied local prequant paths behind operator opt-in
load_prequantized_transformer ends in torch.load(weights_only=False), which executes
arbitrary code from the pickle. The transformer_prequant_path load-request field reached
that unpickle for any local file an authenticated caller named, so a request could trigger
remote code execution. Refuse the source.kind=='path' branch unless the operator sets
UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1; the first-party hosted-repo checkpoint stays trusted
and unaffected. Document the requirement on the API field and add gate tests.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 10): reset the global attention backend on native, gate arch-specific kernels, accept sdpa
- apply_attention_backend now restores the native default when no backend is requested or a
kernel fails. diffusers keeps a process-wide active attention backend that
set_attention_backend updates, and a fresh transformer's processors follow it, so a load
that wanted native could silently inherit a backend (e.g. cuDNN) an earlier speed-profile
load pinned, breaking the bit-identical/off guarantee.
- select_attention_backend drops flash3/flash4 up front when the CUDA capability is below
Hopper/Blackwell. diffusers only checks the kernels package at set time, so an explicit
request on the wrong card set fine then crashed mid-generation; it now falls back to native.
- Add the sdpa alias to the attention_backend Literal so an API request with sdpa (already a
valid alias of native) is accepted instead of 422-rejected by Pydantic.
- Drop the dead replace('-','_') normalization (no alias uses dashes/underscores).
- perf_levers_probe.py output dir is now relative to the script, not a hardcoded path.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 12): only engage FBCache on context-aware transformers; quantized threshold for GGUF
- apply_step_cache now engages only via the transformer's native enable_cache (the diffusers
CacheMixin path), which exists exactly when the pipeline wraps the transformer call in a
cache_context. The standalone apply_first_block_cache fallback installed on non-CacheMixin
transformers too (e.g. Z-Image), whose pipeline opens no cache_context, so the load reported
transformer_cache=fbcache and then the first generation crashed inside the hook. Such a model
now runs uncached per the best-effort contract.
- GGUF transformers are quantized (the default Studio load path), so they now use the higher
quantized FBCache threshold when the caller leaves it unset, instead of the dense default
that could keep the cache from triggering.
- fbcache_flux_probe.py: compile cached runs with fullgraph=False (FBCache is a graph break, so
fullgraph=True failed warmup and silently measured an eager cached run); output dir is now
relative to the script, not a hardcoded path.
* Studio diffusion (Phase 11): keep professional RTX cards on the fp8 ladder
_is_consumer_gpu treated professional parts (RTX PRO 6000 Blackwell, RTX 6000 Ada) as
consumer because their names carry no datacenter token, so the auto ladder moved int8 ahead
of fp8 and the fp8 path chose fast accumulate for them. The rest of the backend already
classifies these as datacenter/professional (llama_cpp.py _DATACENTER_GPU_RE), so detect the
same RTX PRO 6000 / RTX 6000 Ada markers here and keep fp8 first with precise accumulate.
Also fix the consumer-Blackwell test to use compute capability (10, 0) instead of (12, 0).
* Studio diffusion (Phase 8): tolerate missing torch.float8_e4m3fn in the mxfp8 config
Accessing torch.float8_e4m3fn raises AttributeError on a torch build without it (not just
TypeError on older torchao), which would break the mxfp8 config helper instead of falling
back to the default. Catch both so the fallback is robust.
quant_probe.py: same AttributeError fallback; run LPIPS on CPU so the scorer never holds
CUDA memory during the per-row VRAM probe; output dir relative to the script.
* Studio diffusion (Phase 7): robust backend-flag snapshot/restore and restore on failed speeded load
- snapshot_backend_flags reads each flag defensively (getattr + hasattr), so a build/platform
missing one (no cuda.matmul on CPU/MPS) still captures the rest instead of skipping the
whole snapshot. restore_backend_flags restores each flag independently so one failure can't
leave the others leaked process-wide.
- load_pipeline restores the flags (and clears the GPU cache) when the build fails after
apply_speed_optims mutated the process-wide flags but before _state captured them for unload
to restore -- otherwise a failed default/max load left cudnn.benchmark/TF32 on and
contaminated later off generations.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 4): enforce the sd-cli timeout while reading output
Iterating proc.stdout directly blocks until the stream closes, so a sd-cli that hangs
without producing output (or without closing stdout) would never reach proc.wait and the
wall-clock timeout was silently bypassed. Drain stdout on a daemon thread and wait on the
PROCESS, so the main thread always enforces the timeout and kills a hung process (which
closes the pipe and ends the reader). Add a test that times out even when stdout blocks,
and make the no-binary test hermetic so a host-installed sd-cli can't leak in.
* Studio diffusion (Phase 14): guard the int8 exclusion filter against a None fqn
The filter callback can be invoked without a module name, so fqn.lower() would raise
AttributeError on None. Fall back to an empty name (nothing matches the exclusion tokens,
so the linear is kept) instead of crashing the quantise pass.
* Studio diffusion (Phase 16) review fixes: native engine robustness
- sd_cpp_backend: stop truncating explicit seeds to 53 bits (mask to int64);
a large requested seed was silently collapsed (2**53 -> 0) and distinct seeds
aliased to the same image. Random seeds stay 53-bit (JS-safe).
- sd_cpp_backend: sanitize empty/whitespace hf_token to None so HfApi/hf_hub
fall back to anonymous instead of failing auth on a blank token.
- sd_cpp_backend: a superseding load now cancels the in-flight generation, so the
old sd-cli can no longer return/persist an image from the previous model.
- diffusion_engine_router: run the previous engine's unload() OUTSIDE the lock so a
slow 10+ GB free / CUDA sync does not block engine selection.
- diffusion_engine_router: probe sd-cli runnability (version()) before committing to
native, so a present-but-unrunnable binary falls back to diffusers at selection.
- diffusion_device: resolve a torch-free CPU target when torch is unavailable, so a
CPU-only install can still reach the native sd.cpp engine instead of failing load.
- tests updated for the runnability probe + a not-runnable fallback case.
* Studio diffusion (Phase 9) review fixes: prequant safety + validation
- SECURITY: a request-supplied local pre-quant path is now unpickled only when it
resolves inside an operator-configured ALLOWLIST of directories
(UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH = dir[:dir...]). The previous boolean opt-in,
once enabled for one trusted checkpoint, allowed torch.load(weights_only=False) on
any path a load request named (arbitrary code execution). realpath() blocks symlink
escapes; a bare on/off toggle is no longer a wildcard.
- Validate the checkpoint's min_features against the runtime Linear filter, so a
checkpoint that quantised a different layer set is rejected instead of silently
loading a model that mismatches the dense path while reporting the same scheme.
- Tolerant base_model_id compare (exact or same final path/repo segment), so a local
path or fork of the canonical base is accepted instead of falling back to dense.
- _has_meta_tensors uses any(chain(...)) (no intermediate lists).
- prequant verify/probe scripts use repo-relative paths (+ env overrides), not the
author's absolute /mnt paths.
- tests: allowlist-dir opt-in, outside-allowlist refusal, min_features mismatch, fork tail.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 7) review fixes: offload fallback + bench scripts
- diffusion_memory: when group offload is unavailable and the plan falls back to
whole-module offload, enable VAE tiling (the group plan left it off, but the fallback
is the low-VRAM path where the decode spike can OOM). Covers both the group and
sequential fallback branches.
- perf_verify: include the balanced-vs-off PSNR in the pass/fail condition, so a
balanced bit-identity regression actually fails the check instead of exiting 0.
- compare_engines: --vae/--llm default to None (were author-absolute /mnt paths), and
the load-progress poll has a 30 min deadline instead of looping forever on a hang.
- test for the group->model fallback enabling VAE tiling.
* Studio diffusion (Phase 8) review fixes: quant compile + nvfp4 path
- diffusion: a torchao-quantized transformer is committed only compiled. A dense model
resolves to speed_mode=off, which would run the quant eager (~30x slower than the GGUF
it replaced), so when transformer_quant engaged and speed resolved to off, promote to
default (regional compile); warn loudly if compile still does not engage.
- diffusion_transformer_quant: build the nvfp4 config with use_triton_kernel=False so the
CUTLASS FP4 path is used (torchao defaults to the Triton kernel, which needs MSLK);
otherwise the smoke probe fails on CUTLASS-only Blackwell and silently drops to GGUF.
- nvfp4_probe: repo-relative output dir + --out-dir (was an author-absolute /mnt path).
- test asserts the eager-quant -> default-compile promotion.
* Studio diffusion (Phase 10) review fixes: attention gating + probe isolation
- diffusion_attention: gate the auto cuDNN-attention upgrade on SM80+; on pre-Ampere
NVIDIA (T4/V100) cuDNN fused SDPA is accepted at set time but fails at first generation,
so auto now stays on native SDPA there.
- diffusion_attention: _active_attention_backend handles get_active_backend() returning an
enum/None (not a tuple); the old unpack always raised and was swallowed, so
the native-restore short-circuit never fired.
- perf_levers_probe: free the resident pipe on a skipped (attn/fbcache) variant; run LPIPS
on CPU so it isn't charged to every variant's peak VRAM; reset force_fuse_int_mm_with_mul
so the inductor_flags variant doesn't leak into later compiled rows.
- tests for the SM80 cuDNN gate.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 4) review fixes: sd.cpp installer + engine hardening
- install_sd_cpp_prebuilt: download the release archive with urlopen + an explicit
timeout + copyfileobj (urlretrieve has no timeout and hangs on a stalled socket);
extract through a per-member containment check (Zip-Slip guard); expanduser the
--install-dir so a tilde path is not taken literally; and on Windows CUDA also fetch
the separately-published cudart runtime DLL archive so sd-cli.exe can start.
- sd_cpp_engine: find_sd_cpp_binary honors UNSLOTH_STUDIO_HOME / STUDIO_HOME like the
installer, so a custom-root install is discovered without UNSLOTH_SD_CPP_PATH; start
sd-cli with the parent-death child_popen_kwargs so it is not orphaned on a backend
crash; reap the SIGKILLed child (proc.wait) so a cancel/timeout does not leave a zombie.
- tests: Zip-Slip rejection, normal extraction, studio-home discovery.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 4) review round 2: collect sd-cli batch outputs
Codex review: when batch_count > 1, stable-diffusion.cpp's save_results() writes
the numbered files <stem>_<idx><suffix> (base_0.png, base_1.png, ...) instead of
the literal --output path. SdCppEngine.generate checked only the literal path, so
a batch generation would exit 0 and then raise 'no image' (or return a stale
file). generate now returns the literal path when present and otherwise falls
back to the numbered siblings; single-image behavior is unchanged.
Test: a fake sd-cli that writes img_0.png/img_1.png (not img.png) is collected
without error.
* Studio diffusion (Phase 6) review round 2: img2img source dims + upscale repeats
Codex review on the native engine arg builder:
- build_sd_cpp_command emitted --width/--height unconditionally, so an
img2img/inpaint/edit run that left dims unset forced a 1024x1024 resize/crop of
the input. width/height are now Optional (None = unset): an image-conditioned
run (init_img or ref_images) with unset dims omits the flags so sd.cpp derives
the size from the input image (set_width_and_height_if_unset); a plain txt2img
run with unset dims keeps the prior 1024x1024 default; explicit dims are always
honored. width/height are read only by the builder, so the type change is local.
- build_sd_cpp_upscale_command used a truthiness guard (params.repeats and ...)
that silently swallowed repeats=0 into sd-cli's default of one pass, turning an
explicit no-op into a real upscale. It now rejects repeats < 1 with ValueError
and emits the flag for any explicit value != 1.
Tests: img2img unset dims omit width/height (init_img and ref_images), explicit
dims emitted, txt2img keeps 1024; upscale rejects repeats=0 and omits the flag at
the default. (Two pre-existing binary-discovery tests fail only because a real
sd-cli is installed in this dev environment; unrelated to this change.)
* Studio diffusion (Phase 9) review round 2: correct prequant allowlist doc
Codex review: the transformer_prequant_path field description still told operators
to enable local checkpoints with UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH=1, but the
prior security fix made that variable a directory allowlist -- _allowed_prequant_roots
deliberately drops bare on/off toggle tokens (1/true/yes/...). An operator
following the documented =1 would have every transformer_prequant_path request
silently refused. The description now states it must name one or more allowlisted
directories and that a bare on/off value is not accepted.
Test: asserts the field help references UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH, does
not say =1, and describes an allowlist/directory (guards against doc drift).
* Studio diffusion (Phase 10) review round 2: cudnn/flash3 gating + registry reset
Codex review on attention-backend selection:
- Explicit attention_backend=cudnn skipped the SM80 gate that auto applies, so on
pre-Ampere NVIDIA (T4 SM75 / V100 SM70) it set fine then crashed at the first
generation with no fallback. select_attention_backend now applies
_cudnn_attention_supported() to an explicit cuDNN request too.
- flash3 used a minimum-only capability gate (>= SM90), so an explicit flash3 on a
Blackwell B200 (SM100) passed and then failed at generation -- FlashAttention 3
is a Hopper-SM90 rewrite with no Blackwell kernel. The arch gate is now a
(min, max-exclusive) range: flash3 is SM9x-only, flash4 stays SM100+.
- apply_attention_backend's success path left diffusers' process-wide active
backend pinned to the kernel it set; a later component whose processors are
unconfigured (backend None) would inherit it. It now resets the global registry
to native after a successful per-transformer set (the transformer keeps its own
backend), best-effort. Also fixed _active_attention_backend: get_active_backend()
returns a (name, fn) tuple, so the prior code stringified the tuple and never
matched a name, defeating the native-restore short-circuit.
Tests: explicit cudnn dropped below SM80; flash3 dropped on SM100 and allowed on
SM90; global registry reset after a successful set; _active_attention_backend
reads the tuple return.
* Studio diffusion (Phase 11) review round 2: keep GH200/B300 on the fp8 ladder
Codex review: _DATACENTER_GPU_TOKENS omitted GH200 (Grace-Hopper) and B300
(Blackwell Ultra), though it has the distinct GB200/GB300 superchip tokens. So
_is_consumer_gpu returned True for 'NVIDIA GH200 480GB' / 'NVIDIA B300', and the
auto ladder moved int8 ahead of fp8 on those data-center parts -- contradicting
llama_cpp.py's datacenter regex, which lists both. Added GH200 and B300 so they
are treated as data-center class and keep the intended fp8-first behavior.
Test: extends the datacenter parametrize with 'NVIDIA B300' and
'NVIDIA GH200 480GB' (now _is_consumer_gpu False).
* Studio diffusion (Phase 14) review round 2: apply int8 M=1 exclusion in the builder
Codex review: the M=1 modulation/embedder exclusion was wired only into the dense
runtime quantiser; the offline builder scripts/build_prequant_checkpoint.py called
make_filter_fn(min_features) with no exclusion. So an int8 prequant checkpoint
quantised the AdaLN modulation and conditioning-embedder linears, and loading it
via transformer_prequant_path (the load path only loads already-quantised tensors,
it can't re-skip them) reintroduced the torch._int_mm M=1 crash this phase fixes
for the runtime path.
Extracted int8_exclude_name_tokens(scheme) as the single source of truth (int8 ->
the M=1 exclusion, every other scheme -> none) and use it in both the runtime
quantiser and the builder, so a prequant artifact's quantised-layer set always
matches the runtime. fp8/fp4/mx artifacts are byte-identical (empty exclusion).
Test: int8_exclude_name_tokens returns the exclusion for int8 and () for
fp8/nvfp4/mxfp8.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion (Phase 16) review round 2: native CPU arbiter, status offload, load race
Codex review on the native-engine routing:
- The /images/load route took the GPU arbiter (acquire_for(DIFFUSION) -> evict chat)
unconditionally after engine selection. A native sd.cpp load on a pure-CPU host
never touches the GPU, so that needlessly tore down the resident chat model. The
handoff is now gated: diffusers always takes it, a force-native sd.cpp load on a
CUDA/XPU/MPS box still takes it, but a native sd.cpp load on a CPU host skips it.
- sd_cpp status() hardcoded offload_policy 'none' / cpu_offload False even when
_run_load computed real offload flags (balanced/low_vram/cpu_offload off-CPU), so
the setting was unverifiable. status now derives them from state.offload_flags
(still 'none' on CPU, where the flags are empty).
- _run_load committed the new state without cancelling/waiting on a generation that
started during the (slow) asset download, so a stale sd-cli run against the OLD
model could finish afterward and persist an image from the previous model once the
new load reported ready. The commit now signals the in-flight cancel and waits on
_generate_lock before swapping _state (taken only at commit, so the download never
serialises against generation), mirroring the diffusers load path.
Tests: CPU native load skips the arbiter while a GPU native load takes it; status
reports offload active when flags are set; _run_load cancels and waits for an
in-flight generation before committing.
* Studio diffusion (Phase 14) review round 2: align helper name with the stack
Rename the int8 exclusion helper to exclude_tokens_for_scheme, matching the
identical helper already present higher in the diffusion stack (Phase 16). The
helper definition, the runtime quantiser call, and the offline builder are now
byte-identical to that version, so the two branches no longer introduce a
divergent name for the same single-source-of-truth and the stack merges without
a conflict on this fix. No behavior change.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion: eager patches + torch.compile cache speed phase
Adds the opt-in speed path for the GGUF diffusion transformer behind a
selectable speed mode (default off, so output is unchanged until a profile
is chosen):
- diffusion_eager_patches.py: shared eager fast-paths (channels_last,
attention/backend selection, fused norms and QKV) installed at load and
rolled back on unload or failed load.
- diffusion_compile_cache.py / diffusion_gguf_compile.py: a persistent
torch.compile cache and the GGUF-transformer compile wiring.
- diffusion_arch_patches.py: architecture-specific patches.
- diffusion_patch_backend.py: shared install/restore plumbing.
- diffusion_speed.py: speed-profile planning.
Tests for each module plus the benchmarking and probe scripts used to
measure speed, memory, and accuracy of the path.
* Studio diffusion: image workflows (safetensors, image-conditioned, editing) + Images UI
Backend:
- Load non-GGUF safetensors models: full bnb-4bit pipelines and single-file
fp8 transformers, gated to the unsloth org plus a curated allowlist.
- Image-conditioned workflows built with Pipeline.from_pipe so they reuse the
loaded transformer/VAE/text-encoder with no extra VRAM: img2img, inpaint,
outpaint, and a hires-fix upscale pass.
- Instruction editing as its own family kind (Qwen-Image-Edit-2511,
FLUX.1-Kontext-dev) and FLUX.2-klein reference conditioning (single and
multi-reference) plus klein inpaint.
- Auto-resize odd-sized inputs to a multiple of 16 (and resize the matched
mask) so img2img/inpaint/edit no longer reject non-/16 uploads. Bound the
decoded image size and cap upscale output to avoid OOM on large inputs.
- Fixes: from_pipe defaulting to a float32 recast that crashed torchao
quantized transformers; image-conditioned calls forcing the slider size
onto the input image. Native sd.cpp engine rejects image-conditioned and
reference requests it cannot serve.
Frontend:
- Redesigned Images page with capability-gated workflow tabs (Create,
Transform, Inpaint, Extend, Upscale, Reference, Edit), a brush mask editor,
client-side outpaint, and a multi-reference picker.
- Advanced options moved to a right-docked panel mirroring Chat: closed by
default, toggled by a single fixed top-bar button that stays in place.
sd.cpp installer: pin the release, verify each download's sha256, add a
download timeout, and make the source repo configurable for a future mirror.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio Images: correct the Advanced panel comment (closed by default, fixed toggle)
* Studio diffusion: LoRA adapters for the Images workflow
Add community LoRA support across both diffusion backends, the single
biggest step toward broad image-workflow coverage.
Backend
- New shared module core/inference/diffusion_lora.py: adapter discovery
(local scan + curated catalog + owner/name[:file] Hub refs), download
via hf_hub_download_with_xet_fallback, alias sanitization, native
managed-dir materialization with collision-broken aliases, prompt-tag
injection (deduped against user-typed tags), and a supports_lora gate.
- Native sd-cli: resolve + materialize selected LoRAs into a per-run
managed dir, inject <lora:ALIAS:w> tags, pass --lora-model-dir with
--lora-apply-mode auto. The arg builder already emitted these flags.
- Diffusers: non-fused load_lora_weights + set_adapters manager, tracked
on the pipe so an unchanged selection is a no-op and a model swap
resets; cleared on unload. Never fuses (breaks quantized transformers
and blocks live weight tweaks).
- Gated off where unsupported: torchao fp8/int8 dense, GGUF-via-diffusers,
and native Qwen-Image (no LoRA name-conversion branch upstream).
- Request contract: optional loras on DiffusionGenerateRequest; empty or
omitted is identical to today. supports_lora surfaced in status; chosen
LoRAs persisted in gallery recipe metadata.
- New GET /api/models/diffusion-loras for the picker (family-filtered).
Frontend
- Repeatable multi-LoRA picker (adapter select + weight slider 0..2 +
remove), gated by the loaded model's supports_lora and family, max 8.
Tests
- New test_diffusion_lora.py (14): helpers, request validation, native
tag/dir wiring, diffusers set_adapters manager, supports_lora matrix.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio diffusion: ControlNet for the Images workflow (diffusers)
Add ControlNet conditioning, the #2 most-used diffusion workflow after
LoRA, on the diffusers backend for the families with ControlNet pipelines
(FLUX.1 and Qwen-Image), with Union models as the default picks.
Backend
- New core/inference/diffusion_controlnet.py: family-gated discovery
(curated Union models + local dirs + bare owner/name repos), resolution
to a loadable repo/dir, control-image preprocessing (passthrough +
a dependency-free canny edge map), and a supports_controlnet gate.
- diffusion.py: a ControlNet manager parallel to the LoRA one. Loads the
(small) ControlNet model once via from_pretrained (cached by id) and
builds the family's ControlNet pipeline via Pipeline.from_pipe(base,
controlnet=model), reusing the resident base modules at their loaded
dtype (no reload, no recast). Passes the control image + conditioning
scale + guidance start/end at generate time; cleared on unload.
- Families: FLUX.1 -> FluxControlNetPipeline/Model, Qwen-Image ->
QwenImageControlNetPipeline/Model. Others declare none (gated off).
- Gated off for the native engine, GGUF-via-diffusers, and torchao
fp8/int8 dense (same rule as LoRA). v1 conditions txt2img only.
- Request contract: optional controlnet on DiffusionGenerateRequest;
supports_controlnet in status; the choice persisted in gallery meta.
- New GET /api/models/diffusion-controlnets for the picker.
Frontend
- A ControlNet control in the Images rail (model select + control-image
upload + control-type select + strength slider), gated by the loaded
model's supports_controlnet + family, shown for text-to-image.
Tests
- New test_diffusion_controlnet.py (10): discovery/resolve/preprocess/gate
helpers, request validation, family wiring, and the diffusers pipe
manager (loads once, caches, from_pipe with controlnet, rejects
unsupported families).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio ControlNet: show the picker on the Create tab (workflow id is 'create', not 'txt2img')
The ControlNet control gated on workflow === 'txt2img', but the Images workflow tab ids are create/transform/inpaint/extend/upscale/reference/edit -- there is no 'txt2img'. So the picker never rendered even with a ControlNet-capable model loaded. Gate on 'create' (the text-to-image tab) for both the picker and the request wiring. Found via a live Playwright capture of the running Studio.
* Studio: do not force diffusers pipelines cross-tagged gguf into the GGUF variant expander
Some diffusers image repos (e.g. unsloth/Qwen-Image-2512-unsloth-bnb-4bit) carry a
stray "gguf" tag on the Hub but ship no .gguf files. The model search classified
them as GGUF from the bare tag, so the picker rendered the GGUF variant expander,
which then dead-ended at "No GGUF variants found." Trust the bare gguf tag only when
the repo is not a diffusers pipeline; the -GGUF name suffix and real gguf metadata
(populated via expand=gguf) remain authoritative, so genuine GGUF repos are unaffected.
* Studio Images: load non-curated unsloth/on-device diffusers repos instead of no-op
handleModelSelect only loaded curated safetensors ids and GGUF variant picks; any other
non-GGUF pick (an on-device diffusers folder, or a future unsloth diffusers image repo
surfaced by search) silently did nothing. Treat such a pick as a full diffusers pipeline
load when the id is unsloth-hosted or on-device (the backend infers the family + base repo
and gates loads to unsloth/* or local paths), and show a clear message otherwise instead
of silently ignoring the click. Curated and GGUF paths are unchanged.
* Studio Images: keep curated safetensors models in Recommended after download
The curated bnb-4bit / fp8 diffusion rows were filtered out of the Images picker's
Recommended list once cached (curatedSafetensorsRows dropped anything in downloadedSet),
so they vanished from the picker after the first load and could only be found by typing an
exact search. The row already renders a downloaded badge, matching how GGUF Recommended
rows stay visible when cached. Drop the exclusion so the curated safetensors always list.
* Studio diffusion LoRA: sanitize dots out of adapter aliases
The LoRA alias is used as the diffusers PEFT adapter name, and PEFT rejects names
containing "." (module name can't contain "."). sanitize_alias kept dots, so a LoRA whose
filename carries a version tag (e.g. Qwen-Image-2512-Lightning-8steps-V1.0-bf16) failed to
apply with a 400. Replace dots too; the alias stays a valid native <lora:NAME:w> filename
stem. Adds regression coverage for internal dots.
* Studio Images: clarify the GGUF transformer-quant Advanced control
Renamed the confusing "Transformer quant / GGUF default" control to "GGUF speed mode"
with an "Off (run the GGUF)" default, and reworded the hint to state plainly that FP8/INT8/
FP4 load the FULL base model (larger download + more VRAM) rather than re-packing the GGUF,
falling back to the GGUF if it can't fit. Behavior unchanged; labels/hint only.
* Studio Images: list on-device unsloth diffusion models in the picker
The Images picker's On Device tab hid every non-GGUF cached repo whenever a
task filter was active, so downloaded unsloth diffusion pipelines (bnb-4bit
and FP8 safetensors) never showed up there. List cached repos that pass the
task gate, limited under a filter to unsloth-hosted ones so base repos (which
fail the diffusion load trust gate) don't appear only to dead-end on click.
Chat behavior is unchanged: the task gate still drops image repos there.
* Studio: hide single-file image checkpoints from the chat model picker
The chat picker treats a cached repo as an image model, and hides it, only
when it ships a diffusers model_index.json. Single-file, ComfyUI, and
ControlNet image checkpoints (an FP8 Qwen-Image, a z-image safetensors, a
Qwen-Image ControlNet) carry none, so they surfaced as loadable chat models.
Fall back to resolving the repo id against the known diffusion families, the
same resolver the Images backend loads from, so these checkpoints are tagged
text-to-image and stay in the Images picker only.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio Images: add the FLUX.2-dev model family
Loading unsloth/FLUX.2-dev-GGUF failed because detect_family knew only the
Qwen3-based FLUX.2-klein, so FLUX.2-dev (the full, Mistral-based Flux2Pipeline)
resolved to nothing and the load errored. Add a flux.2-dev family: Flux2Pipeline
+ Flux2Transformer2DModel over the black-forest-labs/FLUX.2-dev base repo (gated,
reachable with an HF token), with its FLUX.2 32-channel VAE and Mistral text
encoder wired for the sd-cli path from the open Comfy-Org/flux2-dev mirror.
text-to-image only: diffusers 0.38 ships no Flux2 img2img / inpaint pipeline for
dev. Frontend gets sensible dev defaults (28 steps, guidance 4), distinct from
klein's turbo defaults. Verified live: GGUF load resolves the family + gated base
repo and generates a real 1024x1024 image on GPU.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio Images: clearer error for an unsupported diffusion model
When a repo id resolves to no diffusion family the load raised 'Could not infer a
diffusion family... Pass family_override (z-image)', which points at an unrelated
family and doesn't say what is supported. Replace it with a message that lists the
supported families (from a new supported_family_names helper) and notes that video
models and image models whose diffusers transformer has no single-file loader are
not supported. Applies to both the diffusers and native sd.cpp load paths. Also
refreshes two stale family-registry comments that still called FLUX.2-dev omitted.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add SDXL diffusion family (U-Net pipeline support)
SDXL is the first U-Net family in the diffusion backend: its denoiser is
pipe.unet (UNet2DConditionModel), not a DiT pipe.transformer, and a single-file
.safetensors is the whole pipeline rather than a transformer-only file. The
backend previously assumed a DiT transformer everywhere, so add the two hooks a
U-Net family needs and register SDXL.
DiffusionFamily gains denoiser_attr ("transformer" for DiT, "unet" for SDXL) and
single_file_is_pipeline (SDXL loads a single file via pipeline_class.from_single_file
with the base repo as config, instead of transformer_class.from_single_file plus a
companion assembly). _align_vae_dtype now reads the denoiser generically so img2img
and inpaint keep the VAE and U-Net dtypes aligned.
The non-GGUF trust gate is extended with a short, exact-match, safetensors-only
allowlist of official base repos (the SDXL base/refiner and sdxl-turbo), because
SDXL ships only as a full pipeline and has no unsloth-hosted GGUF. Local paths stay
trusted as before; a random repo, even one that detects as SDXL, is still rejected.
The image-conditioned and ControlNet workflows are the standard SDXL pipelines,
built around the resident modules via from_pipe like every other family, so SDXL
gets txt2img, img2img, inpaint, outpaint, upscale, LoRA and ControlNet. There is no
native sd.cpp mapping yet, so the no-GPU route falls back to diffusers.
Frontend catalog gains SDXL Base 1.0 and SDXL Turbo entries with SDXL step/guidance
defaults (Turbo: few steps, no CFG; base: ~30 steps, real CFG).
Tests: new test_diffusion_sdxl.py (family shape, detection, trust allowlist, model
kind, U-Net VAE-dtype alignment, LoRA gate) plus loader-branch tests in
test_diffusion_backend.py (pipeline-kind from_pretrained, single-file whole-pipeline
from_single_file, allowlist accept/reject). Verified live on GPU: sdxl-turbo loads
both as a pipeline and as a single file and generates coherent txt2img + img2img.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Images: LoRA free-text Hugging Face entry + recipe round-trip
The backend has always accepted a bare Hugging Face repo id (owner/name, or
owner/name:weight-file.safetensors) as a LoRA, downloading and applying it. But the
picker only rendered when the curated catalog had entries, and the catalog is empty,
so there was no UI path to apply any LoRA. Show the LoRA section whenever the loaded
model supports LoRA, and replace the curated-only dropdown with a text input: type a
Hub repo id, or pick a discovered adapter from a datalist of suggestions when the
catalog is populated.
Also restore LoRAs when loading a recipe. restoreSettings now parses the recipe's
"id:weight" strings (splitting on the last colon, since the id itself may contain one
for a specific weight file) back into the selection, so replaying a saved image
reproduces its adapters. The generate payload trims hand-typed ids and drops empty /
zero-weight rows, and a model swap clears the selection (a LoRA is family-specific)
without discarding a free-text pick that is not in the curated list.
* Add diffusion LoRA training (SDXL text-to-image)
First diffusion training path in Studio: train a LoRA on the SDXL U-Net from an
image + caption dataset and export it as a diffusers .safetensors that the existing
diffusion LoRA loader (and any diffusers pipeline) can load.
core/training/diffusion_lora_trainer.py:
- DiffusionLoraConfig with validation/defaults (rank, alpha, targets, lr, steps, grad
accumulation, resolution, min-SNR gamma, gradient checkpointing, lr scheduler, seed,
mixed precision).
- discover_image_caption_pairs: captions from metadata.jsonl / captions.jsonl, per-image
.txt/.caption sidecars, or a dreambooth instance_prompt fallback (pure, unit-tested).
- run_diffusion_lora_training: the loop -- freeze base, PEFT-wrap the U-Net attention
projections, VAE-encode (fp32 VAE to avoid the SDXL fp16 overflow), sample noise +
timesteps, predict, MSE loss with optional min-SNR weighting (epsilon / v-prediction),
AdamW + get_scheduler + grad accumulation + grad clipping, then export via
save_lora_weights. Emits worker-protocol events (model_load_*, progress, complete) and
polls should_stop for a clean stop with a partial save.
- run_diffusion_training_process: mp.Queue subprocess adapter (event_queue / stop_queue),
so the training worker can spawn it; plus a CLI entry point.
Only SDXL (U-Net) is trained here; DiT families and the Studio UI form + route wiring are
follow-ups. The trainer is decoupled and worker-ready.
Tests: test_diffusion_lora_trainer.py covers caption discovery (metadata / sidecar /
instance prompt / skip-uncaptioned / errors), config normalisation + validation, the SDXL
add-time-ids, and the dict->config adapter. Verified live on GPU: a 60-step SDXL LoRA run
lowers the loss, exports a ~45 MB adapter, and loading it back shifts generation from
baseline (mean abs pixel diff ~55/255).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* diffusion trainer: emit learning_rate in progress events (Studio pump compatibility)
The Studio training pump reads 'learning_rate' from progress events; the diffusion
trainer emitted 'lr'. Rename the field (and the CLI reader) so the trainer's events are
directly consumable by the existing training status/SSE machinery when it is wired into
the worker, without a translation shim.
* Wire diffusion LoRA training into the Studio API
Make the SDXL LoRA trainer reachable from the app with a small, self-contained job
service and JSON routes, deliberately separate from the LLM TrainingBackend (whose
lifecycle -- LLM config build, per-run SQLite rows, matplotlib plots, transfer-to-chat-
inference -- is text-training specific and would mis-handle a diffusion run).
core/training/diffusion_training_service.py: DiffusionTrainingService runs one job at a
time -- validate the config cheaply (before any spawn), spawn the trainer subprocess
(spawn context, parent-lifetime bound), pump its events (model_load_* / progress /
complete / error) into an in-memory status snapshot, and support a clean stop. The
subprocess context and target are injectable so the full start -> pump -> status ->
complete path is unit-tested without real multiprocessing or torch.
routes/training.py: POST /api/train/diffusion/start (400 on a bad config, 409 when a job
is already running), POST /api/train/diffusion/stop, GET /api/train/diffusion/status
(JSON poll). models/training.py: DiffusionTrainingStartRequest + response schemas
mirroring DiffusionLoraConfig, so model_dump() passes straight through.
Tests: test_diffusion_training.py -- service happy path, bad-config-before-spawn,
concurrent-job rejection, clean stop, crash-without-terminal-event, event transitions;
plus route wiring via the FastAPI TestClient (start / 422 / 400 / 409 / status / stop)
with a mocked service. The diffusion trainer's progress events already use the field
names this path expects.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Import diffusion training schemas from models.training directly
The import-hoist lint flags newly re-exported names in the models/__init__.py hub as
unused (it does not treat __all__ membership as a use). Import the three diffusion
training schemas straight from models.training in routes/training.py, where they are
used in the route annotations and calls, and drop the __init__ re-export.
* Remove stray async task scratch outputs committed by mistake
* ControlNet: reject filesystem-like ids and do not cache a model past an unload race
Two review findings on the ControlNet path:
- resolve_controlnet's bare-repo fallback accepted any id with a slash, so a
path-shaped id (/tmp/x, ../x) reached from_pretrained as a local directory.
Restrict the fallback to a strict owner/name HF repo id shape.
- _controlnet_pipe now re-checks the cancel event after the blocking
from_pretrained: an unload that raced the download had already cleared the
caches, so caching the late module would pin it past the unload.
* Pipeline prefetch: fetch only the default torch weights
A full-pipeline prefetch kept every repo file outside assets/, so an official
repo that ships multiple formats (SDXL Base: fp16 variants, ONNX, OpenVINO,
Flax, a top-level single-file twin) downloaded tens of GB from_pretrained never
loads. Skip non-torch exports and dtype-variant twins in
_pipeline_file_downloaded, and drop a component .bin when the same directory
carries a picked safetensors weight (diffusers' own preference).
* Diffusion LoRA training: fall back to fp16 when CUDA lacks bf16
The default mixed_precision=bf16 hard-fails on pre-Ampere GPUs (T4 / V100 /
RTX 20xx) which have no bf16 compute; check torch.cuda.is_bf16_supported()
and drop to fp16 there.
* Diffusion training service: join the old pump outside the lock
start() joined a finished job's pump thread while holding the service lock,
but the pump's final state writes need that same lock, so the join always
burned its full timeout and a stale pump could then overwrite the new job's
state. Join outside the lock (with a re-check after), and fence _apply_event
and the exit handler by process identity so a superseded pump can never touch
the current job's state. Adds regression tests for both.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion LoRA training: harden config handling, cancellation, SDXL conditioning, and safety
Addresses review findings on the SDXL LoRA trainer:
- Gate the base model with the same trust check as inference (unsloth/*, allowlisted
official bases, or a local path) before from_pretrained, so an untrusted remote repo
is never fetched or deserialised.
- Check the stop signal before the (slow) model load, not only between steps, so a
cancel during download is honoured; a stop may carry save=False to cancel without
leaving a partial adapter.
- Per-sample SDXL add_time_ids from the actual crop (original size + crop offset, with
the offset mirrored on horizontal flip) instead of a fixed uncropped-square tensor.
- Apply EXIF orientation before resize/crop so rotated photos train upright.
- Skip gradient clipping when max_grad_norm <= 0 (the Studio 'disable' value) instead
of scaling every gradient to zero.
- Coerce Studio config strings/blanks: learning_rate string to float, blank hf_token to
anonymous, gradient_checkpointing 'none'/'true'/'unsloth' to bool; reject a zero/negative
lora_alpha or learning_rate.
- Alias the generic Studio training payload keys (model_name/max_steps/batch_size/lora_r/
lr_scheduler_type/random_seed) onto the diffusion field names.
- Mirror the trained adapter into loras/diffusion so the Images LoRA picker discovers it.
- Report worker exceptions in both message and error keys so the failure is not lost.
Adds regression tests for the config coercion/validation and aliasing.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* ControlNet: address review findings on the diffusers path
- resolve_controlnet enforces catalog family compatibility so a direct API call
cannot load a ControlNet built for another family through the wrong pipeline.
- Unknown ControlNet ids now surface as a 400 (call site maps FileNotFoundError
to ValueError) instead of a generic 500.
- strength 0 disables ControlNet entirely, so a no-op selection never pays the
download / VRAM cost; the control image is decoded and validated BEFORE the
ControlNet is resolved or built, so a malformed image fails fast for the same reason.
- ControlNet loads use the base compute dtype (state.dtype is a display string,
not a torch.dtype, so it silently fell back to float32) and honor the base
offload policy via group offloading instead of forcing the module resident.
- Empty/malformed HF token coerced to anonymous access.
- Flux Union ControlNet control_mode mapped from the selected control type.
- resolve_controlnet drops the unused hf_token/cancel_event params.
- ControlNetSpec validates guidance_start <= guidance_end (clean 422).
- Images UI ControlNet Select shows its placeholder when nothing is selected.
Adds regression tests for family enforcement and the union control-mode map.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion training API: LLM interlock, pre-spawn VRAM free, path containment, no dropped knobs
Four review findings on the diffusion training start path:
- It spawned the SDXL trainer without checking the LLM TrainingBackend, so a
start while an LLM run was active put two trainers on the same GPU. Add a
symmetric interlock: diffusion start returns 409 when LLM training is active,
and LLM start refuses while a diffusion job is active.
- It went straight to service.start() without freeing GPU residents. Add a
pre-spawn free of the export subprocess, the resident Images pipeline (with an
arbiter release), and chat models, mirroring the LLM start path.
- data_dir / output_dir were passed through unresolved, so Studio-relative names
failed and absolute paths bypassed containment. Resolve them with
resolve_dataset_path / resolve_output_dir before spawn (400 on an uncontained
path).
- The request model dropped max_grad_norm and lora_target_modules, so runs that
set them trained with defaults. Add both fields.
The gemini pump-join deadlock was already fixed earlier (join outside the lock +
proc-identity fence). Note: honoring a stop DURING model load is a trainer-loop
change owned by the diffusion training engine PR (should_stop polled before the
first optimizer step). Adds route + model regression tests.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion LoRA: harden resolution, native tag precedence, and diffusers teardown
Address review findings on the LoRA path:
- resolve_one: normalise a blank/whitespace hf_token to None (anonymous access)
and reject a client-supplied weight file with traversal / absolute path.
- resolve_specs: convert FileNotFoundError from an unknown/stale id to ValueError
so the route returns 400 instead of a generic 500.
- _scan_local: disambiguate local adapters that share a stem (foo.safetensors vs
foo.gguf) so each is uniquely addressable.
- inject_prompt_tags: the backend-validated weight now wins over a user-typed
<lora:ALIAS:...> for a selected adapter; unselected user tags are left alone.
- diffusers _apply_loras: reject a .gguf adapter with a clear error before touching
the pipe (diffusers loads safetensors only).
- _unload_locked: drop the explicit unload_lora_weights() on teardown; the pipe is
dropped wholesale (freeing adapters), so the previous call could race an in-flight
denoise on the same pipe.
- Images page: use a stable LoRA key and clear the selection (not just the options)
when the catalog refresh fails.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Diffusion: guard trust check against OSError and validate conditioning inputs
- _is_trusted_diffusion_repo: wrap Path.exists() so a repo id with invalid
characters (or a bare owner/name id) can't raise OSError; treat any failure as
not-a-local-path and fall through to the unsloth/ allowlist. validate_load_request
still raises the clear FileNotFoundError for a genuinely missing local pick.
- generate(): reject mask_image / upscale / reference_images supplied without an
input image, and reject reference_images on a family that does not support
reference conditioning, instead of silently degrading to txt2img / img2img.
* SDXL: reject GGUF up front, skip unused base weights, drop refiner, and harden helpers
Addresses review findings on the SDXL family:
- Reject a GGUF load for single_file_is_pipeline families (SDXL) in validate_load_request,
before the route evicts the current model; SDXL has no transformer-only GGUF variant.
- Skip base-repo weight files when a whole-pipeline single file is loaded: from_single_file
(config=base) needs only the base config/tokenizer/scheduler, so a local .safetensors no
longer triggers a multi-GB base download.
- Remove the SDXL refiner from the non-GGUF trust allowlist: it is an img2img-only pipeline
but this backend loads every sdxl repo as the base txt2img pipeline.
- Normalize a blank/whitespace hf_token to None once in load_pipeline so every load branch
degrades to anonymous instead of erroring on a malformed token.
- Read the denoiser dtype from a parameter (compile-wrapped modules may lack .dtype) and
access state.family.denoiser_attr directly.
Adds/updates regression tests for the trust allowlist, GGUF rejection, and base-config filter.
* Images: preserve restored LoRAs through model load and never send hidden LoRAs
- The LoRA effect cleared the selection on every load->capable transition, which
wiped adapters restored from a gallery recipe before the model finished loading.
Track the previously-loaded family in a ref and clear only on a real family swap;
keep the selection on the initial load and on unload.
- Gate the generate payload's loras on loraCapable so a restored selection that is
hidden (loaded model does not support LoRA) is never sent to the backend.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Merge diffusion-sdxl into diffusion-lora-ux; keep options-only LoRA catch
The catalog-refresh .catch from the lower branch clears the selected adapters
too, which is right for its catalog-only picker but wrong here: this picker
holds free-text HF repo ids that are valid without being in the catalog, so a
transient refresh failure must not wipe them. Family swaps still clear the
selection and hidden LoRAs are never sent.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Align the VAE to the denoiser's first FLOATING dtype, not its first parameter
A GGUF-quantized transformer's leading parameters are packed uint8 storage,
so reading next(parameters()).dtype handed nn.Module.to() an integer dtype
and every image-conditioned generation on a GGUF model (Qwen-Image-Edit)
failed with a 500. Probe the parameters for the first floating dtype, treat
an all-integer module as a no-op, and also catch TypeError so an unexpected
dtype can never break generation. Regression test included.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Count LR scheduler warmup/decay in optimizer steps, not micro-steps
lr_sched.step() runs once per outer optimizer step (after the gradient
accumulation inner loop), for train_steps total. The scheduler was
configured with num_warmup_steps and num_training_steps multiplied by
gradient_accumulation_steps, so with accumulation > 1 a warmup or
non-constant schedule stretched past the run and never reached the
intended decay. Count both in optimizer steps.
* Address Codex review findings on the image-workflows PR
Keep diffusion.py importable without torch: the compile/arch patch modules
import torch at module level, so import them lazily at their load/unload
call sites instead of at module load. This restores the torchless contract
so get_diffusion_backend() works on a CPU/native sd.cpp install.
Match family reject keywords and aliases as whole path/name segments, not
raw substrings, so an unrelated word like edited, edition, or kontextual no
longer misroutes or hides a valid base image model, while supported edit
families (Qwen-Image-Edit, FLUX Kontext) still resolve. Mirror the same
segment matching in the picker task filter.
Route FLUX.2-dev native guidance through --guidance like the other FLUX
families rather than --cfg-scale. Reject native upscale requests that have
no input image. Read image header dimensions and reject over-limit inputs
before decoding pixels, so a crafted small-payload image cannot spike
memory. Reject an upscale that would shrink the source below its input
size. Validate the model_kind against the filename extension before the
GPU handoff. Estimate a local diffusers pipeline's size from its on-disk
weights so auto memory planning does not skip offload and OOM. Report
workflows: [txt2img] from the native backend status so the Create tab
stays enabled for a loaded native model. Clamp the outpaint canvas to the
backend's 4096px decode limit.
Adds regression tests for segment matching and kind/extension validation.
* Guard inference loads and worker lifetime against diffusion training
Teach the chat and image load guards about an active diffusion (SDXL) LoRA
job: a chat load is refused (its footprint cannot be fit-checked against the
trainer) and an image load is refused outright, mirroring the existing LLM
training guards, so a load can no longer allocate GPU memory alongside the
trainer and undo the pre-start cleanup.
Bind the diffusion trainer subprocess to the parent's lifetime and scrub the
native path lease secret from it by running the child through
run_without_native_path_secret, matching the inference/export/LLM workers, so
a Studio crash or kill no longer leaves the trainer holding the GPU.
Reset in_model_load on the complete and error terminal events: a stop or
failure during model loading otherwise leaves the status reporting a stale
loading indicator after the job has ended.
* Harden diffusion LoRA handling on the diffusers and native paths
Reject LoRA on a torch.compile'd diffusers transformer (Speed=default/max):
diffusers requires the adapter loaded before compilation, so applying one to
the already-compiled module fails with adapter-key mismatches. The status
gate now hides the picker and generate raises a clear message instead.
Convert a cancelled Hub LoRA download (RuntimeError Cancelled) to the
diffusion cancellation sentinel in resolve_specs, so an unload/superseding
load during resolution maps to a 409 instead of a generic server error.
Drop weight-0 LoRA rows before the native support gate so a request carrying
only disabled adapters stays a no-op on families where native LoRA is
unsupported, matching the diffusers path.
Reject duplicate LoRA ids in the request model: both apply paths suffix
colliding names, so a repeated id would stack the same adapter past its
per-adapter weight bound.
Strip all user-typed <lora:...> prompt tags on the native path (only the
selected adapters are materialized in the managed lora-model-dir, so an
unselected tag can never resolve), and restore saved LoRA selections from a
gallery recipe so restore reproduces a LoRA image.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Harden ControlNet resolve, gallery metadata, and the control-type picker
Check cancellation immediately after a ControlNet from_pretrained and before
any device placement, so an unload/eviction that raced the download does not
allocate several GB onto the GPU after the load was already cleared.
Require a loadable weight or shard index (not just config.json) before a local
ControlNet folder is advertised, so an interrupted copy is hidden instead of
failing deep in from_pretrained as a generic 500.
Do not record a strength-0 ControlNet in the gallery recipe: it is treated as
disabled and skipped, so the image is unconditioned and the metadata must not
claim a ControlNet was applied.
Build the control-type picker from the selected ControlNet's advertised
control_types instead of a hardcoded passthrough/canny pair, so a union model
with a precomputed depth or pose map sends the correct control_mode.
* Address further Codex findings on the image-workflows PR
- Persist the actual output image size in the gallery recipe instead of the
request sliders: Transform/Inpaint/Edit derive the size from the uploaded
image, Extend grows the canvas, and Upscale resizes it, so the sliders
recorded (and later restored) the wrong dimensions for those workflows.
- Reject a remote '*-GGUF' repo loaded as a full pipeline (no single-file
name) in validate_load_request, so the unloadable pick fails before chat is
evicted rather than deep in from_pretrained.
- Only publish an image-conditioned from_pipe wrapper to the shared aux cache
when the load is still current: from_pipe runs under the generate lock but
not the state lock, so an unload racing its construction could otherwise
cache a wrapper over torn-down modules that a later load would reuse.
- Verify the Windows CUDA runtime archive checksum before extracting it, like
the main sd-cli archive, so a corrupt or tampered runtime is rejected rather
than extracted next to the binary.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Refuse non-SDXL base models at diffusion training start
The trainer only supports the SDXL U-Net, but a FLUX / Qwen-Image / Z-Image
repo or a GGUF filename passed as base_model was accepted and then failed
minutes later inside StableDiffusionXLPipeline.from_pretrained with an
unrelated-looking error. Add a name-based guard in normalized() so known
DiT-family names and .gguf checkpoints are rejected up front, which the API
start route surfaces as an immediate 400 with a message that says exactly
which bases are trainable. Unrecognisable names still pass through so custom
local SDXL checkpoints keep working.
* Add diffusion dataset upload and training info endpoints
Training an image LoRA required knowing the Studio home layout and copying
files onto the server by hand, which is the most confusing step of the whole
flow. Two small endpoints fix that:
- GET /api/train/diffusion/info reports the datasets and outputs roots plus
every dataset folder that contains images (with image/caption counts), so
the UI can offer a picker instead of a blind free-text path.
- POST /api/train/diffusion/dataset uploads images and optional caption
.txt / metadata.jsonl files into a named folder under the datasets root,
creating it on first use and accumulating on repeat uploads so large sets
can arrive in batches. Names are validated to a single path component and
files stream to disk under the same per-upload size cap as LLM dataset
uploads. The returned name is a valid data_dir for /diffusion/start.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Validate diffusion training config before freeing the GPU
The start route freed resident GPU workloads (export, Images pipeline, chat)
before the service validated the config, so a start that was then refused,
now including a non-SDXL base model, tore down the user's loaded model for
nothing. Run the same cheap normalise pass first; the LLM path already
follows this rule via its before_spawn hook.
* [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>
Co-authored-by: oobabooga <112222186+oobabooga@users.noreply.github.com>
* (feat) Add project names to studio training runs to avoid models being overwritten when doing similar training runs
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update studio/frontend/src/features/export/export-page.tsx
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
* Update studio/frontend/src/features/export/export-page.tsx
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
* Update studio/frontend/src/features/export/export-page.tsx
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
* better project name sanitization, removed duplicated project name normalization
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* implement checkpoint scanning utilities and tests for base model inference
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Guard project_name against null and use leading important modifiers
* Fix/adjust training project names for PR #6512
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix/adjust training project names for PR #6512
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address project-name review feedback
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Show project names in training recents
* Keep GGUF export directories source-specific
---------
Co-authored-by: NZ-Linix <nz-linix@outlook.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: NZ-Linix <linus.ordowski@outlook.com>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
* Studio: require signed capability tokens for /p preview links
The public /p preview routes added in #6486 run model load and chat
generation as the admin user with no authentication. The only gate is the
preview ref, a deterministic outputs-root path (run or run/checkpoint) that
is guessable rather than secret. On a network-reachable Studio (--secure
tunnel or -H 0.0.0.0), an unauthenticated caller who guesses a ref can
consume GPU and probe a private fine-tuned checkpoint.
Make the share link an unguessable, revocable capability:
- Sign the canonical ref with a dedicated server-side secret (HMAC-SHA256,
stored in app_secrets, independent of the JWT/login secret).
- Require a valid token on every /p chat, models, and page request before
resolving a checkpoint or loading a model; missing or invalid tokens get a
generic 404 so the surface never confirms a ref exists.
- Accept the token via ?k= (browser link and preview page) or
Authorization: Bearer (OpenAI-compatible clients).
- Rotate the secret to revoke every outstanding link
(POST /api/settings/preview-links/rotate).
- Clamp preview generation (max_tokens/max_completion_tokens <= 1024, n = 1)
and set Referrer-Policy: no-referrer on the page so the token is not
leaked via Referer.
Training history hands the authenticated owner the signed token, and the
copy-link button builds /p/{ref}?k={sig}.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: honor a lower caller token limit in the preview clamp
Codex review: when only the legacy max_tokens was sent, the clamp left
max_completion_tokens at the 1024 default, and _effective_max_tokens prefers
max_completion_tokens, so a request like max_tokens=16 could still generate up
to 1024 tokens. Derive one effective limit (max_completion_tokens wins, else the
legacy max_tokens) and pin both fields to it so a caller's lower limit is kept.
* Studio: add preview kill switch, rate limit, and revoke-links UI
Follow-ups to the /p preview capability work:
- Public-sharing kill switch: a persisted setting (default on) gates the public
/p surface. When off, every preview request 404s even with a valid token, and
the owner UI stops offering share links. GET/PUT /api/settings/preview-sharing;
enforced in _verify_or_404.
- Per-IP rate limit on the preview chat route: a coarse in-process sliding-window
limiter (20 req/min/IP) returns 429 + Retry-After before the GPU lock is taken.
Client IP honors X-Forwarded-For only when UNSLOTH_STUDIO_TRUST_FORWARDED is
set, matching the login limiter's trust model.
- Settings UI: a "Preview sharing" section with the public-sharing toggle and a
"Revoke all preview links" button (confirm dialog) that rotates the secret.
Tests cover the kill switch (404 when off), the 429 path, the sliding window,
client-IP trust behavior, and the setting default.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: fix preview-fields sharing arg and refresh sigs after revoke
Codex review:
- P1: get_training_run_detail and update_training_run called _preview_fields
with only output_dir after it gained a required sharing_on parameter, raising
a 500 TypeError once get_run succeeded. Pass get_preview_sharing_enabled() at
both sites; add a detail-endpoint regression test.
- P2: after rotating the preview secret from settings, the history grid still
held stale preview_sig values, so a freshly copied link would 404. Emit
emitTrainingRunsChanged() after a successful revoke so the grid refetches
freshly signed refs.
* Studio: harden preview sharing controls (Codex review)
- Fail closed: a read failure on the preview-sharing kill switch now returns
False instead of defaulting to enabled, so an unavailable settings DB can't
reopen the public surface. A missing key still defaults to enabled.
- Per-IP rate limit behind the managed Cloudflare tunnel: client_ip now honors
CF-Connecting-IP when the socket peer is loopback, so tunneled visitors are
keyed by their real IP instead of collapsing onto the local cloudflared peer.
- GET /p no longer mints key/share_url when sharing is disabled; it returns
sharing_enabled=false so clients don't distribute links that 404.
- Settings UI: toggling public sharing emits the training-runs-changed event so
the history grid shows/hides Copy preview link without a manual refresh.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: harden preview rate limiter and IP keying (Opus review)
From a two-agent review of the PR:
- Rate limiter no longer evicts an active bucket when the table is full: a flood
of distinct keys could otherwise cycle out a throttled bucket and reset its
counter. Evict only aged-out buckets; if the table is full of live clients,
fail closed (deny the new key) instead.
- client_ip keys on the rightmost (proxy-appended) X-Forwarded-For hop when the
trust env is set; the leftmost is client-spoofable. Documented the
append/overwrite-proxy assumption.
- _verify_or_404 checks the capability token before the kill-switch DB read, so
unauthenticated /p spam can't be used as an unbounded settings-DB sink and the
response is identical regardless of the sharing on/off state.
Tests: nested run/checkpoint happy path + wrong-ref rejection, the eviction
fail-closed behavior, and route-level coverage for the rotate / preview-sharing
settings endpoints.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* checkpoint preview endpoint
* harden new preview endpoints
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* address review
* Studio preview: pin adapter, guard streaming submit, robust copy-link
Harden the public per-checkpoint preview surface:
- Pin use_adapter=True in the preview payload sanitizer. Otherwise an
unauthenticated /p caller can POST use_adapter=false, which calls
disable_adapter_layers() on the shared in-memory model without restoring
it; since load_model skips reloads for the same checkpoint, every later
visitor (the page never sends the field) keeps getting base-model output
instead of the fine-tuned checkpoint. Forcing it on also re-enables a
previously disabled adapter and no-ops on merged checkpoints.
- Ignore preview-page submits while a response is streaming. The send
button was disabled but the Enter handler still called requestSubmit(),
so a second request could start before the first reply landed in msgs and
reorder the chat history. Both the keydown and submit handlers now honor
the disabled button.
- Keep the cloudflare-URL polling loop alive across transient startup fetch
errors instead of letting one rejection halt it.
- Build the copy-link from a backend preview_ref (output dir relative to
outputs_root, gated on previewability and the two-segment /p route limit)
so a nested output dir no longer copies a basename-only link that 404s.
Expose preview_ref on training run summaries.
Add route-level security tests (path traversal, payload sanitization,
asset containment, CSP header, HTML title escaping, streaming lock held
until drained) and preview_ref unit tests.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio preview: Safari-safe submit and adapter pin only for LoRA
Follow-ups from cross-browser and route simulations:
- Preview page: send the message from a shared send() helper called by both
the form submit and the Enter key, instead of form.requestSubmit(). The
latter throws on Safari < 16 and older iOS, which broke Enter-to-send there.
Verified across Chromium, Firefox and WebKit with Playwright.
- Only pin use_adapter=True when the resolved checkpoint is a LoRA adapter
(adapter_config.json present); for a merged checkpoint strip it to None.
A merged model has no adapter to toggle, so forcing it on only produced a
per-request "not a PeftModel" warning. The cross-request base-model
contamination fix still holds for LoRA previews.
Add a merged-checkpoint test asserting use_adapter is stripped to None.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio preview: trim verbose comments
Tighten comments across the preview routes, page, checkpoint helpers, and tests
to short single-line notes; drop ones that just restate the code. No behavior
change (verified comment/docstring-only with comment_tools.py check).
* Harden preview routes for PR #6486
- Return a generic 400 detail on a rejected preview path so the public /p
route never echoes the absolute install path (the real reason is logged
server-side instead).
- Strip confirm_tool_calls, session_id and rag_scope in the preview payload
sanitizer so the public surface stays inert regardless of the tool gate.
- Use Path.is_relative_to for the asset containment check, matching the rest
of the codebase.
- Add img-src 'self' and font-src 'self' to the preview page CSP.
- Preview page: on a mid-stream error keep the streamed text, flag the break,
and restore the prompt so the user can retry; drop the unused --font-sans var.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Lee Jackson <130007945+Imagineer99@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
* Add HF dataset streaming mode to Studio
* Added default value for datasetStreaming in training-config-store.ts
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Handle None max_steps for streaming validation
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: fast-fail streaming validation and guard incompatible modes
Reject dataset_streaming at the API boundary when hf_dataset is empty,
the dataset is vision/audio, or max_steps is not set. Probe eval split
with get_dataset_split_names before the streaming load so typos fail
immediately instead of mid-training. Guard column_names=None after map
on iterables. Hide the UI toggle for non-text configurations and clear
the stale flag when config becomes incompatible.
* studio: add streaming dataset tests, iterable helper, and streaming template/format support (WIP)
Work-in-progress on top of feat/studio-dataset-streaming-mode (PR #4946):
- new test_training_streaming.py and iterable.py dataset helper
- streaming support in chat_templates.py and format_conversion.py
- additional streaming guards in trainer.py / models / routes
- frontend streaming wiring in params-section and training-config-store
Committed to preserve uncommitted work before merging latest main.
* studio: fix review-team findings for streaming + main merge
BLOCKER: streaming + raw-text/CPT crashed on len(IterableDataset). Guard it in the
start route (reject format_type=="raw" or training_type=="Continued Pretraining")
and in isStreamingSupported (datasetFormat !== "raw").
Also:
- models/training.py: validate hf_dataset/subset/split (charset+length, block ..//);
cap dataset slice indices (le=1e9); note validator ordering
- chat_templates.py: guard _apply_custom_mapping .map() for streaming
- trainer.py: warn when packing+streaming
- training-config-store.ts: persist-migration bump to v11 (standalone datasetStreaming
backfill); add isVisionModel to NON_PERSISTED; toast on silent streamingCompatiblePatch
mutations in the 4 indirect setters
- tests: route rejections (max_steps, raw/cpt), slice cap, unsafe hf_dataset
* studio: enable raw-text/CPT dataset streaming + streaming UX polish
- raw_text: keep the lazy filter but skip len()-based row counting for
IterableDatasets so raw-text / CPT can stream; guard the eval-size log
- routes/trainer: drop the raw/CPT streaming block; add a defensive
not-streaming guard on the eval auto-split (train_test_split)
- dataset-section: streaming toggle is visible-but-disabled and lists the
exact unmet requirement(s) in its tooltip; block embedding models
- training-start-overlay: show "streaming (no full download)" instead of a
stuck download bar for streaming runs
- trim the streaming test suite to the high-value cases
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: address streaming review (MLX/embedding guards, sliced eval split, rehydrate timing)
- routes: reject dataset_streaming for embedding training and on Apple Silicon
(MLX); both loaders materialize the full dataset instead of streaming
- trainer: validate the base eval split name so streaming eval accepts HF slice
syntax such as "validation[:1000]"
- training-config-store: defer the onRehydrateStorage setState to a microtask so
it doesn't hit the store's TDZ during synchronous hydration
- test: streaming start rejects embedding models
* studio: harden HF dataset streaming (column_names, split slicing, empty/eval bounds, gating)
Address a deeper streaming review:
- raw_text: resolve_column_names() guards IterableDataset.column_names=None
(from_generator / unresolved features) so raw-text and CPT streaming no longer
raise TypeError before training
- models/routes: reject HF slice syntax in train_split/eval_split when streaming
(load_dataset(streaming=True) raises "Bad split"); reject mixed sources
(local/S3) and embedding/MLX streaming at the API, not just in the UI
- trainer: an empty post-slice/filter stream fails preflight with a clear message;
streaming eval is capped (STREAMING_EVAL_MAX_SAMPLES) so each eval terminates;
the manual-slice shortcut falls back to a regular load when train_split is sliced
- format_conversion: streaming conversions preflight the first mapped row so
format errors surface before training, not mid-iteration
- frontend: block streaming on Apple Silicon; clear datasetStreaming when a
dataset is detected as image/audio at start
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: fix CI for streaming PR (lint blocker + no-torch sandbox + preflight test)
- trainer.py: drop unused `IterableDataset` import (hoist safety-net blocker).
- test_training_streaming.py: only select real classes (isinstance type) when
locating the trainer class, so a MagicMock-stubbed global is never passed to
object.__new__ (fixes TypeError on the Python 3.10-3.13 jobs).
- no-torch import sandboxes (test_e2e_no_torch_sandbox.py,
test_studio_import_no_torch.py): teach the chat_templates/format_conversion
exec stubs and the full-import-chain copy list about the new `.iterable`
module so the AFTER/runtime cases import without torch again.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
* Harden model fetching: consent gate for trust_remote_code
Add a load-path consent gate that scans a model's auto_map repository code
before it executes and blocks CRITICAL/HIGH findings unless the user pins
approval of that exact code version. Capability detection stays code-free,
reading raw config.json instead of AutoConfig.
- Scan config.json and tokenizer_config.json auto_map, nested local helpers,
and external owner/name--module repos; fail closed on partial downloads.
- Gate inference, training, and export workers, including the MLX path and a
LoRA's base model, and report requires_trust_remote_code from the raw config
so chat and auto-load surface the dialog.
- Verify trusted-org auto-enable against the Hub with the request token and key
the verdict cache by token; reject local-path and spoofed names.
- Add a consent dialog showing the flagged file, line, and surrounding code.
- Thread hf_token through the scan and load paths for gated repos.
* Address review: token handling, tokenizer/LoRA scan coverage, rollback
- Send the HF token for remote-code scans in the POST body, not the URL, so it
never lands in a log or browser history.
- Collect tokenizer_config.json auto_map files directly instead of relying only
on the repo file listing.
- Resolve a LoRA's base model for the validate flag and the scan endpoint so the
dialog scans the code the workers actually gate.
- Pass the request token to the training YAML trusted-org auto-enable.
- Resend a previously approved fingerprint when rolling back to a custom-code
model after a failed switch.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Consent UX: drop legacy chat toggle, fix decline copy, purge declined downloads
The per-model consent dialog is now the single approval path for custom
(auto_map) code in chat, so three leftovers from before it existed are removed:
- Remove the "Enable custom code" switch from Chat Settings and stop persisting
trust_remote_code, so a previously saved blanket-on cannot linger and load a
model without going through per-version review. The flag stays as an internal
YAML/preset default (e.g. first-party auto-enable); the load path still gates
every custom-code load on a fingerprint only the dialog produces.
- Reword the decline message and the auto-load toast to describe approving the
model's code from the dialog, not a missing settings toggle.
- On decline, purge the repo the scan downloaded so untrusted code is not left
on disk. A new /api/models/discard-remote-code endpoint deletes only a
metadata-only cache entry the scan created; it refuses local paths, loaded
models, and any repo with weight files cached, so a model the user already had
or pre-downloaded is always left untouched. The frontend only calls it when
the scan reported created_by_scan.
Adds discard-endpoint tests (delete metadata-only, refuse on weights/gguf,
refuse local, no-op when not cached) and a created_by_scan payload assertion.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Export: remove the user-facing trust remote code toggle
The Export page kept a "Trust remote code" switch (default on) next to the HF
token field. Like chat, custom (auto_map) code should be approved per model
through the load-time review dialog, not a persistent blanket switch, so the
toggle is removed. The export load path already routes through the same consent
dialog: an HF source now starts with trust_remote_code off and only enables it
when the user approves the scanned code in the dialog (a local checkpoint the
user exported stays trusted by default). With the dialog unreachable and no
approval, an HF source loads with trust_remote_code off, which fails closed
rather than running unreviewed code.
* Block loads of repos with unsafe files using Hugging Face's security scan
The trust_remote_code consent gate covers one load-time RCE vector (a repo's
auto_map Python). It does not cover the other: a malicious pickle inside a weight
file (pytorch_model.bin, *.pkl, *.dat) deserializes during from_pretrained even
with trust_remote_code False, so a repo with a normal config plus a poisoned
pickle slips past the existing gate.
Add a metadata-only malware gate that uses Hugging Face's own scan (picklescan +
ClamAV), read via model_info(securityStatus=True).security_repo_status. It never
downloads, opens, or unpickles the flagged files; it only reads the Hub's verdict
and surfaces the flagged file names. New evaluate_file_security runs
unconditionally (independent of trust_remote_code) in every load path (inference,
training SFT/MLX, export), blocking the load when a file is flagged
unsafe/suspicious/malicious. The /remote-code-scan preflight and the validate
endpoint also report the result so the consent dialog opens as a hard block (no
override) listing the flagged files, even for a repo with no custom code.
Policy: hard block with no user override; fail open when the scan is unavailable
(offline/unscanned) so legitimate loads are not broken; no first-party exemption
(a poisoned pickle in a compromised trusted repo still blocks); local paths and
GGUF are skipped (no Hub scan, non-pickle format). Blocking does not gate on
scansDone, since that is often false for clean repos and a file already flagged
unsafe is unsafe regardless.
Adds test_file_security.py covering the block/allow/fail-open/skip matrix.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Address review: scan list-form tokenizer auto_map, gate unsafe files on all load paths
Fixes from a 10-reviewer pass on the model-fetching hardening:
- The remote-code scanner skipped tokenizer auto_map encoded as a [slow, fast]
list (transformers' standard tokenizer shape, e.g.
{"AutoTokenizer": ["owner/repo--tokenization_x.Slow", null]}). External
tokenizer code in that form was never fetched, scanned, or fingerprinted, so an
AutoTokenizer(trust_remote_code=True) load could run it. _auto_map_refs now
flattens string, list, and nested values. Adds a regression test.
- Compare-mode chat loads and background auto-load only gated on
requires_trust_remote_code, so a repo flagged unsafe by the Hub scan but with no
custom code skipped the hard-block dialog. Both now also gate on
requires_security_review, matching the main chat path.
- The /remote-code-scan and /validate routes collapsed a LoRA adapter to its base
before the malware scan, so unsafe files in the adapter repo itself were missed
in the pre-load review (the workers already scan both). Both routes now run the
file-security scan over the adapter and the base.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Require approval for all HIGH remote code, fail closed when unscannable
Tighten the load-time security gates based on review:
Consent gate
- HIGH-severity auto_map code now requires explicit, per-version approval for
every repo, including first-party unsloth/nvidia. The org is no longer a
blanket bypass: a compromised first-party repo with HIGH code still warrants
review. CRITICAL stays a hard block; clean code still loads after the consent
prompt.
- Fail closed when auto_map code is present but cannot be fully fetched or
listed to scan (gated, offline, transient, or a repo-listing failure that
could hide an imported helper). We cannot fingerprint code we cannot see, so
this is a non-approvable block, retryable once the repo is reachable.
- Scan auto_map from every config that can carry one (model, tokenizer, image
and feature processor, processor, video processor), not just config.json and
tokenizer_config.json, so a custom-processor model is not missed. The file
list is the single source of truth in remote_code_scan and is pinned to the
transformers filename constants by a guard test.
- Distinguish a genuine 404 (config truly absent) from a transient error: only
the latter forces a scan, so a repo with no config is correctly a no-op.
Malware gate
- Scan a remote repo even when its name ends in .gguf; only local paths skip the
Hub scan, so a repo cannot dodge the scan by naming itself "*.gguf".
- Correct the docstring: a file already flagged unsafe blocks regardless of
scansDone; the only fail-open path is an unavailable scan.
Coverage
- Resolve a remote LoRA adapter's base model (not just local directories) so the
base, where the code and weights actually execute, is scanned in validate,
the scan route, and the training and export workers.
- Gate the embedding training path (FastSentenceTransformer) with the malware
and consent checks, matching the other load paths.
Tests updated and added for each change.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Scope malware gate to the load-path vector; stop false-blocking first-party models
Follow-up hardening from a second review pass + a broad live model matrix
(unsloth/* , nvidia/* , third-party, and the eicar malware repo).
Malware / unsafe-file gate
- Scope the block to the actual RCE vector: a root-level file in a code-executing
format. from_pretrained deserializes weight files at the repo ROOT, so a flag is
only a load-path pickle vector there. Two exclusions, because neither is loaded:
inert formats (safetensors is tensor-only, gguf is non-pickle, configs/text/
images) and files in subdirectories. This keeps eicar blocked (its *.pkl/*.dat/
eicar_test_file sit at the repo root) while no longer false-blocking legitimate
first-party repos: nvidia/Nemotron-H-8B-Base-8K ships root safetensors plus NeMo
pickle checkpoints under nemo/ that the loader never touches, and the Hub flags
both; the gate previously hard-blocked it.
- Unknown / future non-"safe" levels now fail closed (block) instead of being
silently allowed, so Hub schema drift cannot introduce a bypass; in-progress
("pending"/"scanning"/"error") levels stay non-blocking to avoid false blocks.
Consent gate
- Ignore a STALE own-repo auto_map target that is absent from the repo listing (an
older config pointing at a file the repo no longer ships) instead of failing the
whole repo closed as unscannable. The present .py are still fully scanned, which
is the stronger coverage, and a file that is not there cannot execute. This
unblocks first-party models like unsloth/PaddleOCR-VL (its tokenizer_config.json
names processing_ppocrvl.py while the repo ships processing_paddleocr_vl.py). A
referenced .py that IS present but cannot be fetched, and a repo-listing failure,
still fail closed.
Remote LoRA base resolution
- Distinguish a genuine 404 (not a LoRA / repo absent -> None) from a transient
error: the transient case is retried once, then logged as a WARNING (a missed
base is scanned by neither gate) rather than silently skipped.
Discard endpoint
- Treat .onnx and .ckpt as weights so a repo whose only heavy artifact is one of
those is never eligible for the declined-download purge.
Tests added for each: load-path scoping (safetensors/subdir/Nemotron-H shapes,
unknown-level fail-closed, pending non-block), stale own-repo auto_map ref, remote
LoRA transient retry, and the empty-config-list (all-404 -> []) semantics.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Make LoRA-base transient-warning test robust to logging backend
Assert on the logger object directly instead of capsys, so the test does not
depend on whether the real structlog logger or the module-stub logger is active
(which varies with test collection order).
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Allow a repo with auto_map but no executable code (e.g. GGUF) instead of blocking
A config can declare an auto_map yet the repo ship NO executable .py -- most
commonly a GGUF repo whose config.json carries an auto_map copied from the original
model (e.g. unsloth/Llama-3_1-Nemotron-Ultra-253B-v1-GGUF references
modeling_decilm.py, which the GGUF-only repo does not contain). A GGUF model loads
through llama.cpp, which never executes auto_map, and transformers cannot run a file
that is not present, so there is nothing to scan and trust_remote_code is a no-op.
The fail-closed change treated this empty result the same as "code is present but we
could not fetch it" and hard-blocked the load. Distinguish the two: repo_remote_code_files
now RAISES RemoteCodeUnscannable when code is present but cannot be fully fetched or
listed (offline / gated / transient / a present .py that 404s / a listing failure),
and returns an empty dict only when the listing succeeded and the repo genuinely ships
no executable .py. The consent gate blocks on the exception (fail closed) and allows the
empty case as a no-op. Real unscannable code still hard-blocks; eicar and CRITICAL/HIGH
custom code are unaffected.
Verified against all 37 unsloth/*Nemotron* models (two GGUF repos were false-blocked,
now load) and the existing matrix (eicar still blocks; DeepSeek-OCR / NVLM-D-72B still
prompt approvable consent). Tests updated to expect the raise for unscannable cases and
added for the no-executable-code no-op.
* Ignore vestigial auto_map in GGUF repos (llama.cpp never runs it)
A GGUF repo's config.json is often copied verbatim from the original
transformers model, auto_map and all, but a GGUF load goes through
llama.cpp which never executes auto_map, so the config is inert. Treat
a direct .gguf reference, and a repo that ships .gguf weights with no
.safetensors, as having no remote code so the consent flow is never
triggered. A mixed repo with both .gguf and .safetensors is still gated,
since the safetensors variant would load through transformers where
auto_map does run. The check sits behind the existing auto_map-present
gate so normal models pay no extra repo listing.
* Add scanner-result copy to the remote-code consent dialog
Make the consent dialog state the scan outcome in plain language for
every model. When the static scan finds nothing, reassure the user with
'Our automatic scanner did not flag any worrying files, but please
double check.' (shown only for the clean, approvable case). When the
scan flags custom code or unsafe files, label the list with 'Our
automatic scanner flagged issues including:'. The Hugging Face
attribution for unsafe files stays in the dialog description.
* Close GGUF-suffix consent bypass for repo ids ending in .gguf
The .gguf short-circuit in _config_has_auto_map skipped the scan for any
model name ending in .gguf, including a bare two-segment repo id like
'evil/model.gguf'. Such a repo can still ship safetensors plus auto_map
Python that transformers would execute, so skipping the scan was an
asymmetric bypass (file_security already scans those repos). Restrict the
short-circuit to genuine direct GGUF file references via
_is_direct_gguf_file_ref: a local .gguf path, or a remote repo_id plus
filename (three or more segments). A two-segment repo id named *.gguf now
falls through to the config scan and _is_gguf_repo file inspection, so it
only skips consent when it actually ships .gguf weights and no safetensors.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Align consent dialog body with the title and fix narrow-width overflow
The scan results (the 'Our automatic scanner...' label, finding/unsafe
cards, and the clean-scan reassurance) sat at the dialog's left padding
while the title and description were indented past the status icon, so
the body did not line up under the description. Move the title,
description and results into one column to the right of the icon so they
share a left edge, and let that column fill its width so the description
no longer wraps early.
Also stop a wide code snippet from pushing the dialog off-screen on
narrow viewports: AlertDialogHeader is a grid with place-items-center,
which sized the content row to its content; give the row w-full so it
fills the track, and add min-w-0 down the results chain so the snippet
scrolls inside its card instead of widening the dialog. Verified aligned
and contained from mobile portrait through ultrawide.
* Treat a repo as GGUF-only only when it ships no transformers weights
_is_gguf_repo excluded only .safetensors, so a repo with a .gguf and a
pytorch_model.bin (or .pt/.pth/.h5/.msgpack/.onnx/.ckpt) and no
safetensors was treated as GGUF-only and skipped the consent scan, even
though transformers can load that weight set and execute the repo's
auto_map code. Require the absence of ANY transformers-loadable weight
before treating the repo as a llama.cpp-only GGUF load. A genuine
GGUF-only repo (only .gguf) is still inert; a mixed repo with any pickle
or safetensors weight is gated. Adds a regression test across all the
non-safetensors weight formats.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Block flagged subdir weight shards referenced by a root index
The malware gate treated every subdirectory file as non-loadable, but
from_pretrained deserializes a subdir shard a root index references
(pytorch_model.bin.index.json -> shards/...-00001-of-00002.bin). Read the
root weight indexes and block a flagged subdir pickle the weight_map
points at; a flagged subdir pickle no index lists (NeMo nemo/*.distcp)
stays non-blocking, and an inconclusive index lookup fails closed.
* Pass hf_token to the export checkpoint load
ExportBackend.load_checkpoint scanned with hf_token in the worker but
loaded the weights unauthenticated, so a gated/private checkpoint passed
preflight then 401'd at from_pretrained. Add hf_token to load_checkpoint
and forward token to every from_pretrained branch; the worker passes the
command's hf_token.
* Scope created_by_scan to every HF cache the discard searches
created_by_scan used get_cache_path (active HF_HUB_CACHE only) while
/discard-remote-code deletes across active, legacy, and default caches. A
repo the user already had in a legacy/default cache was marked
scan-created and deleted on decline. Check all three caches for the repo
dir before declaring the scan created it.
* Scan the full .py closure of external auto_map repos
An auto_map cross-repo ref (owner/name--module.Class) only had its entry
file downloaded, but transformers also fetches that file's relative
imports from the same repo, so a dangerous helper.py was left outside the
scanned fingerprint. List each external repo's .py and scan the whole set
(plus the referenced entry files); fail closed if the repo cannot be
listed or fetched.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fail closed when a weight index cannot be fully read
_indexed_shard_paths treated a partial result as definitive: if one weight
index read cleanly but another failed transiently, it returned the shard
paths it did see. A flagged subdirectory pickle listed only by the index we
could not read would then be classed as "not a load input" and skipped,
re-opening the very fail-open this guard was added to close.
Return None whenever any index read is inconclusive, even if another read
cleanly, so the caller blocks the already-flagged subdir pickle. A repo that
ships no index files raises EntryNotFoundError for each (never inconclusive)
and still returns an empty set.
* Match cached repos case-insensitively in the created_by_scan guard
_repo_in_any_hf_cache resolved casing only against the active cache and then
probed every cache with an exact directory name. A case-variant already
present in a legacy or default cache (models--Unsloth--Foo for a scan of
unsloth/foo) was missed, so the repo was marked created_by_scan and deleted
on decline -- but discard_remote_code_download deletes case-insensitively,
so that delete would hit the user's pre-existing cache entry. Detect
case-insensitively too, mirroring the deletion path.
* Skip remote-code and security review for selected GGUF variants
validate_model ran the trust_remote_code and Hugging Face security-scan
preflight against the repo even when the selected artifact is a .gguf. A
GGUF loads through llama.cpp, which never executes the repo's auto_map
Python and never deserializes root pickle weights, so repo-level Transformers
artifacts (a config.json with auto_map, or an unsafe pytorch_model.bin next
to the .gguf in a mixed repo) are inert for that load. Gating the GGUF on
them is a false positive. Run both preflights only for non-GGUF loads.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Scope the malware gate to actual load roots and serialized files
Two fixes to evaluate_file_security so it neither misses a load-path pickle nor
false-blocks an inert file:
- Honor subdirectory load roots. Spark-TTS / BiCodec call from_pretrained on the
snapshot's LLM subdirectory, so a flagged pickle directly under it is a
root-level load artifact there. A new load_subdirs parameter (set from the
model's audio type via security_load_subdirs) reclassifies those files relative
to the load root and looks for weight indexes under it, so a flagged shard in
that subdir is no longer skipped as "not root-level".
- Exempt source files. A root .py is never deserialized by from_pretrained;
executable repo code runs only through auto_map, which the remote-code consent
gate scans. Flagging a Python helper here would false-block a repo that merely
ships a build or train script.
* Scan a LoRA adapter and base as one consent unit, and gate MEDIUM code
A LoRA load runs both the adapter's and the base's repo code. The consent gate
scanned them separately and pinned one fingerprint per repo, so an adapter that
shipped its own auto_map code was either never shown in the dialog (which only
saw the base) or impossible to approve with the base's fingerprint.
evaluate_remote_code_consent_for_targets now scans all of a load's repos as a
single combined unit and pins ONE fingerprint over the union of their code, so
approving the load approves every repo's code together. evaluate_remote_code_consent
becomes a thin single-target wrapper, and an unscannable target fails the whole
load closed.
Also gate MEDIUM findings: like HIGH they now block pending pinned approval, so a
direct API caller cannot run flagged code by setting trust_remote_code=True
without consenting. Only a clean scan loads without a fingerprint.
* Preflight a LoRA load's adapter and base as one combined consent scan
scan_model_remote_code rewrote a LoRA adapter to its base and scanned only the
base for remote code, so the dialog never surfaced an adapter's own auto_map
code. Scan the adapter and base together through
preflight_remote_code_consent_for_targets, which pins one combined fingerprint
the worker gate accepts. The malware preflight is also scoped to each target's
load subdirectories.
* Apply combined consent and subdir-aware malware scan in load workers
Each load worker (inference, export, training) evaluated remote-code consent
once per target with a single shared fingerprint, so a LoRA adapter that ships
its own auto_map code could not be approved by the base's fingerprint. They now
scan the adapter and base together via evaluate_remote_code_consent_for_targets,
which pins one combined fingerprint over the union of their code. The malware
scan in each worker is also scoped to the model's load subdirectories so a
flagged pickle under a from_pretrained load subdir is not missed.
* Report a consistent trust_remote_code requirement after a model loads
validate_model reports requires_trust_remote_code from the YAML default OR the
raw auto_map, but the load, already-loaded, and status responses reported only
the YAML default. A custom-code model approved and loaded via auto_map was then
reported as not requiring trust_remote_code, so the frontend stored false and a
later retry or rollback sent trust_remote_code=false and failed.
A shared resolver reports the same requirement for a loaded model (a value
stored at load time, else the trust_remote_code the load used, else the YAML
default, else the raw auto_map check), and the load response persists it so the
status and already-loaded paths stay consistent. The selected-GGUF security
review is also scoped to the model's load subdirectories.
* Run the consent gate on training resume and for YAML-only trust_remote_code
Three frontend gaps left a model loading without the trust_remote_code it needs:
- The shared consent helper returned early when the scan found no auto_map and no
unsafe files, dropping a requirement that comes from a model's Studio YAML
default (e.g. GLM-4.7-Flash). It now grants the caller's requirement with an
empty pin instead of sending trust_remote_code=false.
- Resume-from-history called startTraining directly with no consent gate, so a
resumed run whose model needs custom code (or an old run with no approved
fingerprint) hit the worker block with no dialog. It now runs the same gate as
a fresh start.
- HF export passed requiresTrustRemoteCode=false for every HF source, so a
YAML-only model could not flip the flag before export. It now signals the
requirement for HF sources.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Cover both LoRA repos in validate, report GGUF as inert, purge all declined repos
Three follow-on gaps from the combined adapter+base consent work:
- validate_model resolved requires_trust_remote_code from the base alone, so a
LoRA adapter that ships its OWN auto_map code (with a plain base) was reported
as not needing trust_remote_code and the consent dialog never opened. It now
checks the [adapter, base] target set, matching the scan route and the workers
(which already gate both) and the security review already running over both.
- The already-loaded, loaded, and status responses for a selected GGUF reported
requires_trust_remote_code from the model's YAML default. A GGUF loads through
llama.cpp, which never executes the repo's auto_map Python, so the requirement
is inert for that load. They now report False, matching validate_model (which
already skips both gates for GGUF) so a status refresh cannot flip the flag
back on.
- The remote-code scan downloads both the adapter's and the base's config, but
created_by_scan tracked only the primary, so a base the scan was first to pull
into the cache was left on disk when the user declined. The scan now reports
scan_created_repos (every repo it newly cached) and the decline cleanup purges
each; created_by_scan stays for older clients. The frontend falls back to the
primary flag when the list is absent.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Scan the repo the load fetches, purge external code on decline, harden consent pins
Six follow-on hardening fixes from a fresh review pass over the gate:
- The malware gate scanned the literal "Spark-TTS-0.5B/LLM" alias, but the trainer
downloads it as unsloth/Spark-TTS-0.5B and loads LLM/, so the alias 404'd and
failed open, missing a flagged LLM/ pickle. evaluate_file_security now resolves
the alias to the repo the loader fetches and scans LLM/ as a load root.
- security_load_subdirs relied only on tokenizer detection, which fails on an
unresolved alias or offline; it now also honors the Studio YAML audio_type
default, so a BiCodec LLM/ load root is not missed.
- The remote-code scan downloads external auto_map repos (owner/name--module.Class),
but the decline cleanup tracked only the model/adapter/base, leaving the external
untrusted code cached. The scan now enumerates external auto_map repos and reports
the ones it created in scan_created_repos, so a decline purges them too.
- External auto_map refs failed the whole load closed on a stale or mis-derived
dotted ref (sub.mod.py vs the real sub/mod.py) even though the actual file was
present and scanned. They now drop such refs when the repo listing is real, exactly
like the own-repo path; an empty/incomplete listing still fetches and fails closed.
- The combined consent fingerprint keyed code by the raw target string, so the scan
endpoint's canonicalized casing and a worker's raw user input produced different
pins for identical code, rejecting a valid approval. Hub repo ids are now folded to
lowercase in the key (local paths stay case-sensitive), so the pin tracks the code.
- Export threaded hf_token into the weight load but not into detect_audio_type /
is_vision_model, so a gated multimodal base 404'd in detection and fell through to
the text loader. Both probes now use the same token.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Thread the token through check-vision and guard the gate's parallel sites
The /check-vision endpoint classified a model without the hf_token, so a gated or
private vision model 404'd in the probe and was reported as a plain text model --
the same dropped-token shape as the export probes, at a sibling site. It now passes
the token like the neighboring /check-embedding endpoint.
Add deterministic consistency guards (tests/test_security_gate_consistency.py) that
enumerate the gate's parallel sites mechanically instead of relying on a review to
spot a missed sibling: every is_vision_model / is_embedding_model / detect_audio_type
caller under routes/ and core/ must thread the token, every GGUF response must report
trust_remote_code via the resolver or False (never the raw YAML default), and every
load worker that runs the malware or consent gate must resolve the LoRA base. A new
site that drops the token or mis-reports the requirement now fails CI directly.
* Narrow the LLM alias rewrite and make audio detection token-aware
Three fixes from the confirmatory review, one a regression from the previous round:
- _load_scan_target rewrote EVERY remote repo ending in "/LLM" to unsloth/<parent>,
so a real third-party repo named "<owner>/LLM" was scanned as unsloth/<owner>
while the loader still fetched the real repo -- a fail-open hole introduced when
the Spark-TTS alias handling was added. It now rewrites only a registry-known
bicodec alias; every other "/LLM" repo is scanned as itself.
- detect_audio_type cached results under the bare model name, so an unauthenticated
probe of a gated/private repo cached None and poisoned a later authenticated call
with the token. The cache is now keyed by (normalized_name, token_fingerprint),
matching the vision cache.
- The training fallback /check-vision call dropped the hf_token, misclassifying a
gated/private VLM when the config endpoint failed. It now passes the token, like
the getModelConfig call it falls back from; checkEmbeddingModel takes the token too.
Extend the consistency guards: every capability cache must be keyed by a tuple
including the token, so a cache re-declared as Dict[str, ...] fails CI.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Document the broad .py scan as deliberate and enforce it with a test
The remote-code scanner scans every .py in a repo once an auto_map exists, not
just the auto_map entry's static import closure. This is intentional: the entry
module can reach a sibling via an absolute import, importlib, or exec, none of
which a static relative-import closure follows, so closure-only scanning would be
a real bypass of a load-time RCE gate. The broad scan never under-scans; the cost
is that an unrelated benign script can over-block, which is the safe failure
direction (HIGH stays approvable; only CRITICAL hard-blocks).
Spell this out at both the local and remote scan sites so the choice reads as
deliberate, and add a test asserting an unrelated, never-imported .py is still
scanned -- so a future narrowing to the static closure fails CI.
* Purge a declined remote LoRA adapter the scan downloaded
scan_model_remote_code probed the created-by-scan state AFTER resolving the base,
but get_base_model_from_lora_identifier downloads a remote adapter's own
adapter_config.json, so the adapter looked already-cached and was dropped from
scan_created_repos. On decline the adapter -- including the auto_map .py the
preflight fetched -- was left on disk, defeating the "untrusted code is not left
on disk" guarantee for the adapter itself.
Snapshot the primary's cache state BEFORE base resolution and use it when marking
the adapter scan-created; on any probe error treat it as pre-existing so a decline
never deletes it. The base and external repos are unaffected (their configs are not
downloaded before their own probe). Add a test that models the mid-scan download
side effect, which the prior static-stub tests did not.
* Clear remote-code approval when the training model changes
Switching the training model from an approved custom-code model to a clean one
kept the previous model's trust_remote_code=true and approved fingerprint in the
store: setSelectedModel reset visionImageSize on a true switch but not the
remote-code approval. The clean model then trained with trust_remote_code=true,
which bypasses the compiler and disables fused cross-entropy.
Reset trustRemoteCode and approvedRemoteCodeFingerprint on a true model switch.
The new model's own YAML default is re-applied by loadAndApplyModelDefaults, and a
custom-code model still re-opens the consent dialog before training starts, so the
only change is that a clean model no longer inherits a stale approval.
* Trim verbose comments across the model-fetching hardening changes
Condense the explanatory comments and docstrings introduced across the
trust_remote_code consent gate, the malware/unsafe-file gate, the remote-code
scanner, the load workers, the model routes, and the security frontend into
fewer, tighter lines while preserving every security rationale (fail-open vs
fail-closed direction, the deliberate broad-scan anti-bypass note, the
empty-vs-unscannable distinction, stale-ref handling, and the alias-rewrite
spoof guard).
Comments and docstrings only. No code, logic, identifiers, or test behaviour
changed; verified comment-only via the AST/TypeScript checker (40/40), with the
backend test suite and frontend tsc green.
* Do not cache transient audio-detection failures
detect_audio_type cached _detect_audio_from_tokenizer's result
unconditionally, so a transient read failure (network error or 5xx,
returned as None) poisoned the cache and the later successful probe never
ran. Mirror the vision cache: _detect_audio_from_tokenizer now returns
(audio_type, definitive) and the caller caches only definitive results.
A read that succeeds with no audio tokens, or clean 404s for every
tokenizer path, stays a cacheable None; only a genuine transient failure
(connection error, timeout, 5xx, malformed body) skips the cache so the
next call retries.
---------
Co-authored-by: danielhanchen <michaelhan2050@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* Expose MLX grad value clipping in Studio
* update test
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* dataset ordering + wd
* fix mlx smoke step expectations
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* cast norm activation output back to original input dtype
* address mlx studio review feedback
* Fix present-but-None seed override for PR #5656
studio/backend/core/training/worker.py
`config.get("model_random_state", random_seed)` only fills the
default when the key is absent. When a caller passes
`config["model_random_state"] = None` explicitly (which happens
any time a JSON payload sends an explicit `null`), the old code
forwarded `None` to FastMLXModel and disabled deterministic init
silently. Same for `lora_random_state`. Treat absent and explicit
None the same way: fall back to random_seed.
studio/backend/tests/test_training_raw_support.py
Update the source-string assertions to match the new lines.
* Guard optional MLXTrainingConfig fields and normalize random_seed for PR #5656
The MLX worker now passes `cast_norm_output_to_input_dtype` and
`dataset_order` only when the linked unsloth-zoo dataclass actually
declares them. Released zoo trees that predate the paired PR can still
construct `MLXTrainingConfig` without raising
`TypeError: unexpected keyword argument`. Once the dependency floor is
bumped to a release that contains both fields, the feature-detect
guards become no-ops.
`random_seed = config.get("random_seed", 3407)` was unguarded against
explicit `None` from raw / backend callers. The same value seeded the
trainer and was the fallback target for `model_random_state` /
`lora_random_state`. Normalize once at the top of the function and use
the normalized value everywhere so an explicit `None` cannot reach
FastMLXModel / get_peft_model / MLXTrainingConfig.
Existing seed source-pattern test updated to match the new normalize
helper. New test asserts the feature-detection guards exist and that
the unconditional kwargs do not include the gated fields.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Normalize seed / cast / max_grad_value at TrainingBackend for PR #5656
Round-3 review consensus: the per-field guards that landed in the MLX
worker only protect the MLX path. The same `TrainingBackend.start_training`
config still reaches the CUDA/text trainer at `worker.py:2267`, the
embedding LoRA init at `worker.py:2450`, and embedding TrainingArguments
at `worker.py:2624` with raw `None` values, so an explicit
`random_seed=None` from a raw / backend caller still breaks non-MLX
training even after the previous fix.
Move the normalization into `TrainingBackend.start_training` itself,
where it runs once for every training mode:
- `_coerce_seed(value)`: explicit `None`, non-int, or absent all become
3407. Every downstream worker now sees an int.
- `_coerce_optional_bool(value, default)`: explicit `None` falls back
to `default` instead of `bool(None) == False`. Also normalizes the
common raw-config / YAML string aliases ("true" / "false" / "0" /
"1"). Used for `cast_norm_output_to_input_dtype`.
- `_coerce_optional_nonneg_float(name, value)`: rejects negative
numerics from raw / backend callers, matching the Pydantic
`ge=0` constraint the HTTP route already enforces. Used for
`max_grad_value`.
worker.py MLX path: the existing `bool(config.get(key, True))` for
`cast_norm_output_to_input_dtype` was changed to also fall back on
explicit `None`, so direct worker callers (bypassing
`TrainingBackend.start_training`) are equally safe. `max_grad_value`
also raises on negative values inside the worker for the same reason.
TrainingStartRequest.random_seed default bumped from 42 to 3407 so
direct REST callers that omit the field receive the same default as
the Studio frontend and the MLX worker.
New regression test exercises the three new helpers across explicit
None, valid values, string aliases, and negative-value rejection.
* Tighten feature-detect test paren tracking for PR #5656
The block-extraction used , which stops at the
first inner closing paren (e.g. )
and would silently miss a future unconditional
/ added later in the same dict literal. Switched to
proper paren-depth tracking so the unconditional block is checked end-to-end.
* Shorten verbose comments in MLX Studio backend
* Handle MLX Studio EOS appending by mode
* Wire MLX leaf norm clipping through Studio
* Respect VLM layer filters for explicit LoRA targets
Rationale / guardrails for the local Studio/vision push:
When callers provide explicit VLM LoRA target_modules together with layer filters, FastVisionModel still needs to route the explicit targets through get_peft_regex. Otherwise the layer filters are ignored and adapters can be attached outside the requested language/vision scope.
Do not revert this to plain list(target_modules) for explicit module lists. The CUDA/Studio-facing contract is that explicit targets and layer filters compose: target_modules selects module names, while finetune_language_layers / finetune_vision_layers / finetune_attention_modules / finetune_mlp_modules constrain where those targets are allowed.
The regression test covers the language-only explicit q_proj case and source-checks that explicit targets are wrapped through get_peft_regex when filters are active.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Refresh MLX smoke clip-config note for leaf_norm default
Trim the 11-line comment block to 5 lines and correct the stale claim
that MLXTrainingConfig defaults to max_grad_value=1.0. The new default
is max_grad_leaf_norm=1.0 (same memory profile as elementwise but
direction-preserving). The smoke still pins max_grad_value=1.0
explicitly to keep the 13-seed pass-rate fixture stable.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Forward max_grad_leaf_norm through the training route and warn when layer filters constrain explicit target_modules for PR #5656
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han-Chen <info@unsloth.ai>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
* feat(studio): add S3 dataset configuration foundation (#4539)
Add foundational types and configuration for S3 bucket dataset loading:
- Add S3Config type to frontend training types
- Add S3Config Pydantic model to backend training models
- Add "s3" as a DatasetSource option
- Add s3Config state and setS3Config action to training config store
- Add i18n translations for S3 configuration (English and Chinese)
This provides the type definitions and UI text for S3 integration.
Full implementation requires boto3 dependency and data loading logic.
Refs: #4539
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Wire S3 config into training pipeline and prevent secrets persistence
- Pass s3_config from request into training_kwargs so it flows to training subprocess
- Add s3Config to NON_PERSISTED_STATE_KEYS to prevent AWS secrets from being
saved to localStorage
Addresses code review feedback on PR #5951.
* Exclude S3 config from database persistence to protect secrets
Filter out s3_config (which contains secret_access_key) from the
config_json stored in training_runs table, preventing AWS credentials
from being persisted to disk.
Addresses P1 security feedback on PR #5951.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Re-raise HTTPException in start_training and defer s3 DatasetSource widening for PR #5951
* Redact s3_config from W&B run config and accept camelCase S3 credential aliases for PR #5951
* feat(studio): implement S3 dataset loading end-to-end
Builds the actual S3 loader on top of the hardened #5951 foundation,
turning the 501-gated scaffold into a working dataset source.
Backend:
- Add core/training/s3_dataset.py: lists and downloads supported dataset
files (parquet/json/jsonl/csv) from an S3 bucket to a temp dir, using
IAM-role or access-key credentials. boto3 is imported lazily (optional dep).
- Wire s3_config into UnslothTrainer.load_and_format_dataset (downloads then
reuses the existing local-file path) and thread it through worker.py.
- Replace the 501 "not implemented" gate with a boto3-availability guard so
S3 works when boto3 is present and fails clearly when it is not.
- Add boto3 to studio.txt requirements.
- Add tests/test_s3_dataset.py (8 tests) covering download/filtering,
collisions, missing-boto3, and S3Config camelCase/IAM validation.
Frontend:
- Widen DatasetSource to include "s3"; add s3_config to the training payload
type and mapper; add an S3 validation branch and selectS3Source store action.
- Add s3-config-form.tsx (bucket/region/prefix/keys/IAM toggle) reusing the
existing studio.dataset.s3.* i18n strings.
- Add a Hugging Face / Local / Amazon S3 source toggle in dataset-section;
the S3 config card replaces the dataset combobox when S3 is selected.
- Fix DatasetPreviewDialog to accept the widened DatasetSource type.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix S3 dataset loader for PR #6222
* Fix S3 dataset edge cases for PR #6222
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix S3 IAM payload handling for PR #6222
* Block multimodal S3 datasets for PR #6222
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: Ash <ash@MacBook-Pro.local>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: wasimysaid <wasimysdev@gmail.com>
Trim and tighten code comments and docstrings across the repository. Comment-only: every changed file verified code-identical to main via AST/token comparison.
Trim and tighten code comments and docstrings across studio/ Python. Comment-only: every changed file verified code-identical to main via AST/token comparison.
Raise ruff line-length to 100 and extend the local pre-commit format pipeline (def-signature magic-comma normalization, short multi-line assert collapse, kwarg '=' spacing, blank-line-after-short-import removal, adjacent string-literal / f-string+plain merge, redundant-pass pruning). Every transform re-checks the file AST and is dropped if it would differ; the whole-repo reformat is verified AST-identical per file and idempotent.
* Studio: add VLM image-size control for training
Studio vision fine-tuning had no explicit way to cap image resolution, so
users could not trade visual detail against context and memory use from the
training UI, YAML config, or API payload. :) Add a nullable `vision_image_size`
setting that keeps the current model default when unset and applies a
max-side resize when provided.
- Add `vision_image_size` to the training request model, route payload, backend
training config, and frontend API/types plumbing.
- Validate the value server-side as either null or an integer in the supported
256-2048 range.
- Surface an Image Size selector for vision LoRA training with Default plus
common preset sizes.
- Include the value in training start payloads only for image-dataset vision
models, and serialize it into vision-aware YAML configs.
- Map backend model defaults back into the training store and reset the value
when reapplying model defaults.
- Pass the resize through the Torch trainer via `UnslothVisionDataCollator`
using max-dimension semantics.
- Apply the same max-dimension resize in the MLX VLM path before mlx-vlm's
internal collation, preserving aspect ratio and avoiding upscaling.
- Add backend validation coverage and MLX resize-size tests for the new
behavior.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: thread vision_image_size into DeepSeek OCR + writable MLX ndarray
- trainer.py: DeepSeek OCR collator now honors the new vision_image_size
setting as image_size. Falls back to 640 when null. base_size stays at
1024 and crop_mode stays True so the Gundam preset's dynamic cropping
of large documents keeps working.
- worker.py: _resize_mlx_vlm_image returns np.array(image, copy=True)
instead of np.asarray(image). The PIL view from np.asarray is not
writable, which makes HF VLM processors emit "The given NumPy array
is not writable, and PyTorch does not support non-writable tensors..."
when they call torch.from_numpy. copy=True keeps the same shape and
dtype but produces a writable buffer.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: align YAML export gate with API mapper + extend Image Size dropdown
- training-section.tsx: handleSaveConfig now passes
isVisionModel && isDatasetImage === true to serializeConfigToYaml,
matching buildTrainingStartPayload. Stops vision_image_size from
leaking into exported YAML for text-only datasets where the API
would have sent null.
- params-section.tsx: add 256 to visionImageSizePresets so the
dropdown spans the validator's full [256, 2048] range. Also render
a synthetic SelectItem for the current value when it was loaded
from YAML or model defaults and is not in the preset list, so the
controlled Select always shows the active size.
* Studio: validate vision_image_size in YAML/model-default loader
mapBackendModelConfigToTrainingPatch now mirrors the backend validator
at studio/backend/models/training.py:169 by dropping any value that is
not an integer in [256, 2048]. Pre-fix, an imported YAML like
vision_image_size: 4096 or 640.5 would land in the store and the UI
would happily display it, only to fail when Start Training posted to
the backend. With this guard the store never holds a value the backend
would reject.
* Studio: precise error messages for invalid vision_image_size inputs
Switch the field_validator to mode="before" so True/False surface as
bool (not Pydantic's coerced 1/0) and give a precise
"must be an integer or null" message instead of the misleading
"must be in [256, 2048] (got 1)". Also explicitly accepts numpy
Integral and integral Real scalars so YAML or programmatic callers
using numpy ints keep working.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: test that bool inputs yield the precise 'integer or null' error
Regression guard for the validator switch to mode="before". Pre-fix,
vision_image_size: True was rejected with "must be in [256, 2048]
(got 1)" because Pydantic coerced before our check ran. New test
asserts the message now reads "integer or null".
* Studio: tighten vision_image_size loader + YAML save + MLX rounding
Round 2 of follow-up review surfaced three usability issues:
- model-defaults.ts: switching to a model whose backend YAML omits
vision_image_size now explicitly resets the store value to null.
Pre-fix, a stale 2048 from a previous model would silently apply
to the new run because every checked-in model-default file omits
the key.
- training-section.tsx: handleSaveConfig now includes vision fields
unless isDatasetImage is definitively false. isDatasetImage is null
during dataset checks, after dataset edits, and on import; treating
unknown as "drop" would silently lose the user's selection in those
windows. Confirmed-text-only datasets still drop the value.
- worker.py: _mlx_vlm_max_resized_size now mirrors the Torch collator's
integer formula (w * size + size_func // 2) // size_func instead of
Python round(), which uses banker's rounding and disagreed by 1px on
half-pixel inputs like 333x1000 with target 500 (was 166, now 167).
Test_mlx_training_worker_config gains parity assertions.
* Studio: reset vision_image_size in the model-config error fallback path
mapBackendModelConfigToTrainingPatch resets stale image size on the
success path, but if the /api/models/config endpoint throws,
training-config-store.ts falls through to checkVisionModel and only
updates capability flags. Pre-fix that left a stale 2048 (or any
prior selection) in the store, so once dataset detection marked the
new dataset as image, the next training start would silently apply
the previous model's size. The error branch now also resets to the
DEFAULT_HYPERPARAMS.visionImageSize sentinel.
* Studio: revert DeepSeek OCR Image Size knob + move missing-key reset
Round 3 of the parallel-reviewer pass surfaced two issues that I had
introduced earlier in this PR's follow-ups.
- trainer.py: my prior change threaded vision_image_size into the
DeepSeek OCR collator's image_size argument. The collator's
(image_size, base_size, crop_mode) is a single preset
(Tiny / Small / Base / Large / Gundam); changing image_size in
isolation desynchronizes the per-crop pixel grid from num_queries
downstream and produces wrong token grids on documents larger than
the per-crop tile. The fix pins the collator back at the Gundam
preset and logs a clear "ignored for DeepSeek OCR" notice when the
user has selected a non-default Image Size.
- model-defaults.ts + training-config-store.ts: the round 4 fix that
reset visionImageSize when a model YAML omitted the key also fired
on same-model reloads (ensureModelDefaultsLoaded re-fires on page
refresh), wiping a value the user had just selected. The reset is
now in setSelectedModel, gated on selectedModel != previousModel,
so true model switches still clear stale values while reloads keep
the user's selection.
* Studio: extend DeepSeek OCR Image Size exclusion to MLX + frontend
Round 4 of the parallel-reviewer pass flagged that the Torch trainer
exclusion I added did not have a matching MLX guard, and that the UI
still offered the dropdown for DeepSeek OCR even though the backend
ignores it.
- worker.py: _run_mlx_training now mirrors the Torch exclusion. When
the model name matches DeepSeek OCR, vision_image_size is forced
back to None before _adapt_for_mlx_vlm sees it, so dataset images
pass through unchanged just like the Torch path. Emits a clear
status line when this happens.
- params-section.tsx: the Image Size Row is now gated on
showVisionImageSize (showVisionLora && !isDeepseekOcr) instead of
showVisionLora alone, so DeepSeek OCR users no longer see a control
that silently has no effect.
- mappers.ts: buildTrainingStartPayload sends null for vision_image_size
whenever the selected model is DeepSeek OCR, so the backend log line
about ignoring the value never fires from a UI-driven start.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Studio: tighten YAML import/save for vision_image_size
Two YAML-path asymmetries that could leak a stale image size into
training:
- parseYamlConfig now treats a missing training.vision_image_size as
null. Without this, importing a YAML saved before this feature (or
any config that omits the key) preserved whatever value the user had
previously set on a different model. The model-defaults reload path
still uses Object.hasOwn so same-model defaults reloads do not wipe
a manual selection; only file import normalises the missing key.
- handleSaveConfig now passes a DeepSeek-OCR-specific guard to
serializeConfigToYaml so saved YAML matches what the API mapper
actually sends. Previously a state with visionImageSize set could
emit the key even though Studio ignored it at training time for
DeepSeek OCR, and a later import for a non-DeepSeek vision model
would activate the stale value.
serializeConfigToYaml gains an optional third parameter
includeVisionImageSize defaulting to includeVisionFields, preserving
the existing 2-arg call signature for backwards compatibility.
* Studio: also reset vision_image_size when YAML lacks a training section
Round 9's parseYamlConfig normalization only fired when the YAML had a
training mapping that omitted vision_image_size. A lora-only or
logging-only YAML (or one with `training: null`) still left trainingObj
unset, the mapper saw no vision_image_size key, and the previously
selected store value persisted into the next training run.
Now an absent or null training section is synthesised as
{ vision_image_size: null } so model-defaults.ts always patches
visionImageSize back to Default on file import. Same-model defaults
reloads still preserve manual choices via the existing Object.hasOwn
gate in mapBackendModelConfigToTrainingPatch.
* Studio: unify parseYamlConfig non-object training handling
A fresh static review (Opus subagent) flagged P3-1: parseYamlConfig
only synthesised vision_image_size: null when raw.training was either
absent or a plain object missing the key. If raw.training is a scalar
or an array (malformed but still parseable), the value was passed
through unchanged, the mapper's Object.hasOwn returned false, and any
previously selected visionImageSize persisted - the same stale-state
leak the lora-only fallback was added to close.
Treat any non-plain-object raw.training (null, array, scalar) as a
malformed/missing section and reset to { vision_image_size: null }.
* Studio: tighten code comments for vision_image_size path
* Studio: tighten vision_image_size validator + restore lost comment context
Two issues surfaced by a fresh adversarial review of the validator:
1. v.strip().lstrip("+-").isdigit() let "++512" / "--256" / "+-+512"
slip past the gate, then int("++512") raised an uncaught ValueError
and Pydantic surfaced "invalid literal for int() with base 10: '++512'"
instead of the contracted "vision_image_size must be an integer or null".
2. str.isdigit() returns True for Unicode digit families (full-width '５１２',
Arabic-Indic '٥١٢', Devanagari '१०२४'), and int() coerces them, so the
value reaching the backend wasn't the ASCII the user typed.
Replaced the lstrip+isdigit pair with re.fullmatch(r'[+-]?[0-9]+', stripped),
which rejects both shapes with the precise error and accepts the documented
ones ('256', '+512', ' 1024 '). Added 8 regression test cases covering
multi-sign strings, lone sign, and the three Unicode digit families.
Also restored comment context lost in f9c39331:
- model-defaults.ts: name studio/backend/models/training.py:_check_vision_image_size
as the spec the [256, 2048] range mirrors, so a maintainer changing the
cap in one file can find the other.
- training-section.tsx: enumerate the three windows in which isDatasetImage
is null (before a check, after dataset edits, on import) so a future
maintainer doesn't simplify the gate to `isCheckingDataset`.
- worker.py: qualify the writable-ndarray comment with "when a resize is
requested" so it doesn't misadvertise the resize=None early-return.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
* studio: drop unused max_grad_value schema + route plumbing
The MLX worker hardcodes max_grad_value to 5.0 after PR #5340. The
schema field, frontend payload type, route forwarder, and start_training
kwarg threading were all left in place as a transitional buffer for old
clients. The field is now genuinely unused everywhere except inside the
MLX worker, so the schema, route forwarder, and config-build entries can
go. Pydantic still tolerates older clients that send max_grad_value
because TrainingStartRequest's model_config defaults to extra=ignore.
* [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>
* mlx fixes
* Fix studio integration, local dataset files, chat templates without the torch gpu imports
* pass grad norm in mlx worker
* fix(studio): pass MLX grad clipping settings
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* mlx: update grad value
* fix(mlx): address ci and clipping review
* fix backward compatibility and CI tests
* unsloth local is mlx function
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* dont reference runtime
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio mlx: hardcode value clipping, drop max_grad_value from frontend
Simplifies the MLX grad-clipping plumbing now that we are standardising on
elementwise value clipping at [-5, 5] for the compiled MLX path and norm
clipping disabled. The MLX worker no longer reads max_grad_norm /
max_grad_value from the request; both are pinned in one place. Frontend
stops sending the field at all, and the TypeScript request type drops it
to match. Non-MLX (CUDA/AMD/Intel) is untouched and continues to pick up
HF TrainingArguments' default max_grad_norm = 1.0.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
* studio: allow huggingface.co and datasets-server.huggingface.co in CSP connect-src
The security hardening pass (0881a7a5) added connect-src 'self', which
blocked the Training page's direct browser calls to HuggingFace. Model
search (@huggingface/hub listModels/modelInfo/whoAmI -> huggingface.co)
and dataset subset/split discovery (datasets-server.huggingface.co/splits)
both returned nothing as a result.
Extend connect-src to permit the two HF hosts the SPA actually talks to.
No other directive changes; HF tokens still stay client-side.
* studio: format FastAPI 422 detail arrays in training error messages
readError in train-api.ts stringified payload.detail directly. On a 422
the detail is an array of {loc, msg} objects, which JS coerces to
'[object Object],[object Object]' -- the UI showed that instead of the
actual validator message.
Format the array into 'field.path: msg; ...' so the offending field and
the validator's message surface in the UI and toast.
* studio: allow num_epochs/max_steps = 0 sentinel through TrainingStartRequest
The hyperparameter validators added in the security pass rejected 0 for
both num_epochs and max_steps. But Studio's steps-vs-epochs toggle uses
0 as a sentinel: when training by max_steps the frontend sends
num_epochs=0, and when training by epochs it sends max_steps=0. The
trainer expects this and ignores the zeroed field.
Widen both validators to [0, MAX]. They still catch the actual
out-of-range and non-integer inputs they were added for.
* studio: reject TrainingStartRequest when num_epochs and max_steps are both 0
Each field's validator accepts 0 as a "use the other one" sentinel, but
on their own they don't catch the case where both are 0 (or max_steps
is None and num_epochs is 0). That payload would otherwise produce a
no-op training job. Add a model-level validator that rejects it with a
clear 422 message.
* studio: add Optional[int] type hints to _check_max_steps and _check_warmup_steps
Brings these two validators in line with the rest of the TrainingStartRequest
validators in the same file, which all carry explicit cls/v/return hints.
* studio: contain export and dataset paths under their configured roots
resolve_under_root and resolve_dataset_path previously returned absolute
paths unchanged, so an authenticated client could supply
save_directory="/tmp/escape" (or any other absolute path) and have the
exporter drop adapter files anywhere the server user could write. This
turned up during a recent audit pass where an authenticated POST to
/api/export/export/lora with save_directory="/tmp/lora_escape_test"
returned 200 and wrote adapter_model.safetensors, adapter_config.json,
and tokenizer files under /tmp.
The fix is two-layered:
storage_roots.py adds an _assert_contained(resolved, root) helper that
runs after path resolution and rejects any result whose realpath does
not sit under realpath(root). resolve_under_root now rejects '..'
segments and null bytes outright, and only accepts absolute inputs when
they are already inside the configured root (internal call sites that
re-resolve a stored absolute path stay idempotent;
worker.py:resolve_output_dir(output_dir) etc. continue to work).
resolve_dataset_path picks up the same containment rule, scoped to the
three dataset roots.
models/export.py adds field_validator("save_directory", mode="before")
to ExportCommonOptions and ExportGGUFRequest so bad input fails fast at
422 with a clear message rather than a 500 deep inside the resolver.
The validator rejects empty/whitespace, null bytes, control chars,
strings longer than 255 chars, absolute paths, and '..' segments.
routes/export.py:_export_details now returns os.path.relpath(output_path,
exports_root()) so the Export Complete dialog and /api/models/loras no
longer leak the absolute install prefix to the UI; the basename is
used as a last-resort fallback.
Verified end to end:
- POST /api/export/export/lora {"save_directory":"/tmp/foo"} -> 422
"save_directory must be a name or relative path under the export
root; absolute paths are rejected". /tmp/foo is not created.
- "../../etc/escape" -> 422 "may not contain '..' segments".
- save_directory="my_subdir" -> still accepted (400 only because the
test had no checkpoint loaded yet, not because of validation).
- Internal idempotent re-resolve via resolve_export_dir(absolute path
that is already under exports_root) returns the same path unchanged.
* studio/sandbox: harden bash + python tool execution
The sandboxed Bash and Python tool channels in Chat ran with a thin
preexec hook (PR_SET_NO_NEW_PRIVS + RLIMIT_FSIZE only). Bash had a
small word blocklist; Python had an AST safety pass aimed at
signal-tampering and shell-escape primitives. An audit pass showed
several gaps that a tool-calling model could trigger inadvertently:
- bash curl/wget/nc reached AWS IMDSv2 and returned live STS
credentials for the instance role.
- python "import socket; s.connect((169.254.169.254, 80))"
reached the same endpoint regardless of the bash blocklist.
- "cat /etc/passwd" was blocked at the bash side (because "passwd"
is in the blocklist), but "open('/etc/passwd').read()" in Python
happily returned its contents.
- "chr(115)+chr(117)+chr(100)+chr(111)" style dynamic-arg
construction slipped through the AST shell-escape check.
- The supervisor used proc.kill() on timeout, which only signals
the immediate pid; bash-backgrounded children survived. A fork
bomb could spawn for the full 300s timeout window.
- Session work directories under ~/studio_sandbox/<id>/ were
created with default umask (0o755), so any other UID on the host
could enumerate them.
- session_id sanitisation used a one-shot str.replace("..",""),
which is non-iterative and a small footgun.
This commit takes a conservative middle path: the sandbox still
runs as the Studio UID with no namespace tricks where the kernel
disallows them, but every chokepoint is tightened.
_sandbox_preexec now:
- calls os.setsid() so children share a process group; the
supervisor uses os.killpg(SIGKILL) on timeout/cancel so
backgrounded children die with the parent (new _kill_process_tree
helper, wired into _cancel_watcher and both _bash_exec /
_python_exec timeout branches).
- calls os.umask(0o077) so files the child writes default to 0o600.
- applies PR_SET_PDEATHSIG=SIGKILL so an orphaned child dies if
Studio exits.
- best-effort unshare(CLONE_NEWNET) for a private network namespace
(failure is logged and swallowed; defense-in-depth is still in
place via the bash blocklist and the AST checker below).
- sets RLIMIT_NPROC=10000 (tunable via UNSLOTH_STUDIO_SANDBOX_NPROC),
RLIMIT_AS=8GB, RLIMIT_CPU=300, RLIMIT_NOFILE=1024. The 10k NPROC
figure is chosen to sit well above the ~500 LWPs a healthy Studio
+ llama-server combination already uses while still capping a
runaway fork bomb. NPROC counts LWPs per real UID, so a lower
figure (e.g. 256) starves legitimate bash forks
("bash: fork: retry: Resource temporarily unavailable").
_get_workdir:
- rejects session_id that doesn't match [A-Za-z0-9_-]{1,64};
non-matching values bucket into a shared "_invalid" dir.
- chmod 0o700 on both the workdir and on ~/studio_sandbox/ so
other UIDs cannot read another session's contents.
_BLOCKED_COMMANDS_COMMON gains: doas, pkexec, halt, poweroff, curl,
wget, nc, ncat, netcat, socat, ssh, scp, sftp, rsync, eval, source.
The intent is to keep general bash usage working (echo, ls, pipes,
loops, for, head, etc.) while denying the obvious egress and
escalation paths.
The AST checker (_check_signal_escape_patterns) is split into the
existing shell/signal/loop checks plus a new narrow IO denylist:
- Always flag non-literal args to anything in _SHELL_EXEC_FUNCS,
not just _STRING_SHELL_FUNCS. Closes the dynamic-arg bypass.
- Reject calls to socket.create_connection, socket.socket().connect,
urllib.request.urlopen, http.client.HTTP*Connection, requests.*,
httpx.* whose literal host argument is in a cloud-metadata
denylist (169.254.169.254 + 169.254.* + 100.64.*, plus the
GCP/Alibaba/ECS metadata hostnames and IPv6 link-local). Public
hosts (example.com, huggingface.co, ...) still work. Dynamic
hosts cannot be statically blocked; mitigated by the bash
blocklist + the netns where the kernel allows it.
- Reject literal open("/etc/passwd"), /etc/shadow, /etc/sudoers,
/etc/ssh/*, and /proc/<pid>/environ. Other files
(/etc/os-release, /etc/hostname, /tmp/*, user dirs) still work.
The _check_code_safety summariser is updated to include the new
network_calls and sensitive_file_reads buckets in its error string.
Regression-checked: echo, sleep, ls /tmp, for loops, piped helpers
(echo a | tr a A), urllib.request.urlopen("http://example.com"),
socket.getaddrinfo("example.com",80), open("/etc/os-release"),
open("/tmp/...","w") all still succeed. curl, wget, nc, ssh, rm,
socket.create_connection(("169.254.169.254",80)),
open("/etc/passwd"), open("/proc/self/environ") all correctly
blocked.
* studio: rate-limit login, rotate refresh tokens, add logout, security headers, gate bootstrap injection
A pass over the auth surface found a cluster of related issues that this
commit closes together.
Login (routes/auth.py):
- Add an in-memory per-IP login rate limiter. Five failed POSTs to
/api/auth/login inside a 60s window produce 429 with Retry-After.
A successful login clears the bucket. Previously 30 wrong passwords
in under one second was accepted as 30x 401, which combined with
the (now fixed) admin-username leak from /api/auth/status made
brute-force trivial against a small password.
Logout (routes/auth.py):
- New POST /api/auth/logout returns 204 and calls
storage.revoke_user_refresh_tokens(subject) so the refresh token
is no longer valid. Previously POST /api/auth/logout returned 405
and there was no way to invalidate refresh tokens short of
changing the password. Frontend session.ts already calls
clearAuthTokens() to drop localStorage; the new endpoint lets the
client also tell the server to revoke server-side state.
Refresh-token rotation (routes/auth.py + auth/storage.py):
- New storage.consume_refresh_token(token) atomically validates +
deletes a refresh token, returning (username, is_desktop). The
/api/auth/refresh handler now mints both a new access AND a new
refresh token; the supplied token becomes invalid. Replaying a
consumed refresh returns 401 "Invalid or expired refresh token".
The previous refresh_access_token helper is left in place for
callers that intentionally want the non-rotating shape; nothing
in the route layer uses it now.
/api/auth/status no longer leaks default_username (models/auth.py +
routes/auth.py):
- AuthStatusResponse.default_username becomes Optional[str] with a
None default; the handler always returns None. The frontend already
hardcodes HIDDEN_LOGIN_USERNAME = "unsloth" (auth-form.tsx:82), so
no UI change is required.
window.__UNSLOTH_BOOTSTRAP__ no longer auto-injects (main.py):
- _inject_bootstrap is now opt-in via the
UNSLOTH_STUDIO_INJECT_BOOTSTRAP env var. The previous default
(inject whenever requires_password_change is true) embedded the
plaintext bootstrap password into the first-boot HTML for any
caller that hit /, /change-password, or any unknown SPA path.
Browser extensions and any XSS payload on the page could read it
trivially. With the new gate the bootstrap password lives only in
the auth/.bootstrap_password file (mode 0o600) where it has always
been; users typing it into a current-password field is the right
UX. routes/auth.py:change_password also clears
app.state.bootstrap_password defensively.
Security headers + server fingerprint (main.py + run.py):
- New SecurityHeadersMiddleware adds Content-Security-Policy,
X-Frame-Options: DENY, X-Content-Type-Options: nosniff,
Referrer-Policy: no-referrer,
Permissions-Policy: camera=(), microphone=(), geolocation=(),
interest-cohort=(), and stamps server: unsloth-studio so the
generic uvicorn banner no longer fingerprints the stack. The
uvicorn.Config gains server_header=False so it stops emitting its
own Server header.
/api/health minimisation (main.py):
- Unauthenticated GET /api/health returns just
{"status":"healthy","timestamp":...} so load-balancer liveness
probes keep working without leaking version, device_type,
chat_only, desktop_protocol_version, or studio_root_id to
arbitrary callers. A request that presents a valid Bearer token
still gets the full diagnostic payload so internal launchers and
sibling-Studio detection (which compares studio_root_id) keep
working.
Verification:
- 30 wrong-password POSTs to /api/auth/login -> first 5 = 401, 6th
through 30th = 429.
- POST /api/auth/logout with a fresh token -> 204. The matching
refresh token then fails 401.
- Login -> R1; /api/auth/refresh with R1 -> new access + R2 (R2 !=
R1); /api/auth/refresh with R1 again -> 401; /api/auth/refresh
with R2 -> still succeeds once and rotates again.
- curl /api/auth/status -> default_username: null.
- curl http://127.0.0.1/ does not contain __UNSLOTH_BOOTSTRAP__.
- curl -I / shows CSP, X-Frame-Options: DENY,
X-Content-Type-Options: nosniff, Referrer-Policy: no-referrer,
Permissions-Policy, and server: unsloth-studio.
- curl /api/health unauthenticated -> {status, timestamp} only.
curl with Authorization: Bearer <valid> -> full payload.
- Existing /api/system, /api/models/list, /api/train/status,
/api/inference/status, /api/auth/api-keys, login flow, SPA root
all still return 200 after the changes (regression smoke).
* studio: add SecurityHeadersMiddleware, MaxBodyMiddleware, /recipes redirect, gate _inject_bootstrap, minimise /api/health
This commit lands the main.py-side changes that share a single
middleware-registration spot. They are kept together because every
change here is either (a) a top-level middleware definition that has
to be added next to LoggingMiddleware, or (b) a route handler at the
same file-level.
SecurityHeadersMiddleware (Content-Security-Policy, X-Frame-Options:
DENY, X-Content-Type-Options: nosniff, Referrer-Policy: no-referrer,
Permissions-Policy, server: unsloth-studio). The previous responses
emitted no CSP, no XFO, no Referrer-Policy and were stamped
server: uvicorn.
MaxBodyMiddleware rejects POST/PUT/PATCH on the inference / dataset /
data-recipe / train / export prefixes when Content-Length exceeds
UNSLOTH_STUDIO_MAX_BODY_MB (default 100). The audit hit this by
attaching a 50 MB plain-text file to a chat message and watching
Studio base64-encode it into the JSON body; uvicorn has no enforced
cap so the only previous guard was the per-file 50 MB ceiling that
data-recipe upload routes already enforce. The new middleware extends
that ceiling to the OpenAI-compat path that the Chat attachments
flow through. Verified: a 200 MB JSON POST to /v1/chat/completions
returns HTTP 413 "Request body too large (209,715,264 bytes; max
104,857,600)". A small valid request continues to reach the handler.
_inject_bootstrap is gated behind UNSLOTH_STUDIO_INJECT_BOOTSTRAP.
The previous default was to inline window.__UNSLOTH_BOOTSTRAP__ =
{username, password} into the first-boot HTML whenever
requires_password_change was true, which exposed the plaintext
bootstrap password to any browser extension, page script, or LAN
caller on -H 0.0.0.0. The bootstrap password remains in the on-disk
.bootstrap_password file (mode 0o600) where it has always lived;
users typing it into a current-password field is the right UX.
/api/health unauthenticated returns {"status":"healthy","timestamp":
...} only; the previous payload (version, device_type, chat_only,
desktop_protocol_version, supports_desktop_auth, studio_root_id,
native_path_leases_supported) is preserved for callers that present
a valid Bearer token, so internal launchers and sibling-Studio
detection (which compares studio_root_id) keep working.
/recipes -> /data-recipes 308 redirect. The Data Recipes page lives
at /data-recipes; users typing /recipes hit the SPA catch-all and
saw "Not Found". The redirect also preserves any tail path, so
/recipes/<rest> -> /data-recipes/<rest>.
Verified end to end with curl: CSP / XFO / X-Content-Type-Options /
Referrer-Policy / Permissions-Policy all present on /, server header
is now unsloth-studio (uvicorn's own banner is suppressed via
server_header=False in run.py from the auth-batch commit). Followed
the /recipes redirect lands on the SPA HTML.
* studio: bound TrainingStartRequest hyperparameters at the schema level
POST /api/train/start accepted any value for learning_rate, batch_size,
max_steps, max_seq_length, warmup_steps, warmup_ratio, num_epochs,
save_steps, weight_decay, gradient_accumulation_steps, lora_r,
lora_alpha and lora_dropout, including -1, 0, 1e9, and non-numeric
strings like 'abc' or 'two' (which silently coerce to 0 in the
trainer). Probing showed the API returning 200 to learning_rate=-1
and batch_size=0; only max_steps had any partial clamping.
This commit adds field_validator on every numeric hyperparameter.
Bounds are chosen wide enough to span realistic single-host
configurations (B200 with 180 GB of memory comfortably fits the
upper end) while rejecting the values that always produce broken
training:
- learning_rate: parses str/float, requires 0 < lr < 1.0. Non-numeric
input raises with "learning_rate must be parseable as float (got
'abc')" instead of silently coercing to 0.
- batch_size: [1, 1024].
- gradient_accumulation_steps: [1, 4096].
- num_epochs: [1, 1000].
- max_steps: [1, 1_000_000].
- max_seq_length: [1, 131072].
- warmup_steps: [0, max_steps].
- warmup_ratio: [0.0, 1.0].
- save_steps: [0, 1_000_000].
- weight_decay: [0, 10] (typical 0..0.1).
- lora_r: [1, 512].
- lora_alpha: [1, 1024].
- lora_dropout: [0.0, 1.0).
Each validator names the offending field in its ValueError message
so the 422 response body identifies which input is bad. The
learning_rate validator returns its result as str (the schema field
type is str("2e-4") for backwards compatibility) so existing call
sites that float() the value continue to work.
Verified:
- learning_rate=-1 -> 422 "learning_rate must be > 0 (got -1.0);
typical range is 1e-6 .. 1e-3".
- learning_rate='abc' -> 422 "must be parseable as float".
- batch_size=-1 / 0 / 999999 -> 422 "batch_size must be in [1, 1024]".
- batch_size='two' -> 422 (pydantic int parser).
- max_steps=0 / -5 -> 422 "must be a positive int".
- max_seq_length=200000 -> 422 "must be in [1, 131072]".
- warmup_ratio=2.5 -> 422 "must be in [0.0, 1.0]".
- lora_dropout=1.5 -> 422 "must be in [0.0, 1.0)".
- Valid request with learning_rate='2e-4', batch_size=1, max_steps=5
passes validation and the training run starts as normal.
* studio: redact image-decode errors, clean checkpoint dirs on cancel, tolerate Stop-button + tool-result message shapes
Three small fixes that fall under "do not let the audit findings
become user-visible papercuts".
routes/inference.py - image-decode error redaction (the audit hit
this with a 0-byte / malformed / wrong-extension image upload). The
three image-normalise sites previously raised HTTPException(400,
detail=f"Failed to process image: {e}"). When PIL raised
UnidentifiedImageError(io.BytesIO(raw)) the message string included
"<_io.BytesIO object at 0x7e40a5d7bf60>", leaking both the Python
class name (confirming the PIL/io stack) and a heap address (mildly
useful for ASLR-bypass chaining if another memory-corruption bug is
ever found). Each site now catches UnidentifiedImageError and
returns the generic "Unsupported or corrupt image format"; the
fall-through generic except returns "Failed to process image". No
exception-repr is interpolated into a response body anywhere along
these paths.
core/training/training.py - checkpoint cleanup on cancel. When a
user clicks Cancel Training, the trainer flips _cancel_requested=True
and the supervisor force-terminates the subprocess. The trainer
writes checkpoint-<step> directories under output_dir every
save_steps; previously these survived the cancel and accumulated on
disk (the audit recorded ~67 MB stuck after a 200-step cancel with
save_steps=20). New helper _cleanup_cancelled_checkpoints(output_dir)
globs checkpoint-<int> entries and removes them. It is gated by a
realpath containment check against outputs_root() so it cannot
accidentally rmtree anything outside the configured outputs root.
force_terminate() invokes the helper after the subprocess join when
_cancel_requested is true. Stop-and-Save runs are unaffected because
that path keeps _cancel_requested=False.
models/inference.py - chat message shape tolerance. Two related
frontend interactions used to crash the request validator:
- After the Stop button truncates a generation, the frontend
retained {role:"assistant", content:""} in the conversation
history and replayed it on the next send. ChatMessage previously
required role="assistant" to have non-empty content or tool_calls,
so the next message returned 422 and the thread was permanently
broken. The validator now normalises empty assistant content to
None so the request round-trips and the trailing empty turn can
be ignored downstream.
- The frontend's second-round tool POST drops the streamed
tool_call_id, hitting the strict-spec check "role=tool requires
tool_call_id". The validator now synthesises an opaque id
(call_<8 hex>) when missing, so the request reaches the handler
and the model's final summarising response gets generated. The
proper fix lives in the frontend (carry the streamed id through
the second POST) and will follow.
Verified end to end with curl: HTTP 400 (model not loaded) on both
the empty-assistant history shape and the tool-result-without-id
shape, instead of HTTP 422 from the schema validator.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: tighten code comments from security-hardening pass
Trim verbose docstrings and inline finding references added in the
previous commits in this branch. Functionality unchanged.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: await get_current_subject in /api/health and make refresh-token consumption atomic
The /api/health auth probe called get_current_subject(creds) without
awaiting it. The coroutine object is truthy, so any caller presenting a
Bearer header (valid or not) received the full diagnostic payload
including version, device_type, studio_root_id, etc. Await the coroutine
and treat HTTPException as 'fall back to the minimal liveness payload'.
consume_refresh_token did SELECT then DELETE WHERE id under default
autocommit isolation. Two concurrent POST /api/auth/refresh requests
could both win the SELECT before either DELETE ran, defeating
single-use refresh-token rotation. Replace with a single
DELETE ... WHERE token_hash = ? AND expires_at >= ? RETURNING ...
statement so the validate-and-delete lands as one atomic op under
SQLite's write lock (3.45.1 supports RETURNING; min was 3.35).
* studio: enforce body cap on chunked uploads and drop unsafe-inline from script-src
MaxBodyMiddleware previously only inspected the declared Content-Length
header; clients omitting it or sending Transfer-Encoding: chunked
bypassed the cap and could still drive an OOM via the downstream
JSON / file readers on /v1/chat/completions, /api/inference, /api/data-recipe,
/api/datasets, /api/train, /api/export. Rewrite as a raw ASGI middleware
that drains and counts http.request frames, replies 413 once the running
total exceeds UNSLOTH_STUDIO_MAX_BODY_MB before invoking the FastAPI
handler, and replays the buffered body to downstream so route code that
calls request.json() / await request.body() works unchanged.
CSP previously included 'unsafe-inline' on script-src, which defeats the
main XSS protection. The frontend bundle does not need inline scripts;
the only inline <script> the backend ever emits is _inject_bootstrap,
which is opt-in via UNSLOTH_STUDIO_INJECT_BOOTSTRAP. Drop 'unsafe-inline'
from script-src by default; when _inject_bootstrap fires, generate a
per-response nonce, embed it on the inlined <script>, and have
SecurityHeadersMiddleware splice 'nonce-XXX' into the CSP for that one
response (the internal x-internal-script-nonce header is popped before
the response leaves the server). 'unsafe-inline' stays on style-src for
Vite-injected styles.
* studio: drop empty assistant sentinel before passthrough
ChatMessage._validate_role_shape normalises role="assistant", content=""
(the post-Stop sentinel emitted by the frontend) to content=None so the
in-process path can drop it via _extract_content_parts. The passthrough
path then ran m.model_dump(exclude_none=True), which strips the now-None
content key entirely, sending {"role":"assistant"} to llama-server / the
OpenAI-compat backend. That fails upstream and leaves the user without a
recoverable Stop->resume.
Add _drop_empty_assistant_sentinels and call it at both passthrough
message origins: _openai_messages_for_passthrough (covers
/v1/chat/completions and the Responses API which routes through it) and
the anthropic_messages_to_openai output before
_anthropic_passthrough_*. Assistant messages that carry only tool_calls
(no content) are preserved.
* studio/tests: cover audit-fix surfaces and rebase pre-existing tests
Adds and updates pytest coverage for the four bot-flagged audit fixes
landed earlier in this branch and rebases two pre-existing tests that
were broken by the relaxed-validator and /api/health auth-gate changes.
studio/backend/tests/test_middleware.py (new)
MaxBodyMiddleware: small protected, large declared, unprotected
passthrough, chunked-upload-over-cap rejection (the regression for
the original Content-Length-only gap), and chunked-under-cap replay.
SecurityHeadersMiddleware: script-src no longer carries
'unsafe-inline', style-src still does, default headers
(XFO/XCTO/Referrer-Policy/Permissions-Policy/server), and the
internal x-internal-script-nonce header is consumed by the
middleware and converted to 'nonce-XXX' in the CSP.
/api/health: no auth -> minimal, invalid Bearer -> minimal
(the await regression), valid Bearer -> full diagnostic payload.
studio/backend/tests/test_desktop_auth.py
consume_refresh_token: second-call returns None, expired returns
None, and a 64-thread concurrent pile-up against the same hash
produces exactly one successful consumer (regression for the
SELECT-then-DELETE race).
test_health_response_reports_desktop_capability_fields: rebase
against the new health_check(request) signature by going through
TestClient with a real bearer instead of asyncio.run-ing the
handler directly.
studio/backend/tests/test_openai_tool_passthrough.py
Pin the new ChatMessage tolerance: assistant without content or
tool_calls is tolerated (normalises content -> None), empty-string
and empty-list assistant content normalise to None, and a missing
/ empty tool_call_id on role='tool' is synthesised as call_<hex>
rather than raising. Tests for _drop_empty_assistant_sentinels
cover the three drop shapes (empty string, empty list, missing
content key), preservation of assistant text and tool_calls-only
messages, and end-to-end through
_openai_messages_for_passthrough.
studio/backend/main.py
SecurityHeadersMiddleware.dispatch used response.headers.pop(...)
for the nonce-header handoff; Starlette's MutableHeaders has no
pop. Read-then-del so the internal handoff header is still
stripped before the response leaves the server.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio/tests: rebase three more pre-existing CI tests against this branch
CI on PR #5375 was red on three tests that were tuned for behaviour
predating this branch. Updates each so the assertions match what the
audit fixes intentionally changed; no production code touched.
studio/backend/tests/test_trained_model_scan.py
test_scan_trained_models_includes_lora_and_full_finetune_outputs
passed an absolute tmp_path through scan_trained_models, which now
runs resolve_output_dir / _assert_contained against outputs_root().
Repoint outputs_root() at tmp_path via monkeypatch so the fixture
dirs land under the configured root and the realpath containment
check passes.
tests/test_studio_install_workspace_guard.py
test_health_endpoint_exposes_studio_root_id_not_raw_path read
the first 1500 bytes after @app.get("/api/health") and asserted on
the studio_root_id literal. The handler grew (unauth short-circuit
+ await dependency gate) and the literal slid past the byte window.
Replace the fixed window with a slice up to the next top-level
@app.* decorator so the test surveys the whole handler regardless
of size.
tests/studio/studio_api_smoke.py
The "login burst (5x wrong pw) -> 401 each" assertion was tagged
"When/if we add one, this assertion updates in the same PR." We
added the per-IP rate-limit in routes/auth.py
(_LOGIN_MAX_FAILS=5/60s) but missed the assertion update. Rewrite
the burst probe to observe the new invariant: at least one 401,
eventual transition to 429, and Retry-After present on the 429.
Adds a small _login_with_headers helper since the existing login()
helper drops response headers.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* ci(studio-ui): set UNSLOTH_STUDIO_INJECT_BOOTSTRAP=1 for Playwright Studios
The Chat UI Playwright test drives the first-boot change-password
form, which (per playwright_chat_ui.py step "1. Change-password
through the UI") pre-seeds the hidden current_password field from
window.__UNSLOTH_BOOTSTRAP__. That global is only emitted when the
backend's _inject_bootstrap path fires, which since the security
pass on this branch is gated behind UNSLOTH_STUDIO_INJECT_BOOTSTRAP
and defaults to off. Without the global, the React form's
current_password validator never satisfies, the submit button stays
disabled, and the composer.wait_for() probe times out on
/change-password.
Re-enable injection only for the CI Studios that drive the chat UI
across linux/mac/windows. Production deployments are unaffected: the
env var has to be explicitly opted into, and the on-disk
auth/.bootstrap_password remains the source of truth for human users
typing the password in by hand.
Covers all eight Studio launch sites: the primary chat-ui boot and
the "extra UI tests" boot for each of the three OSes, plus the
pipeTransport JSON-crash retry relaunches in the macOS workflow that
re-spawn Studio mid-job.
A follow-up frontend PR will add a visible current_password input so
the form satisfies its own validator without needing the bootstrap
auto-fill at all; once that lands this CI knob can come back out.
* studio/sandbox: drop unshare(CLONE_NEWNET); add trusted-host allowlist; block sandbox file uploads; raise CPU rlimit default to 600 s
CLONE_NEWNET inside _sandbox_preexec silently killed every outbound
HTTP request from sandboxed Python whenever the kernel allowed
unprivileged user namespaces. requests.get('https://huggingface.co'),
urllib.request.urlopen('https://en.wikipedia.org/wiki/...'),
socket.connect(('arxiv.org', 443)) all failed despite the AST visitor
intending to allow them. The bash blocklist (curl / wget / nc / ssh /
scp / sftp / rsync / socat / eval / source) plus the AST-level
metadata-host denylist still carry the network policy after this
change; CLONE_NEWNET was redundant with both.
Add _TRUSTED_PUBLIC_HOST_LITERALS + _TRUSTED_PUBLIC_HOST_SUFFIXES
(~100 informational hosts: Wikipedia language subdomains, Wikimedia,
Wikidata, Google search, Bing, DuckDuckGo, HuggingFace, GitHub,
raw.githubusercontent.com, arXiv, StackOverflow / Stack Exchange,
MDN, docs.python.org, PyTorch / TensorFlow / NumPy / pandas docs,
pypi / files.pythonhosted.org / npmjs / crates.io, ReadTheDocs,
arXiv, Britannica, BBC / Reuters / Nature / Science, NASA / CDC /
NIH / WHO open data, api.weather.gov). The visitor now blocks
literal hosts that are neither metadata nor trusted with a short
LLM-readable string so the model can retry with an allowed source
instead of choking on a multi-line error.
Block upload-shape calls regardless of host: requests.post / put /
patch / delete / request with files= or data=open(...) /
data=bytes_literal; httpx equivalents; urllib.request.urlopen /
Request with data=...; HuggingFace upload_file / upload_folder /
upload_large_folder / create_commit (module-level FQ paths AND
method-name match on any receiver). Message: "Blocked: file upload
disallowed in sandbox".
Bump UNSLOTH_STUDIO_SANDBOX_CPU_S default 300 -> 600 s so long
agentic chains that span multiple tool calls don't get SIGXCPU'd
mid-stride. Env-var override path is unchanged.
Host normalisation now strips trailing dot, userinfo @, and explicit
port before allowlist / denylist comparison so trailing-DNS-dot,
userinfo-smuggling, and explicit-:443 URLs are decided correctly.
* studio: raise default request-body cap from 100 MB to 500 MB
UNSLOTH_STUDIO_MAX_BODY_MB default goes 100 -> 500 to comfortably
cover vision + audio + multi-recipe-batch JSON payloads. The
MaxBodyMiddleware stream-counting logic from this branch's earlier
06ec088 already handles chunked bodies up to the new cap; env-var
override path is unchanged for callers that want a tighter limit.
* studio/auth: restore /api/auth/status.default_username to 'unsloth'
This branch's earlier b39e9a4 changed default_username to None on the
public /api/auth/status endpoint so the username field didn't leak to
unauthenticated callers. In practice this regressed third-party
clients (and the in-tree React login form's pre-fill UX) without
adding meaningful security: the bootstrap password is the actual
secret, and the username 'unsloth' is the documented default.
Pin default_username to storage.DEFAULT_ADMIN_USERNAME ('unsloth')
and tighten the response model so the field is required rather than
Optional. Anyone who needs anonymisation can still reach for an
allow-list deployment with auth disabled.
* studio/training: raise max_seq_length / batch_size / lora_r / lora_alpha caps
This branch's 7102815 introduced field validators with conservative
caps. The follow-up loosens them so long-context experiments and
high-rank LoRA exploration aren't gated at the schema layer:
_MAX_BATCH_SIZE 1024 -> 4096
_MAX_SEQ_LENGTH 131_072 -> 2_000_000 (2M tokens)
lora_r cap 512 -> 16_384 (_MAX_LORA_R)
lora_alpha cap 1024 -> 32_768 (_MAX_LORA_ALPHA)
_MAX_GRAD_ACCUM / _MAX_STEPS / _MAX_EPOCHS / lora_dropout /
warmup_ratio / weight_decay are unchanged. Hardware (VRAM, host
RAM, kernel launch latency) is now the binding constraint at the
new caps, which is the correct ordering -- the validator stays a
sanity check on -1 / 0 / 'abc' style garbage, not a usability gate.
* studio/tests: cover sandbox allowlist + upload block + raised training caps
studio/backend/tests/test_sandbox_tools.py (new):
TestMetadataHostDenylist -- short "Blocked: cloud-metadata host"
message on AWS IMDS, GCP metadata,
Alibaba ECS, AWS IPv6 IMDS, 169.254/16.
TestTrustedHostAllowlist -- Wikipedia (any language subdomain),
Google, DuckDuckGo, HF, raw GitHub,
arXiv, StackOverflow / family,
MDN, docs.python.org, pypi, BBC,
api.weather.gov, NumPy / PyTorch docs.
TestUntrustedHostBlock -- example.com / random unlisted host
rejected with the short "Blocked: host
not in sandbox allowlist; use an
allowed informational source" message.
Dynamic URLs (computed var) still pass
-- documented limit of static analysis.
TestHostNormalization -- trailing dot, explicit :443, uppercase,
userinfo-@-smuggle all decided
correctly without false-block /
false-pass.
TestUploadDenylist -- requests / httpx / urllib.urlopen with
files= / data=open / data=bytes,
HfApi().upload_file / upload_folder /
create_commit, module-level
huggingface_hub.upload_folder. POST
json= to trusted host still passes.
TestSandboxCpuRlimitDefault -- pin UNSLOTH_STUDIO_SANDBOX_CPU_S=600
default and confirm CLONE_NEWNET
source line is gone.
TestMaxBodyDefault -- pin UNSLOTH_STUDIO_MAX_BODY_MB=500
default.
studio/backend/tests/test_studio_train_validation.py (new):
Pin at-cap-accepts / over-cap-rejects boundaries for
max_seq_length=2_000_000, batch_size=4_096, lora_r=16_384,
lora_alpha=32_768 so a future regression that tightens them back
without explicit user opt-in is caught.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* studio: tighten code comments across the security-hardening pass
* studio: always inject bootstrap credentials on first boot
The UNSLOTH_STUDIO_INJECT_BOOTSTRAP gate added an extra
terminal-to-browser copy-paste on every fresh install. In practice
the LAN credential leak it guarded against is narrow: the password
is one-time, the user rotates it on the very next click, the
default Studio bind is 127.0.0.1, and -H 0.0.0.0 already exposes
the entire API surface. Drop the gate so the inject fires whenever
a bootstrap password is still pending. The CSP nonce wiring stays
in place; the inline script remains the only inline script the
backend ever emits.
The three Playwright UI smoke workflows lose their
UNSLOTH_STUDIO_INJECT_BOOTSTRAP=1 lines along with the explanatory
comment blocks since the inject now happens by default.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Wasim Yousef Said <wasimysdev@gmail.com>
* Dark theme refactor, right sidebar redesign, and chat UI polish
- Dark theme refactor
- Redesign right sidebar
- Further left sidebar adjustments
- Wider chat and content area; layout tweaks for chat content
- Rounded corners across elements for consistency
- Show chat message menu icons on menu-area hover, not only on message hover
- Assistant message menu icons now always visible; user messages keep on-hover
- Redesigned copy icon used consistently across chat blocks and messages
- Redesigned trash icon, applied consistently
- Unified icon sizing and style with the sidebar
- Adjusted icon colors across chat
- Fix on-hover background design for chat icons
- Fix tooltip from 'more' button staying visible after clicking elsewhere
- Adjust position and design of generation speed info text below messages
- Adjust design of token speed info popup
- Adjust sidebar scrollbar to cover recent chats only
* Recents sidebar rename, UI/theme refactor, layout and chat polish
UI & Theme:
- Dark theme refactor
- Consistent rounded corners across elements
- CSS polish and cleanup
- Remove unused logo image assets
Recents sidebar:
- Add 'more' button for options menu
- Support renaming conversations and training runs
- Confirmation dialog before deleting chats
- Add optional display_name column to training_runs (idempotent ALTER TABLE) so renaming doesn't lose model_name/dataset_name from the run config
- New PATCH /api/train/runs/{run_id} endpoint accepts { display_name: string | null }; empty/whitespace clears the override
- Sidebar shows display_name ?? model_name and exposes Rename in the row's More menu, mirroring the chat rename flow
- Cache last list response in localStorage and hydrate from it on mount, so recents paint instantly on F5 / route revisit; cached items are shape-validated and dropped if malformed
- Optimistic updates on rename and delete (apply locally + cache before background refresh)
- Visible toast on rename/delete failure instead of swallowed errors
Layout:
- Redesigned right sidebar
- Further left sidebar adjustments
- Updated chat content layout; chat and content area slightly widened
- Sidebar scrollbar covers recent chats only
Icons:
- Redesigned copy icon, unified across chat blocks and messages
- Redesigned trash icon to match
- Consistent icon sizing and style across chat and sidebar
- Adjusted icon colors across chat
- Fix icon on-hover background design
Chat messages:
- Menu icons now appear on hover over the menu area, not just the message
- Assistant message menu icons always visible; user messages keep on-hover (next/previous response stays visible for edited prompts)
- Repositioned and restyled generation speed info text below messages
- Restyled token generation speed popup
Tooltips:
- Removed tooltip on hover for previous/next assistant response icons
- Unified tooltip design across sidebars and chat
- Removed tooltip animations (also fixes related lag)
Model & Chat Template config:
- Merged Chat Template config into Model Configuration section
- Added revert-to-original for chat template
- Fix Chat Template config disappearing on page refresh until model reload
Performance & scroll:
- Removed chatbox movement animations across pages/navigation (fixes related UI lag)
- Fix scroll flicker at end of streaming when a code block is the final element
- Additional chat scroll improvements
Bug fixes:
- Fix 'more' button tooltip remaining visible after clicking elsewhere
* Remove sidebar localStorage cache and optimistic updates
Drops the localStorage hydration and optimistic rename/delete logic from the recents sidebar; reverts to fetching fresh on mount.
* Fix missing cn import in shared-composer (regression from merge)
* chore(sidebar): import sidebar deps from feature indexes
Re-export deleteChatItem / renameChatItem / useChatSidebarItems / SidebarItem / useChatSearchStore / ChatSearchDialog from @/features/chat, and removeTrainingUnloadGuard from @/features/training. Switch app-sidebar.tsx to consume them via the public feature indexes instead of deep paths, clearing the no-restricted-imports eslint errors. No behavior or UX change.
* fix(studio/frontend): reload training Recents sidebar after F5 refresh
The Recents sidebar showed empty after a hard refresh. The hook's inFlightRef dedup guard collided with React StrictMode's double-mount in dev: the second mount's fetch returned silently with no error, no retry, and no toast — leaving the sidebar empty until navigation.
Replace skip-if-busy dedup with abort-previous via a hook-level AbortController. This also fixes a latent race where a slow poll could resurrect a just-deleted row by clobbering the optimistic update.
Changes (all in use-training-history-sidebar.ts):
- fetchRuns aborts any in-flight request before starting a new one; post-await signal.aborted check drops stale responses.
- Optimistic helpers (applyRunUpdate, removeRun) abort in-flight fetches so they don't depend on caller discipline to invalidate stale data.
- Initial load gets bounded retry-with-backoff (500ms / 1.5s / 3.5s) and surfaces a sonner toast with a Retry action on final failure.
- Failure toast auto-dismisses on any successful load (initial retry, Retry click, or polling recovery).
- Polling pauses while the tab is hidden and catches up on visible, avoiding wasted requests during long training runs.
- Both effects own their teardown explicitly (abort + clear timer).
* Apply unified tooltip design and behavior across remaining pages for consistency
* UI polish: spacing, tooltip on source icons, letter spacing, smaller icons, consistent edit icon
- Adjust tiny spacing between elements around the UI for subtle polish
- Redesign tooltip on source icons for web search / tool use, consistent with the new design
- Adjust chat text letter spacing
- Smaller icon sizes
- Replace 'edit message' icon in chat with the new Rename icon used in Recents for consistency
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Adjust CSS for right sidebar
* Fix scrollbar UI compatibility across browsers
* fix: preserve chat preset settings on model load
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix(studio): remove duplicate chat template status field
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* chore: remove creative preset assumption
* fix(studio): align speculative decoding default
* fix(studio/chat): snap numeric param inputs to step grid
- Type a value in any param input (Temperature, Top K, Max Tokens, etc.)
now clamps to [min, max] and snaps to the slider's step grid, killing
off-grid values like 1.051234 and FP residue from slider drags.
- Branch picker chevrons share the action bar's 32px height + 10px radius
via a new .aui-branch-chevron-btn utility; hover area aligns visually
while staying narrower than the sibling icon buttons.
* fix(studio/chat): keep training-run polls converging and drop dead preset code
- Keep training-run polls converging when responses outrun the 5s interval
(don't unconditionally abort prior in-flight; skip if one is still pending,
mutation race still guarded).
- Drop dead Creative/Precise preset code paths (remove 'builtin-fixed' source
variant + unreachable branches).
* fix(studio): training-run cards show custom name + model + dataset
- Training-run cards now display custom display_name + model + dataset,
with cross-view sync on rename/delete.
- Enhance clarity of borders and colors in dark theme on export etc.
* fix(studio): match active state green to unsloth brand color
* fix(studio): preserve can_resume on training rename
* fix(studio): keep GGUF chat template override distinct
* fix(studio): treat audio input models as multimodal
* fix(studio): cancel numeric draft on Escape
* fix(studio): use default speculative mode on toggle
* fix(studio): detect GGUF audio VLM input models
* fix(studio): address final PR review findings
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix(studio): refresh sidebar/history when a new training run starts so it appears without a manual reload
* fix: API and svg
* fix(studio/sidebar): align run rename dirty check with displayed baseline
* fix(studio/sidebar): use leading-tight on account block to prevent descender clipping with truncate
---------
Co-authored-by: sneakr <hauzin@hotmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: shine1i <wasimysdev@gmail.com>
* feat(studio): add Continued Pretraining (CPT) support
Implements CPT as a first-class training method in Unsloth Studio,
resolving feature request #4565.
Changes:
- frontend/src/types/training.ts: add 'cpt' to TrainingMethod union
- frontend/src/lib/vram.ts: add 'cpt' to VramTrainingMethod (fp16 footprint)
- frontend/src/features/export/constants.ts: add CPT to METHOD_LABELS
- frontend/src/features/training/api/mappers.ts: map 'cpt' -> 'Continued Pretraining',
force packing=true and train_on_completions=false for CPT payloads
- frontend/src/features/studio/sections/model-section.tsx: add 'Continued Pretraining'
option (purple dot) to Method selector; update tooltip
- frontend/src/features/onboarding/.../model-selection-step.tsx: add CPT to
onboarding wizard method dropdown
- backend/models/training.py: update training_type field description
- backend/core/training/worker.py: detect is_cpt flag, force packing=True,
train_on_completions=False, pass is_cpt to _train_worker
- backend/core/training/trainer.py: _train_worker reads is_cpt kwarg, forces
packing on, skips train_on_responses_only for raw-text pretraining
CPT behaviour:
- Full model weights (no LoRA adapters), same as Full Finetuning
- Sequence packing always enabled for GPU efficiency
- Trains on every token (no chat-format masking)
- VRAM estimated at fp16 (2.0 bytes/param)
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update mappers.ts
* Add CPT raw dataset support and UI fixes
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add missing training methods module
* Handle invalid raw-text rows and expose raw in onboarding
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Etherll <61019402+Etherll@users.noreply.github.com>
Co-authored-by: Etherll <mrmrmidessam@gmail.com>
* feat: add checkpoint resume for stopped training runs
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix:add resume checkpoint helpers
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix: use checkpoint parent as resume output dir
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix: save optimizer and scheduler state on stop-and-save
Use Trainer._save_checkpoint instead of save_state so resume restores
optimizer momentum and LR-schedule position via the checkpoint-NNN/
subdir written by HF's official path.
* fix: clean up resume training history and startup progress
* fix: preserve resume output dirs
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix: tighten resume run lookup
* fix: remove stale output-dir lookup
* fix: preserve startup download progress
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
* fix(studio): change default weight_decay from 0.01 to 0.001
The default weight decay across Studio was 0.01 but should be 0.001.
Updated the default in all backend fallbacks, the Pydantic model, the
frontend config, and every YAML preset/model-default config.
* fix(studio): auto-set learning rate based on training method
Default LR should be 2e-4 for LoRA/QLoRA and 2e-5 for full fine-tuning.
Frontend: track whether the user has manually edited the LR field via a
_learningRateManuallySet flag (same pattern as trainOnCompletions).
When switching training method and the user has not touched the LR,
auto-set it to the appropriate default. Reset the flag on model load.
Backend: change trainer.py start_training default from 5e-5 to 2e-4,
update default.yaml fallback from 5e-5 to 2e-4, and fix
full_finetune.yaml from 0.0002 (2e-4) to 2e-5.
* refactor(studio): centralize weight_decay and learning rate defaults
Create studio/backend/core/training/constants.py as the single source of
truth for DEFAULT_WEIGHT_DECAY (0.001), DEFAULT_LEARNING_RATE (2e-4),
DEFAULT_LEARNING_RATE_FULL (2e-5), and DEFAULT_LEARNING_RATE_STR ("2e-4").
All backend modules (trainer.py, training.py, worker.py, models/training.py)
now import from constants.py instead of hardcoding values.
On the frontend, add LR_DEFAULT_LORA and LR_DEFAULT_FULL to
config/training.ts and use them in the store instead of magic numbers.
A comment cross-references the backend constants file.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix model-specific LR override, persist migration, and flag resets
- Preserve model-specific learning rates from YAML configs when the
async autoSelectTrainingMethod callback fires (fixes Qwen2.5-1.5B
getting 2e-4 instead of its configured 1e-5, etc.)
- Bump zustand persist version to 9 with migration so existing users
with weightDecay=0.01 get updated to 0.001
- Clear _learningRateManuallySet in reset() and applyConfigPatch()
for consistency with trainOnCompletions flag behavior
- Add DEFAULT_LEARNING_RATE_FULL_STR to constants.py
* Refine applyConfigPatch to only clear LR flag when patch includes LR
Only reset _learningRateManuallySet when the applied config patch
actually provides a learningRate value. This prevents unrelated config
patches from silently disarming the manual-edit guard, which would
cause a subsequent setTrainingMethod call to overwrite the user's
custom LR.
* Preserve model-specific LR when switching between qlora and lora
Only auto-switch the learning rate when the training category changes
(adapter <-> full fine-tuning). Switching between qlora and lora keeps
the current LR since both methods share the same learning rate range.
This preserves curated per-model defaults (e.g. 1e-5 for
Qwen2.5-1.5B-Instruct) when the user toggles between adapter methods.
* Remove constants.py, use YAML configs as the source of truth
The YAML config files (model-specific + default.yaml) are the intended
config layer for training defaults. The Python backend fallbacks now use
inline values that match the YAML configs, rather than importing from a
separate constants module. This keeps the config architecture simple:
YAML files are the single source of truth, and the inline Python
fallbacks are just safety nets that mirror them.
* fix(studio): preserve model-specific LR when switching training method
Stash YAML-provided learning rate and use it to restore the correct
value when switching between adapter and full fine-tune modes.
- qlora <-> lora no longer overwrites the model's LR
- full -> adapter restores the YAML LR instead of a hardcoded constant
- selecting a model while on full fine-tune uses LR_DEFAULT_FULL
instead of applying the YAML adapter LR
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@users.noreply.github.com>
Co-authored-by: Roland Tannous <rolandtannous@gravityq.ai>
* [WIP] balanced device map for studio
* gpus as a request parameter
* API for multi GPU stuff
* return multi gpu util in new API
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Use balanced_low0 instead of balanced
* Use balanced_low0 instead of balanced
* Fix device_map typo, UUID parsing crash, set() filter bug, and broken tests
- balanced_low0 -> balanced_low_0 (transformers/accelerate rejects the old string)
- get_parent_visible_gpu_ids() now handles UUID/MIG CUDA_VISIBLE_DEVICES
gracefully instead of crashing on int() parse
- _get_backend_visible_gpu_info() set() or None bug: empty set is falsy so
CUDA_VISIBLE_DEVICES=-1 would disable filtering and report all GPUs
- test_gpu_selection.py: add missing get_visible_gpu_utilization import and
add required job_id arg to start_training() calls
* Smart GPU determinism using estimates
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* disallow gpu selection for gguf for now
* cleanup
* Slightly larger baseline
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Treat empty list as auto
* Verbose logging/debug
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Cleanup and revert unnecessary deletions
* Cleanup excessive logs and guard against disk/cpu offload
* auth for visibility API. cleanup redundant imports. Adjust QLoRA estimate
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* support for non cuda gpus
* Fix multi-GPU auto-selection memory accounting
The multi_gpu_factor was applied uniformly to all GPUs including the
first one, which unfairly penalizes single-GPU capacity when
transitioning to multi-GPU. This created a discontinuity where a model
that barely fits 1 GPU would suddenly require 2 GPUs because the first
GPU's free memory was discounted by 20%.
Now the first GPU keeps its full free memory, and only additional GPUs
have an overhead factor (0.85) applied to account for inter-GPU
communication and sharding overhead. This gives more accurate
auto-selection and avoids unnecessary multi-GPU for models that
comfortably fit on one device.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add sandbox tests for multi-GPU selection logic
24 tests covering model size estimation, memory requirements, automatic
GPU selection, device map generation, GPU ID validation, and multi-GPU
overhead accounting. All tests use mocks so they run without GPUs on
Linux, macOS, and Windows.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix reviewer findings: 4bit inference estimate, fallback, GGUF gpu_ids, retry
1. 4-bit inference now uses reduced memory estimate (model_size/3 + buffer)
instead of the FP16 1.3x multiplier. This prevents over-sharding
quantized models across unnecessary GPUs.
2. When model size estimation fails, auto_select_gpu_ids now falls back to
all visible GPUs instead of returning None (which could default to
single-GPU loading for an unknown-size model).
3. GGUF inference route now treats gpu_ids=[] as auto-selection (same as
None) instead of rejecting it as an unsupported explicit request.
4. Training retry path for "could not get source code" now preserves the
gpu_ids parameter so the retry lands on the same GPUs.
5. Updated sandbox tests to cover the new 4-bit inference estimate branch.
* Remove accidentally added unsloth-zoo submodule
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix UUID/MIG visibility and update test expectations
1. nvidia.py: When CUDA_VISIBLE_DEVICES uses UUID/MIG tokens, the
visibility APIs now return "unresolved" with empty device lists instead
of exposing all physical GPUs. This prevents the UI from showing GPUs
that the backend process cannot actually use.
2. test_gpu_selection.py: Updated test expectations to match the new
multi-GPU overhead accounting (first GPU at full capacity, 0.85x for
additional GPUs) and 4-bit inference memory estimation formula.
All 60 tests now pass.
* Add CPU/disk offload guard to audio inference path
The audio model loading branch returned before the common
get_offloaded_device_map_entries() check, so audio models loaded with a
multi-GPU device_map that spilled layers to CPU/disk would be accepted
instead of rejected. Now audio loads also verify no modules are offloaded.
* Improve VRAM requirement estimates
* Replace balanced_low_0 with balanced
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* refine calculations for slightly easier nums
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* adjust estimates
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Use nums instead of obj to avoid seralisation error
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Harden nvidia-smi parsing and fix fallback GPU list
1. nvidia.py: Wrap int() casts for GPU index and memory in try/except
so MIG slices, N/A values, or unexpected nvidia-smi output skip the
unparseable row instead of aborting the entire GPU list.
2. nvidia.py: Handle GPU names containing commas by using the last
field as memory instead of a fixed positional index.
3. hardware.py: fallback_all now uses gpu_candidates (GPUs with verified
VRAM data) instead of raw devices list, which could include GPUs
with null VRAM that were excluded from the ranking.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* cleanup
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* consolidate raise_if_offload
* Improve MoE support. Guard against nvidia-smi failures
* Improve MoE support. Guard against nvidia-smi failures
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix shared-expert LoRA undercount, torch VRAM fallback, and apply_gpu_ids edge case
1. vram_estimation.py: compute_lora_params now includes shared experts
(n_shared_experts) alongside routed experts when computing MoE LoRA
adapter parameters. Previously only n_experts were counted, causing
the estimator to undercount adapter, optimizer, and gradient memory
for DeepSeek/GLM-style models with shared experts.
2. hardware.py: _torch_get_per_device_info now uses mem_get_info (which
reports system-wide VRAM usage) instead of memory_allocated (which
only reports this process's PyTorch allocations). This prevents
auto-selection from treating a GPU as mostly free when another
process is consuming VRAM. Falls back to memory_allocated when
mem_get_info is unavailable.
3. hardware.py: apply_gpu_ids([]) now returns early instead of setting
CUDA_VISIBLE_DEVICES="" which would disable CUDA entirely. Empty
list inherits the parent visibility, same as None.
4. hardware.py: Upgraded fallback_all GPU selection log from debug to
warning so operators are notified when the model likely will not fit
in available VRAM.
* Guard nvidia-smi subprocess calls against OSError and TimeoutExpired
get_visible_gpu_utilization and get_backend_visible_gpu_info now catch
OSError (nvidia-smi not found) and TimeoutExpired internally instead
of relying on callers to wrap every invocation. Returns the standard
available=False sentinel on failure so the torch-based fallback in
hardware.py can take over.
* Guard get_primary_gpu_utilization and reset GPU caches between tests
1. nvidia.py: get_primary_gpu_utilization now catches OSError and
TimeoutExpired internally, matching the pattern already used in
get_visible_gpu_utilization and get_backend_visible_gpu_info. All
three nvidia-smi callers are now self-contained.
2. test_gpu_selection.py: Added _GpuCacheResetMixin that resets the
module-level _physical_gpu_count and _visible_gpu_count caches in
tearDown. Applied to all test classes that exercise GPU selection,
device map, or visibility functions. This prevents stale cache
values from leaking between tests and causing flaky results on
machines with real GPUs.
* Fix nvidia-smi fallback regression and physical GPU count validation
1. hardware.py: get_gpu_utilization, get_visible_gpu_utilization, and
get_backend_visible_gpu_info now check result.get("available") before
returning the nvidia-smi result. When nvidia-smi is unavailable or
returns no data (e.g., containers without nvidia-smi, UUID/MIG masks),
the functions fall through to the torch-based fallback instead of
returning an empty result. This fixes a regression where the internal
exception handling in nvidia.py prevented the caller's except block
from triggering the fallback.
2. hardware.py: resolve_requested_gpu_ids now separates negative-ID
validation from physical upper-bound validation. The physical count
check is only enforced when it is plausibly a true physical count
(i.e., higher than the largest parent-visible ID), since
torch.cuda.device_count() under CUDA_VISIBLE_DEVICES returns the
visible count, not the physical total. The parent-visible-set check
remains authoritative in all cases. This prevents valid physical IDs
like [2, 3] from being rejected as "out of range" when nvidia-smi is
unavailable and CUDA_VISIBLE_DEVICES="2,3" makes torch report only
2 devices.
* Fix UUID/MIG torch fallback to enumerate devices by ordinal
When CUDA_VISIBLE_DEVICES uses UUID or MIG identifiers,
get_parent_visible_gpu_ids() returns [] because the tokens are
non-numeric. The torch fallback in get_visible_gpu_utilization() and
get_backend_visible_gpu_info() previously passed that empty list to
_torch_get_per_device_info(), getting nothing back.
Now both functions detect the empty-list case and fall back to
enumerating torch-visible ordinals (0..device_count-1) with
index_kind="relative". This means the UI and auto-selection still
see real device data in Kubernetes, MIG, and Slurm-style UUID
environments where nvidia-smi output cannot be mapped to physical
indices.
Updated test_uuid_parent_visibility to verify the new torch fallback
path returns available=True with relative ordinals.
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Add type hint for gpu_ids parameter in InferenceOrchestrator.load_model
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
* feat(db): add SQLite storage layer for training history
* feat(api): add training history endpoints and response models
* feat(training): integrate DB persistence into training event loop
* feat(ui): add training history views and card grid
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix(studio): address review issues in training history persistence
- Strip hf_token/wandb_token from config before SQLite storage
- Add UUID suffix to job_id for collision resistance
- Use isfinite() for 0.0 metric handling throughout
- Respect _should_stop in error event finalization
- Run schema DDL once per process, not per connection
- Close connection on schema init failure
- Guard cleanup_orphaned_runs at startup
- Cap _metric_buffer at 500 entries
- Make FLUSH_THRESHOLD a class constant
- Map 'running' to 'training' phase in historical view
- Derive LR/GradNorm from history arrays in historical view
- Fix nested button with div[role=button] in history cards
- Guard String(value) against null/undefined in config popover
- Clear selectedHistoryRunId on auto tab switch
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix(studio): address round-2 review findings across training backend and frontend
Backend (training.py):
- Move state mutation after proc.start() so a failed spawn does not wedge
the backend with is_training=True
- Create DB run row eagerly after proc.start() so runs appear in history
during model loading, not after first metric event
- Rewrite _flush_metrics_to_db() with snapshot-before-insert pattern to
preserve metrics arriving during the write and retain buffer on failure
- Guard eval_loss with float() coercion and math.isfinite(), matching the
existing grad_norm guard
- Increase pump thread join timeout from 3s to 8s to cover SQLite's
default 5s lock timeout
Frontend (studio-page.tsx):
- Fix history navigation: check isTrainingRunning instead of
showTrainingView in onSelectRun so completed runs are not misrouted
- Replace activeTab state + auto-switch useEffect with derived tab to
eliminate react-hooks/set-state-in-effect lint violation
Frontend (historical-training-view.tsx):
- Add explicit "running" branch to message ternary so running runs no
longer fall through to "Training errored"
- Derive loading from detail/error state and move cleanup to effect
return to eliminate react-hooks/set-state-in-effect lint violation
Frontend (progress-section.tsx):
- Derive stopRequested from isTrainingRunning && stopRequestedLocal to
eliminate react-hooks/set-state-in-effect lint violation and remove
unused useEffect import
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix(studio): resolve 3 remaining bugs from round-2 review
1. Stuck on Current Run tab [12/20]: Only force "current-run" tab when
isTrainingRunning is true, not when stale completed-run data exists.
After training ends, users can freely navigate to Configure.
2. Incomplete metric sanitization [7/20]: Apply float() coercion and
isfinite() guards to loss and learning_rate, matching the existing
pattern used by grad_norm and eval_loss. Prevents TypeError from
string values and NaN leaks into history arrays.
3. Stop button state leak across runs [10/20]: Add key={runtime.jobId}
to ProgressSection so React remounts it when a new run starts,
resetting stopRequestedLocal state.
* fix(studio): deduplicate loss/lr sanitization in training event handler
Reuse _safe_loss/_safe_lr from the progress update block instead of
re-sanitizing the same raw event values for metric history.
* fix(studio): restore loss > 0 guard to prevent eval steps injecting 0.0 into metric histories
Round-2/3 fixes relaxed the history append guard from `loss > 0` to
`loss is not None`, which let eval-only log events (where loss defaults
to 0.0) append fake zeros into loss_history and lr_history. Restore the
`loss > 0` check to match the worker's own has_train_loss gate. The
float() coercion and isfinite() sanitization from round-3 remain intact.
* fix(studio): resolve training history bugs — nullable loss/lr, tab nav, sparkline
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.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>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
* user can upload eval dataset, removed bugs
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* resolving merge conflicts
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* resolving gpt comments
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Roland Tannous <115670425+rolandtannous@users.noreply.github.com>
Add end-to-end embedding/sentence-transformer training pipeline using
FastSentenceTransformer, SentenceTransformerTrainer, and
MultipleNegativesRankingLoss with BatchSamplers.NO_DUPLICATES.
Backend:
- Add is_embedding_model() detection via HF tags + pipeline_tag
- Add /check-embedding/ API route and EmbeddingCheckResponse
- Extend derive_model_type() to return "embeddings"
- Add _run_embedding_training() in worker.py with progress callbacks,
stop handling, LoRA (task_type=FEATURE_EXTRACTION), and model saving
- Add is_embedding field to TrainingStartRequest and ModelDetails
- Add YAML configs for 5 models: all-MiniLM-L6-v2, bge-m3,
embeddinggemma-300m, gte-modernbert-base, Qwen3-Embedding-0.6B
Frontend:
- Wire isEmbeddingModel flag through store, API types, and mappers
- Force packing=false, train_on_completions=false, warmup_ratio=0.03
- Hide packing and train_on_completions checkboxes for embedding models
- Auto-set modelType to "embeddings" from backend model_type response
Non-conversational HF datasets (e.g. stanfordnlp/snli) were naively mapped
column→role, producing poor training results. The AI Assist button now runs
a 3-pass advisor using Qwen 7B that:
1. Fetches the HF dataset card/README to understand the dataset purpose
2. Classifies the dataset type and determines if conversion is needed
3. Generates a system prompt, user/assistant templates with {column}
placeholders, and label mappings (e.g. 0→entailment)
4. Validates the conversion quality (score ≥7/10 required)
Architecture: advisor metadata flows as __-prefixed keys in
custom_format_mapping (e.g. __system_prompt, __user_template,
__assistant_template, __label_mapping). The existing _apply_user_mapping()
detects these keys and routes to template-based conversation construction.
No __ keys = existing simple mode (backwards compatible).
Backend: upgraded llm_assist.py (7B default, multi-pass advisor,
HF card fetching), extended API models, added _apply_template_mapping()
to dataset_utils.py.
Frontend: extended store with advisor state fields, wired AI Assist
to store templates/system prompt, inject __ metadata in training request,
show advisor notification banner in mapping card.
Add Start/End index inputs under Advanced in the dataset card,
allowing users to slice a dataset by row range before training.
Wired end-to-end: frontend store, API payload, backend Pydantic
model, and trainer dataset loading (inclusive on both ends).
Add Start/End index inputs under Advanced in the dataset card,
allowing users to slice a dataset by row range before training.
Wired end-to-end: frontend store, API payload, backend Pydantic
model, and trainer dataset loading (inclusive on both ends).
- Changed default eval_steps from 0.01 to 0.0 across backend and frontend
- Fixed UI to allow eval_steps=0 (removed min=0.001 constraint)
- Added conditional eval logic with helpful console messages
- Updated tooltip to explain how to disable evaluation
- Tested: confirmed eval disabled by default with eval_steps=0.0
Daniel Han